feat: Make workspace watching realtime instead of polling

This was leading to performance issues on the frontend, where the page should only be rendered if changes occur. While this could be changed on the frontend, it was always the intention to make this socket ~realtime anyways.
coder · kylecarbs · Nov 7, 2022 · Nov 6, 2022 · Nov 7, 2022 · Nov 6, 2022
commit f2cce85f66a41268b65aa22c6cf4ad7e544e29ea
diff --git a/coderd/activitybump.go b/coderd/activitybump.go
@@ -54,7 +54,7 @@ func activityBumpWorkspace(log slog.Logger, db database.Store, workspace databas
 
 		newDeadline := database.Now().Add(bumpAmount)
 
-		if err := s.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
+		if _, err := s.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
 			ID:               build.ID,
 			UpdatedAt:        database.Now(),
 			ProvisionerState: build.ProvisionerState,

diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -2823,7 +2823,7 @@ func (q *fakeQuerier) UpdateWorkspaceLastUsedAt(_ context.Context, arg database.
 	return sql.ErrNoRows
 }
 
-func (q *fakeQuerier) UpdateWorkspaceBuildByID(_ context.Context, arg database.UpdateWorkspaceBuildByIDParams) error {
+func (q *fakeQuerier) UpdateWorkspaceBuildByID(_ context.Context, arg database.UpdateWorkspaceBuildByIDParams) (database.WorkspaceBuild, error) {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
 
@@ -2835,9 +2835,9 @@ func (q *fakeQuerier) UpdateWorkspaceBuildByID(_ context.Context, arg database.U
 		workspaceBuild.ProvisionerState = arg.ProvisionerState
 		workspaceBuild.Deadline = arg.Deadline
 		q.workspaceBuilds[index] = workspaceBuild
-		return nil
+		return workspaceBuild, nil
 	}
-	return sql.ErrNoRows
+	return database.WorkspaceBuild{}, sql.ErrNoRows
 }
 
 func (q *fakeQuerier) UpdateWorkspaceDeletedByID(_ context.Context, arg database.UpdateWorkspaceDeletedByIDParams) error {

diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspacebuilds.sql b/coderd/database/queries/workspacebuilds.sql
@@ -124,12 +124,12 @@ INSERT INTO
 VALUES
 	($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING *;
 
--- name: UpdateWorkspaceBuildByID :exec
+-- name: UpdateWorkspaceBuildByID :one
 UPDATE
 	workspace_builds
 SET
 	updated_at = $2,
 	provisioner_state = $3,
 	deadline = $4
 WHERE
-	id = $1;
+	id = $1 RETURNING *;
diff --git a/coderd/httpapi/httpapi.go b/coderd/httpapi/httpapi.go
@@ -228,14 +228,20 @@ func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (sendEvent f
 		buf := &bytes.Buffer{}
 		enc := json.NewEncoder(buf)
 
-		_, err := buf.WriteString(fmt.Sprintf("event: %s\ndata: ", sse.Type))
+		_, err := buf.WriteString(fmt.Sprintf("event: %s\n", sse.Type))
 		if err != nil {
 			return err
 		}
 
-		err = enc.Encode(sse.Data)
-		if err != nil {
-			return err
+		if sse.Data != nil {
+			_, err = buf.WriteString("data: ")
+			if err != nil {
+				return err
+			}
+			err = enc.Encode(sse.Data)
+			if err != nil {
+				return err
+			}
 		}
 
 		err = buf.WriteByte('\n')

diff --git a/coderd/provisionerdaemons.go b/coderd/provisionerdaemons.go
@@ -223,6 +223,10 @@ func (server *provisionerdServer) AcquireJob(ctx context.Context, _ *proto.Empty
 		if err != nil {
 			return nil, failJob(fmt.Sprintf("get owner: %s", err))
 		}
+		err = server.Pubsub.Publish(watchWorkspaceChannel(workspace.ID), []byte{})
+		if err != nil {
+			return nil, failJob(fmt.Sprintf("publish workspace update: %s", err))
+		}
 
 		// Compute parameters for the workspace to consume.
 		parameters, err := parameter.Compute(ctx, server.Database, parameter.ComputeScope{
@@ -543,7 +547,7 @@ func (server *provisionerdServer) FailJob(ctx context.Context, failJob *proto.Fa
 		if err != nil {
 			return nil, xerrors.Errorf("unmarshal workspace provision input: %w", err)
 		}
-		err = server.Database.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
+		build, err := server.Database.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
 			ID:               input.WorkspaceBuildID,
 			UpdatedAt:        database.Now(),
 			ProvisionerState: jobType.WorkspaceBuild.State,
@@ -552,6 +556,10 @@ func (server *provisionerdServer) FailJob(ctx context.Context, failJob *proto.Fa
 		if err != nil {
 			return nil, xerrors.Errorf("update workspace build state: %w", err)
 		}
+		err = server.Pubsub.Publish(watchWorkspaceChannel(build.WorkspaceID), []byte{})
+		if err != nil {
+			return nil, xerrors.Errorf("update workspace: %w", err)
+		}
 	case *proto.FailedJob_TemplateImport_:
 	}
 
@@ -657,7 +665,7 @@ func (server *provisionerdServer) CompleteJob(ctx context.Context, completed *pr
 			if err != nil {
 				return xerrors.Errorf("update provisioner job: %w", err)
 			}
-			err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
+			_, err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
 				ID:               workspaceBuild.ID,
 				Deadline:         workspaceDeadline,
 				ProvisionerState: jobType.WorkspaceBuild.State,
@@ -692,6 +700,11 @@ func (server *provisionerdServer) CompleteJob(ctx context.Context, completed *pr
 		if err != nil {
 			return nil, xerrors.Errorf("complete job: %w", err)
 		}
+
+		err = server.Pubsub.Publish(watchWorkspaceChannel(workspaceBuild.WorkspaceID), []byte{})
+		if err != nil {
+			return nil, xerrors.Errorf("update workspace: %w", err)
+		}
 	case *proto.CompletedJob_TemplateDryRun_:
 		for _, resource := range jobType.TemplateDryRun.Resources {
 			server.Logger.Info(ctx, "inserting template dry-run job resource",

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -539,13 +539,17 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
 			Valid: true,
 		}
 		_ = updateConnectionTimes()
+		_ = api.Pubsub.Publish(watchWorkspaceChannel(build.WorkspaceID), []byte{})
 	}()
 
 	err = updateConnectionTimes()
 	if err != nil {
 		_ = conn.Close(websocket.StatusGoingAway, err.Error())
 		return
 	}
+	if !api.publishWorkspaceUpdate(ctx, rw, build.WorkspaceID) {
+		return
+	}
 
 	// End span so we don't get long lived trace data.
 	tracing.EndHTTPSpan(r, http.StatusOK, trace.SpanFromContext(ctx))
@@ -972,6 +976,34 @@ func (api *API) postWorkspaceAppHealth(rw http.ResponseWriter, r *http.Request)
 		}
 	}
 
+	resource, err := api.Database.GetWorkspaceResourceByID(r.Context(), workspaceAgent.ResourceID)
+	if err != nil {
+		httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error fetching workspace resource.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	job, err := api.Database.GetWorkspaceBuildByJobID(r.Context(), resource.JobID)
+	if err != nil {
+		httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error fetching workspace build.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	workspace, err := api.Database.GetWorkspaceByID(r.Context(), job.WorkspaceID)
+	if err != nil {
+		httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error fetching workspace.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	if !api.publishWorkspaceUpdate(r.Context(), rw, workspace.ID) {
+		return
+	}
+
 	httpapi.Write(r.Context(), rw, http.StatusOK, nil)
 }
 

diff --git a/coderd/workspacebuilds.go b/coderd/workspacebuilds.go
@@ -574,6 +574,10 @@ func (api *API) postWorkspaceBuilds(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if !api.publishWorkspaceUpdate(ctx, rw, workspace.ID) {
+		return
+	}
+
 	httpapi.Write(ctx, rw, http.StatusCreated, apiBuild)
 }
 
@@ -632,6 +636,11 @@ func (api *API) patchCancelWorkspaceBuild(rw http.ResponseWriter, r *http.Reques
 		})
 		return
 	}
+
+	if !api.publishWorkspaceUpdate(ctx, rw, workspace.ID) {
+		return
+	}
+
 	httpapi.Write(ctx, rw, http.StatusOK, codersdk.Response{
 		Message: "Job has been marked as canceled...",
 	})

diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -634,6 +634,10 @@ func (api *API) patchWorkspace(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if !api.publishWorkspaceUpdate(ctx, rw, workspace.ID) {
+		return
+	}
+
 	aReq.New = newWorkspace
 	rw.WriteHeader(http.StatusNoContent)
 }
@@ -839,7 +843,7 @@ func (api *API) putExtendWorkspace(rw http.ResponseWriter, r *http.Request) {
 			return err
 		}
 
-		if err := s.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
+		if _, err := s.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
 			ID:               build.ID,
 			UpdatedAt:        build.UpdatedAt,
 			ProvisionerState: build.ProvisionerState,
@@ -883,48 +887,60 @@ func (api *API) watchWorkspace(rw http.ResponseWriter, r *http.Request) {
 	// Ignore all trace spans after this, they're not too useful.
 	ctx = trace.ContextWithSpan(ctx, tracing.NoopSpan)
 
-	t := time.NewTicker(time.Second * 1)
-	defer t.Stop()
+	cancelSubscribe, err := api.Pubsub.Subscribe(watchWorkspaceChannel(workspace.ID), func(_ context.Context, _ []byte) {
+		workspace, err := api.Database.GetWorkspaceByID(ctx, workspace.ID)
+		if err != nil {
+			_ = sendEvent(ctx, codersdk.ServerSentEvent{
+				Type: codersdk.ServerSentEventTypeError,
+				Data: codersdk.Response{
+					Message: "Internal error fetching workspace.",
+					Detail:  err.Error(),
+				},
+			})
+			return
+		}
+
+		data, err := api.workspaceData(ctx, []database.Workspace{workspace})
+		if err != nil {
+			_ = sendEvent(ctx, codersdk.ServerSentEvent{
+				Type: codersdk.ServerSentEventTypeError,
+				Data: codersdk.Response{
+					Message: "Internal error fetching workspace data.",
+					Detail:  err.Error(),
+				},
+			})
+			return
+		}
+
+		_ = sendEvent(ctx, codersdk.ServerSentEvent{
+			Type: codersdk.ServerSentEventTypeData,
+			Data: convertWorkspace(
+				workspace,
+				data.builds[0],
+				data.templates[0],
+				findUser(workspace.OwnerID, data.users),
+			),
+		})
+	})
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error subscribing to workspace events.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	defer cancelSubscribe()
+
+	_ = sendEvent(ctx, codersdk.ServerSentEvent{
+		Type: codersdk.ServerSentEventTypePing,
+	})
+
 	for {
 		select {
 		case <-ctx.Done():
 			return
 		case <-senderClosed:
 			return
-		case <-t.C:
-			workspace, err := api.Database.GetWorkspaceByID(ctx, workspace.ID)
-			if err != nil {
-				_ = sendEvent(ctx, codersdk.ServerSentEvent{
-					Type: codersdk.ServerSentEventTypeError,
-					Data: codersdk.Response{
-						Message: "Internal error fetching workspace.",
-						Detail:  err.Error(),
-					},
-				})
-				return
-			}
-
-			data, err := api.workspaceData(ctx, []database.Workspace{workspace})
-			if err != nil {
-				_ = sendEvent(ctx, codersdk.ServerSentEvent{
-					Type: codersdk.ServerSentEventTypeError,
-					Data: codersdk.Response{
-						Message: "Internal error fetching workspace data.",
-						Detail:  err.Error(),
-					},
-				})
-				return
-			}
-
-			_ = sendEvent(ctx, codersdk.ServerSentEvent{
-				Type: codersdk.ServerSentEventTypeData,
-				Data: convertWorkspace(
-					workspace,
-					data.builds[0],
-					data.templates[0],
-					findUser(workspace.OwnerID, data.users),
-				),
-			})
 		}
 	}
 }
@@ -1213,3 +1229,19 @@ func splitQueryParameterByDelimiter(query string, delimiter rune, maintainQuotes
 
 	return parts
 }
+
+func watchWorkspaceChannel(id uuid.UUID) string {
+	return fmt.Sprintf("workspace:%s", id)
+}
+
+func (api *API) publishWorkspaceUpdate(ctx context.Context, rw http.ResponseWriter, workspaceID uuid.UUID) bool {
+	err := api.Pubsub.Publish(watchWorkspaceChannel(workspaceID), []byte{})
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error publishing workspace update.",
+			Detail:  err.Error(),
+		})
+		return false
+	}
+	return true
+}