coder · kylecarbs · Nov 7, 2022 · Nov 6, 2022 · Nov 7, 2022 · mafredri
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,15 @@ 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
 	}
+	api.publishWorkspaceUpdate(ctx, build.WorkspaceID)
 
 	// End span so we don't get long lived trace data.
 	tracing.EndHTTPSpan(r, http.StatusOK, trace.SpanFromContext(ctx))
@@ -972,6 +974,32 @@ 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
+	}
+	api.publishWorkspaceUpdate(r.Context(), workspace.ID)
+
 	httpapi.Write(r.Context(), rw, http.StatusOK, nil)
 }
 

diff --git a/coderd/workspacebuilds.go b/coderd/workspacebuilds.go
@@ -574,6 +574,8 @@ func (api *API) postWorkspaceBuilds(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	api.publishWorkspaceUpdate(ctx, workspace.ID)
+
 	httpapi.Write(ctx, rw, http.StatusCreated, apiBuild)
 }
 
@@ -632,6 +634,9 @@ func (api *API) patchCancelWorkspaceBuild(rw http.ResponseWriter, r *http.Reques
 		})
 		return
 	}
+
+	api.publishWorkspaceUpdate(ctx, workspace.ID)
+
 	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,8 @@ func (api *API) patchWorkspace(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	api.publishWorkspaceUpdate(ctx, workspace.ID)
+
 	aReq.New = newWorkspace
 	rw.WriteHeader(http.StatusNoContent)
 }
@@ -839,7 +841,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 +885,65 @@ 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 {
+		_ = sendEvent(ctx, codersdk.ServerSentEvent{
+			Type: codersdk.ServerSentEventTypeError,
+			Data: codersdk.Response{
+				Message: "Internal error subscribing to workspace events.",
+				Detail:  err.Error(),
+			},
+		})
+		return
+	}
+	defer cancelSubscribe()
+
+	// An initial ping signals to the request that the server is now ready
+	// and the client can begin servicing a channel with data.
+	_ = 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 +1232,15 @@ 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, workspaceID uuid.UUID) {
+	err := api.Pubsub.Publish(watchWorkspaceChannel(workspaceID), []byte{})
+	if err != nil {
+		api.Logger.Warn(ctx, "failed to publish workspace update",
+			slog.F("workspace_id", workspaceID), slog.Error(err))
+	}
+}