coder · aslilac · Feb 13, 2024 · Feb 7, 2024 · Feb 7, 2024 · Feb 7, 2024
diff --git a/coderd/agentapi/api.go b/coderd/agentapi/api.go
@@ -114,6 +114,7 @@ func New(opts Options) *API {
 	api.StatsAPI = &StatsAPI{
 		AgentFn:                   api.agent,
 		Database:                  opts.Database,
+		Pubsub:                    opts.Pubsub,
 		Log:                       opts.Log,
 		StatsBatcher:              opts.StatsBatcher,
 		TemplateScheduleStore:     opts.TemplateScheduleStore,

diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go
@@ -16,8 +16,10 @@ import (
 	"github.com/coder/coder/v2/coderd/autobuild"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 	"github.com/coder/coder/v2/coderd/schedule"
+	"github.com/coder/coder/v2/codersdk"
 )
 
 type StatsBatcher interface {
@@ -27,6 +29,7 @@ type StatsBatcher interface {
 type StatsAPI struct {
 	AgentFn                   func(context.Context) (database.WorkspaceAgent, error)
 	Database                  database.Store
+	Pubsub                    pubsub.Pubsub
 	Log                       slog.Logger
 	StatsBatcher              StatsBatcher
 	TemplateScheduleStore     *atomic.Pointer[schedule.TemplateScheduleStore]
@@ -130,5 +133,16 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 		return nil, xerrors.Errorf("update stats in database: %w", err)
 	}
 
+	// Tell the frontend about the new agent report, now that everything is updated
+	a.publishWorkspaceAgentStats(ctx, workspace.ID)
+
 	return res, nil
 }
+
+func (a *StatsAPI) publishWorkspaceAgentStats(ctx context.Context, workspaceID uuid.UUID) {
+	err := a.Pubsub.Publish(codersdk.WorkspaceNotifyChannel(workspaceID), codersdk.WorkspaceNotifyDescriptionAgentStatsOnly)
+	if err != nil {
+		a.Log.Warn(ctx, "failed to publish workspace agent stats",
+			slog.F("workspace_id", workspaceID), slog.Error(err))
+	}
+}
diff --git a/coderd/agentapi/stats_test.go b/coderd/agentapi/stats_test.go
@@ -19,6 +19,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbmock"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 	"github.com/coder/coder/v2/coderd/schedule"
 )
@@ -78,8 +79,10 @@ func TestUpdateStates(t *testing.T) {
 		t.Parallel()
 
 		var (
-			now                   = dbtime.Now()
-			dbM                   = dbmock.NewMockStore(gomock.NewController(t))
+			now = dbtime.Now()
+			db  = dbmock.NewMockStore(gomock.NewController(t))
+			ps  = pubsub.NewInMemory()
+
 			templateScheduleStore = schedule.MockTemplateScheduleStore{
 				GetFn: func(context.Context, database.Store, uuid.UUID) (schedule.TemplateScheduleOptions, error) {
 					panic("should not be called")
@@ -124,7 +127,8 @@ func TestUpdateStates(t *testing.T) {
 			AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			Database:                  dbM,
+			Database:                  db,
+			Pubsub:                    ps,
 			StatsBatcher:              batcher,
 			TemplateScheduleStore:     templateScheduleStorePtr(templateScheduleStore),
 			AgentStatsRefreshInterval: 10 * time.Second,
@@ -144,25 +148,25 @@ func TestUpdateStates(t *testing.T) {
 		}
 
 		// Workspace gets fetched.
-		dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
+		db.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
 			Workspace:    workspace,
 			TemplateName: template.Name,
 		}, nil)
 
 		// We expect an activity bump because ConnectionCount > 0.
-		dbM.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{
+		db.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{
 			WorkspaceID:   workspace.ID,
 			NextAutostart: time.Time{}.UTC(),
 		}).Return(nil)
 
 		// Workspace last used at gets bumped.
-		dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
+		db.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
 			ID:         workspace.ID,
 			LastUsedAt: now,
 		}).Return(nil)
 
 		// User gets fetched to hit the UpdateAgentMetricsFn.
-		dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
+		db.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
 
 		resp, err := api.UpdateStats(context.Background(), req)
 		require.NoError(t, err)
@@ -188,7 +192,8 @@ func TestUpdateStates(t *testing.T) {
 
 		var (
 			now                   = dbtime.Now()
-			dbM                   = dbmock.NewMockStore(gomock.NewController(t))
+			db                    = dbmock.NewMockStore(gomock.NewController(t))
+			ps                    = pubsub.NewInMemory()
 			templateScheduleStore = schedule.MockTemplateScheduleStore{
 				GetFn: func(context.Context, database.Store, uuid.UUID) (schedule.TemplateScheduleOptions, error) {
 					panic("should not be called")
@@ -213,7 +218,8 @@ func TestUpdateStates(t *testing.T) {
 			AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			Database:                  dbM,
+			Database:                  db,
+			Pubsub:                    ps,
 			StatsBatcher:              batcher,
 			TemplateScheduleStore:     templateScheduleStorePtr(templateScheduleStore),
 			AgentStatsRefreshInterval: 10 * time.Second,
@@ -225,13 +231,13 @@ func TestUpdateStates(t *testing.T) {
 		}
 
 		// Workspace gets fetched.
-		dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
+		db.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
 			Workspace:    workspace,
 			TemplateName: template.Name,
 		}, nil)
 
 		// Workspace last used at gets bumped.
-		dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
+		db.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
 			ID:         workspace.ID,
 			LastUsedAt: now,
 		}).Return(nil)
@@ -244,7 +250,8 @@ func TestUpdateStates(t *testing.T) {
 		t.Parallel()
 
 		var (
-			dbM = dbmock.NewMockStore(gomock.NewController(t))
+			db  = dbmock.NewMockStore(gomock.NewController(t))
+			ps  = pubsub.NewInMemory()
 			req = &agentproto.UpdateStatsRequest{
 				Stats: &agentproto.Stats{
 					ConnectionsByProto: map[string]int64{}, // len() == 0
@@ -255,7 +262,8 @@ func TestUpdateStates(t *testing.T) {
 			AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			Database:                  dbM,
+			Database:                  db,
+			Pubsub:                    ps,
 			StatsBatcher:              nil, // should not be called
 			TemplateScheduleStore:     nil, // should not be called
 			AgentStatsRefreshInterval: 10 * time.Second,
@@ -290,7 +298,9 @@ func TestUpdateStates(t *testing.T) {
 		nextAutostart := now.Add(30 * time.Minute).UTC() // always sent to DB as UTC
 
 		var (
-			dbM                   = dbmock.NewMockStore(gomock.NewController(t))
+			db = dbmock.NewMockStore(gomock.NewController(t))
+			ps = pubsub.NewInMemory()
+
 			templateScheduleStore = schedule.MockTemplateScheduleStore{
 				GetFn: func(context.Context, database.Store, uuid.UUID) (schedule.TemplateScheduleOptions, error) {
 					return schedule.TemplateScheduleOptions{
@@ -321,7 +331,8 @@ func TestUpdateStates(t *testing.T) {
 			AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			Database:                  dbM,
+			Database:                  db,
+			Pubsub:                    ps,
 			StatsBatcher:              batcher,
 			TemplateScheduleStore:     templateScheduleStorePtr(templateScheduleStore),
 			AgentStatsRefreshInterval: 15 * time.Second,
@@ -341,26 +352,26 @@ func TestUpdateStates(t *testing.T) {
 		}
 
 		// Workspace gets fetched.
-		dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
+		db.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
 			Workspace:    workspace,
 			TemplateName: template.Name,
 		}, nil)
 
 		// We expect an activity bump because ConnectionCount > 0. However, the
 		// next autostart time will be set on the bump.
-		dbM.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{
+		db.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{
 			WorkspaceID:   workspace.ID,
 			NextAutostart: nextAutostart,
 		}).Return(nil)
 
 		// Workspace last used at gets bumped.
-		dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
+		db.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
 			ID:         workspace.ID,
 			LastUsedAt: now,
 		}).Return(nil)
 
 		// User gets fetched to hit the UpdateAgentMetricsFn.
-		dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
+		db.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
 
 		resp, err := api.UpdateStats(context.Background(), req)
 		require.NoError(t, err)

diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -1,6 +1,7 @@
 package coderd
 
 import (
+	"bytes"
 	"context"
 	"database/sql"
 	"encoding/json"
@@ -1343,7 +1344,48 @@ func (api *API) watchWorkspace(rw http.ResponseWriter, r *http.Request) {
 		<-senderClosed
 	}()
 
-	sendUpdate := func(_ context.Context, _ []byte) {
+	sendUpdate := func(_ context.Context, description []byte) {
+		// The agent stats get updated frequently, so we treat these as a special case and only
+		// send a partial update. We primarily care about updating the `last_used_at` and
+		// `latest_build.deadline` properties.
+		if bytes.Equal(description, codersdk.WorkspaceNotifyDescriptionAgentStatsOnly) {
+			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
+			}
+
+			workspaceBuild, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspace.ID)
+			if err != nil {
+				_ = sendEvent(ctx, codersdk.ServerSentEvent{
+					Type: codersdk.ServerSentEventTypeError,
+					Data: codersdk.Response{
+						Message: "Internal error fetching workspace build.",
+						Detail:  err.Error(),
+					},
+				})
+				return
+			}
+
+			_ = sendEvent(ctx, codersdk.ServerSentEvent{
+				Type: codersdk.ServerSentEventTypePartial,
+				Data: struct {
+					database.Workspace
+					LatestBuild database.WorkspaceBuild `json:"latest_build"`
+				}{
+					Workspace:   workspace,
+					LatestBuild: workspaceBuild,
+				},
+			})
+			return
+		}
+
 		workspace, err := api.Database.GetWorkspaceByID(ctx, workspace.ID)
 		if err != nil {
 			_ = sendEvent(ctx, codersdk.ServerSentEvent{

diff --git a/codersdk/serversentevents.go b/codersdk/serversentevents.go
@@ -20,9 +20,10 @@ type ServerSentEvent struct {
 type ServerSentEventType string
 
 const (
-	ServerSentEventTypePing  ServerSentEventType = "ping"
-	ServerSentEventTypeData  ServerSentEventType = "data"
-	ServerSentEventTypeError ServerSentEventType = "error"
+	ServerSentEventTypePing    ServerSentEventType = "ping"
+	ServerSentEventTypeData    ServerSentEventType = "data"
+	ServerSentEventTypePartial ServerSentEventType = "partial"
+	ServerSentEventTypeError   ServerSentEventType = "error"
 )
 
 func ServerSentEventReader(ctx context.Context, rc io.ReadCloser) func() (*ServerSentEvent, error) {

diff --git a/codersdk/workspaces.go b/codersdk/workspaces.go
@@ -497,6 +497,8 @@ func (c *Client) UnfavoriteWorkspace(ctx context.Context, workspaceID uuid.UUID)
 	return nil
 }
 
+var WorkspaceNotifyDescriptionAgentStatsOnly = []byte("agentStatsOnly")
+
 // WorkspaceNotifyChannel is the PostgreSQL NOTIFY
 // channel to listen for updates on. The payload is empty,
 // because the size of a workspace payload can be very large.

diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts
diff --git a/site/src/hooks/useTime.ts b/site/src/hooks/useTime.ts
@@ -0,0 +1,32 @@
+import { useEffect, useState } from "react";
+
+/**
+ * useTime allows a component to rerender over time without a corresponding state change.
+ * An example could be a relative timestamp (eg. "in 5 minutes") that should count down as it
+ * approaches.
+ *
+ * This hook should only be used in components that are very simple, and that will not
+ * create a lot of unnecessary work for the reconciler. Given that this hook will result in
+ * the entire subtree being rerendered on a frequent interval, it's important that the subtree
+ * remains small.
+ *
+ * @param active Can optionally be set to false in circumstances where updating over time is
+ * not necessary.
+ */
+export function useTime(active: boolean = true) {
+  const [, setTick] = useState(0);
+
+  useEffect(() => {
+    if (!active) {
+      return;
+    }
+
+    const interval = setInterval(() => {
+      setTick((i) => i + 1);
+    }, 1000);
+
+    return () => {
+      clearInterval(interval);
+    };
+  }, [active]);
+}
diff --git a/site/src/modules/workspaces/activity.ts b/site/src/modules/workspaces/activity.ts
@@ -0,0 +1,45 @@
+import dayjs from "dayjs";
+import type { Workspace } from "api/typesGenerated";
+
+export type WorkspaceActivityStatus =
+  | "ready"
+  | "connected"
+  | "inactive"
+  | "notConnected"
+  | "notRunning";
+
+export function getWorkspaceActivityStatus(
+  workspace: Workspace,
+): WorkspaceActivityStatus {
+  const builtAt = dayjs(workspace.latest_build.created_at);
+  const usedAt = dayjs(workspace.last_used_at);
+  const now = dayjs();
+
+  if (workspace.latest_build.status !== "running") {
+    return "notRunning";
+  }
+
+  // This needs to compare to `usedAt` instead of `now`, because the "grace period" for
+  // marking a workspace as "Connected" is a lot longer. If you compared `builtAt` to `now`,
+  // you could end up switching from "Ready" to "Connected" without ever actually connecting.
+  const isBuiltRecently = builtAt.isAfter(usedAt.subtract(1, "second"));
+  // By default, agents report connection stats every 30 seconds, so 2 minutes should be
+  // plenty. Disconnection will be reflected relatively-quickly
+  const isUsedRecently = usedAt.isAfter(now.subtract(2, "minute"));
+
+  // If the build is still "fresh", it'll be a while before the `last_used_at` gets bumped in
+  // a significant way by the agent, so just label it as ready instead of connected.
+  // Wait until `last_used_at` is after the time that the build finished, _and_ still
+  // make sure to check that it's recent, so that we don't show "Ready" indefinitely.
+  if (isUsedRecently && isBuiltRecently && workspace.health.healthy) {
+    return "ready";
+  }
+
+  if (isUsedRecently) {
+    return "connected";
+  }
+
+  // TODO: It'd be nice if we could differentiate between "connected but inactive" and
+  // "not connected", but that will require some relatively substantial backend work.
+  return "inactive";
+}