From 940aad0831a38442feedfb9f07ef58b877e0b9c8 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Fri, 10 Feb 2023 12:55:55 +0000
Subject: [PATCH 1/3] fix(api): Allow workspace agent coordinate to report
 disconnect

---
 coderd/coderd.go          | 11 +++++++++++
 coderd/workspaceagents.go | 16 ++++++++++++----
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/coderd/coderd.go b/coderd/coderd.go
index bfce5a5fb1a88..bc8fd763dae74 100644
--- a/coderd/coderd.go
+++ b/coderd/coderd.go
@@ -223,7 +223,11 @@ func New(options *Options) *API {
 	)
 
 	r := chi.NewRouter()
+	ctx, cancel := context.WithCancel(context.Background())
 	api := &API{
+		ctx:    ctx,
+		cancel: cancel,
+
 		ID:          uuid.New(),
 		Options:     options,
 		RootHandler: r,
@@ -669,6 +673,11 @@ func New(options *Options) *API {
 }
 
 type API struct {
+	// ctx is canceled immediately on shutdown, it can be used to abort
+	// interruptable tasks.
+	ctx    context.Context
+	cancel context.CancelFunc
+
 	*Options
 	// ID is a uniquely generated ID on initialization.
 	// This is used to associate objects with a specific
@@ -703,6 +712,8 @@ type API struct {
 
 // Close waits for all WebSocket connections to drain before returning.
 func (api *API) Close() error {
+	api.cancel()
+
 	api.WebsocketWaitMutex.Lock()
 	api.WebsocketWaitGroup.Wait()
 	api.WebsocketWaitMutex.Unlock()
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
index 5c14f5ea217a5..8a3f94ce72050 100644
--- a/coderd/workspaceagents.go
+++ b/coderd/workspaceagents.go
@@ -601,7 +601,7 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
 		Valid: true,
 	}
 	disconnectedAt := workspaceAgent.DisconnectedAt
-	updateConnectionTimes := func() error {
+	updateConnectionTimes := func(ctx context.Context) error {
 		err = api.Database.UpdateWorkspaceAgentConnectionByID(ctx, database.UpdateWorkspaceAgentConnectionByIDParams{
 			ID:               workspaceAgent.ID,
 			FirstConnectedAt: firstConnectedAt,
@@ -620,15 +620,23 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
 	}
 
 	defer func() {
+		// If connection closed then context will be canceled, try to
+		// ensure our final update is sent. By waiting at most the agent
+		// inactive disconnect timeout we ensure that we don't block but
+		// also guarantee that the agent will be considered disconnected
+		// by normal status check.
+		ctx, cancel := context.WithTimeout(api.ctx, api.AgentInactiveDisconnectTimeout)
+		defer cancel()
+
 		disconnectedAt = sql.NullTime{
 			Time:  database.Now(),
 			Valid: true,
 		}
-		_ = updateConnectionTimes()
+		_ = updateConnectionTimes(ctx)
 		_ = api.Pubsub.Publish(watchWorkspaceChannel(build.WorkspaceID), []byte{})
 	}()
 
-	err = updateConnectionTimes()
+	err = updateConnectionTimes(ctx)
 	if err != nil {
 		_ = conn.Close(websocket.StatusGoingAway, err.Error())
 		return
@@ -668,7 +676,7 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
 			Time:  database.Now(),
 			Valid: true,
 		}
-		err = updateConnectionTimes()
+		err = updateConnectionTimes(ctx)
 		if err != nil {
 			_ = conn.Close(websocket.StatusGoingAway, err.Error())
 			return

From 4ce0b84747e4ff22412abab1b30fe491e2823c4c Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Fri, 10 Feb 2023 13:01:59 +0000
Subject: [PATCH 2/3] Use ctx for publish

---
 coderd/workspaceagents.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
index 8a3f94ce72050..f8f40d362e999 100644
--- a/coderd/workspaceagents.go
+++ b/coderd/workspaceagents.go
@@ -633,7 +633,7 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
 			Valid: true,
 		}
 		_ = updateConnectionTimes(ctx)
-		_ = api.Pubsub.Publish(watchWorkspaceChannel(build.WorkspaceID), []byte{})
+		api.publishWorkspaceUpdate(ctx, build.WorkspaceID)
 	}()
 
 	err = updateConnectionTimes(ctx)

From ad1a46f6ea4ecdba22cc6042071b21e97a2fd362 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Fri, 10 Feb 2023 13:03:55 +0000
Subject: [PATCH 3/3] I type I typo

---
 coderd/coderd.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/coderd/coderd.go b/coderd/coderd.go
index bc8fd763dae74..11fbcf9432f7f 100644
--- a/coderd/coderd.go
+++ b/coderd/coderd.go
@@ -674,7 +674,7 @@ func New(options *Options) *API {
 
 type API struct {
 	// ctx is canceled immediately on shutdown, it can be used to abort
-	// interruptable tasks.
+	// interruptible tasks.
 	ctx    context.Context
 	cancel context.CancelFunc