coder · johnstcn · Jan 16, 2024 · Jan 12, 2024 · Jan 15, 2024 · Jan 15, 2024
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -695,6 +695,15 @@ func (q *querier) ArchiveUnusedTemplateVersions(ctx context.Context, arg databas
 	return q.db.ArchiveUnusedTemplateVersions(ctx, arg)
 }
 
+func (q *querier) BatchUpdateWorkspaceLastUsedAt(ctx context.Context, arg database.BatchUpdateWorkspaceLastUsedAtParams) error {
+	// Could be any workspace and checking auth to each workspace is overkill for the purpose
+	// of this function.
+	if err := q.authorizeContext(ctx, rbac.ActionUpdate, rbac.ResourceWorkspace.All()); err != nil {
+		return err
+	}
+	return q.db.BatchUpdateWorkspaceLastUsedAt(ctx, arg)
+}
+
 func (q *querier) CleanTailnetCoordinators(ctx context.Context) error {
 	if err := q.authorizeContext(ctx, rbac.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
 		return err

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -1549,6 +1549,13 @@ func (s *MethodTestSuite) TestWorkspace() {
 			ID: ws.ID,
 		}).Asserts(ws, rbac.ActionUpdate).Returns()
 	}))
+	s.Run("BatchUpdateWorkspaceLastUsedAt", s.Subtest(func(db database.Store, check *expects) {
+		ws1 := dbgen.Workspace(s.T(), db, database.Workspace{})
+		ws2 := dbgen.Workspace(s.T(), db, database.Workspace{})
+		check.Args(database.BatchUpdateWorkspaceLastUsedAtParams{
+			IDs: []uuid.UUID{ws1.ID, ws2.ID},
+		}).Asserts(rbac.ResourceWorkspace.All(), rbac.ActionUpdate).Returns()
+	}))
 	s.Run("UpdateWorkspaceTTL", s.Subtest(func(db database.Store, check *expects) {
 		ws := dbgen.Workspace(s.T(), db, database.Workspace{})
 		check.Args(database.UpdateWorkspaceTTLParams{

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -963,6 +963,31 @@ func (q *FakeQuerier) ArchiveUnusedTemplateVersions(_ context.Context, arg datab
 	return archived, nil
 }
 
+func (q *FakeQuerier) BatchUpdateWorkspaceLastUsedAt(_ context.Context, arg database.BatchUpdateWorkspaceLastUsedAtParams) error {
+	err := validateDatabaseType(arg)
+	if err != nil {
+		return err
+	}
+
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	// temporary map to avoid O(q.workspaces*arg.workspaceIds)
+	m := make(map[uuid.UUID]struct{})
+	for _, id := range arg.IDs {
+		m[id] = struct{}{}
+	}
+	n := 0
+	for i := 0; i < len(q.workspaces); i++ {
+		if _, found := m[q.workspaces[i].ID]; !found {
+			continue
+		}
+		q.workspaces[i].LastUsedAt = arg.LastUsedAt
+		n++
+	}
+	return nil
+}
+
 func (*FakeQuerier) CleanTailnetCoordinators(_ context.Context) error {
 	return ErrUnimplemented
 }

diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
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/workspaces.sql b/coderd/database/queries/workspaces.sql
@@ -357,6 +357,14 @@ SET
 WHERE
 	id = $1;
 
+-- name: BatchUpdateWorkspaceLastUsedAt :exec
+UPDATE
+	workspaces
+SET
+	last_used_at = @last_used_at
+WHERE
+	id = ANY(@ids :: uuid[]);
+
 -- name: GetDeploymentWorkspaceStats :one
 WITH workspaces_with_jobs AS (
 	SELECT

diff --git a/coderd/httpapi/websocket.go b/coderd/httpapi/websocket.go
@@ -4,8 +4,9 @@ import (
 	"context"
 	"time"
 
-	"cdr.dev/slog"
 	"nhooyr.io/websocket"
+
+	"cdr.dev/slog"
 )
 
 // Heartbeat loops to ping a WebSocket to keep it alive.

diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -1626,7 +1626,7 @@ func TestWorkspaceAgentExternalAuthListen(t *testing.T) {
 		cancel()
 		// We expect only 1
 		// In a failed test, you will likely see 9, as the last one
-		// gets cancelled.
+		// gets canceled.
 		require.Equal(t, 1, validateCalls, "validate calls duplicated on same token")
 	})
 }
diff --git a/coderd/workspaceapps/apptest/apptest.go b/coderd/workspaceapps/apptest/apptest.go
@@ -64,6 +64,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			// reconnecting-pty proxy server we want to test is mounted.
 			client := appDetails.AppClient(t)
 			testReconnectingPTY(ctx, t, client, appDetails.Agent.ID, "")
+			assertWorkspaceLastUsedAtUpdated(t, appDetails)
 		})
 
 		t.Run("SignedTokenQueryParameter", func(t *testing.T) {
@@ -92,6 +93,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			// Make an unauthenticated client.
 			unauthedAppClient := codersdk.New(appDetails.AppClient(t).URL)
 			testReconnectingPTY(ctx, t, unauthedAppClient, appDetails.Agent.ID, issueRes.SignedToken)
+			assertWorkspaceLastUsedAtUpdated(t, appDetails)
 		})
 	})
 
@@ -117,6 +119,9 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			body, err := io.ReadAll(resp.Body)
 			require.NoError(t, err)
 			require.Contains(t, string(body), "Path-based applications are disabled")
+			// Even though path-based apps are disabled, the request should indicate
+			// that the workspace was used.
+			assertWorkspaceLastUsedAtNotUpdated(t, appDetails)
 		})
 
 		t.Run("LoginWithoutAuthOnPrimary", func(t *testing.T) {
@@ -142,6 +147,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			require.NoError(t, err)
 			require.True(t, loc.Query().Has("message"))
 			require.True(t, loc.Query().Has("redirect"))
+			assertWorkspaceLastUsedAtUpdated(t, appDetails)
 		})
 
 		t.Run("LoginWithoutAuthOnProxy", func(t *testing.T) {
@@ -179,6 +185,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			// request is getting stripped.
 			require.Equal(t, u.Path, redirectURI.Path+"/")
 			require.Equal(t, u.RawQuery, redirectURI.RawQuery)
+			assertWorkspaceLastUsedAtUpdated(t, appDetails)
 		})
 
 		t.Run("NoAccessShould404", func(t *testing.T) {
@@ -195,6 +202,8 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			require.NoError(t, err)
 			defer resp.Body.Close()
 			require.Equal(t, http.StatusNotFound, resp.StatusCode)
+			// TODO(cian): A blocked request should not count as workspace usage.
+			// assertWorkspaceLastUsedAtNotUpdated(t, appDetails.AppClient(t), appDetails)
 		})
 
 		t.Run("RedirectsWithSlash", func(t *testing.T) {
@@ -209,6 +218,8 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			require.NoError(t, err)
 			defer resp.Body.Close()
 			require.Equal(t, http.StatusTemporaryRedirect, resp.StatusCode)
+			// TODO(cian): The initial redirect should not count as workspace usage.
+			// assertWorkspaceLastUsedAtNotUpdated(t, appDetails.AppClient(t), appDetails)
 		})
 
 		t.Run("RedirectsWithQuery", func(t *testing.T) {
@@ -226,6 +237,8 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			loc, err := resp.Location()
 			require.NoError(t, err)
 			require.Equal(t, proxyTestAppQuery, loc.RawQuery)
+			// TODO(cian): The initial redirect should not count as workspace usage.
+			// assertWorkspaceLastUsedAtNotUpdated(t, appDetails.AppClient(t), appDetails)
 		})
 
 		t.Run("Proxies", func(t *testing.T) {
@@ -267,6 +280,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			require.NoError(t, err)
 			require.Equal(t, proxyTestAppBody, string(body))
 			require.Equal(t, http.StatusOK, resp.StatusCode)
+			assertWorkspaceLastUsedAtUpdated(t, appDetails)
 		})
 
 		t.Run("ProxiesHTTPS", func(t *testing.T) {
@@ -312,6 +326,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			require.NoError(t, err)
 			require.Equal(t, proxyTestAppBody, string(body))
 			require.Equal(t, http.StatusOK, resp.StatusCode)
+			assertWorkspaceLastUsedAtUpdated(t, appDetails)
 		})
 
 		t.Run("BlocksMe", func(t *testing.T) {
@@ -331,6 +346,8 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			body, err := io.ReadAll(resp.Body)
 			require.NoError(t, err)
 			require.Contains(t, string(body), "must be accessed with the full username, not @me")
+			// TODO(cian): A blocked request should not count as workspace usage.
+			// assertWorkspaceLastUsedAtNotUpdated(t, appDetails.AppClient(t), appDetails)
 		})
 
 		t.Run("ForwardsIP", func(t *testing.T) {
@@ -349,6 +366,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			require.Equal(t, proxyTestAppBody, string(body))
 			require.Equal(t, http.StatusOK, resp.StatusCode)
 			require.Equal(t, "1.1.1.1,127.0.0.1", resp.Header.Get("X-Forwarded-For"))
+			assertWorkspaceLastUsedAtUpdated(t, appDetails)
 		})
 
 		t.Run("ProxyError", func(t *testing.T) {
@@ -361,6 +379,9 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			require.NoError(t, err)
 			defer resp.Body.Close()
 			require.Equal(t, http.StatusBadGateway, resp.StatusCode)
+			// An valid authenticated attempt to access a workspace app
+			// should count as usage regardless of success.
+			assertWorkspaceLastUsedAtUpdated(t, appDetails)
 		})
 
 		t.Run("NoProxyPort", func(t *testing.T) {
@@ -375,6 +396,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 			// TODO(@deansheather): This should be 400. There's a todo in the
 			// resolve request code to fix this.
 			require.Equal(t, http.StatusInternalServerError, resp.StatusCode)
+			assertWorkspaceLastUsedAtUpdated(t, appDetails)
 		})
 	})
 
@@ -1430,16 +1452,12 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 	t.Run("ReportStats", func(t *testing.T) {
 		t.Parallel()
 
-		flush := make(chan chan<- struct{}, 1)
-
 		reporter := &fakeStatsReporter{}
 		appDetails := setupProxyTest(t, &DeploymentOptions{
 			StatsCollectorOptions: workspaceapps.StatsCollectorOptions{
 				Reporter:       reporter,
 				ReportInterval: time.Hour,
 				RollupWindow:   time.Minute,
-
-				Flush: flush,
 			},
 		})
 
@@ -1457,10 +1475,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 		var stats []workspaceapps.StatsReport
 		require.Eventually(t, func() bool {
 			// Keep flushing until we get a non-empty stats report.
-			flushDone := make(chan struct{}, 1)
-			flush <- flushDone
-			<-flushDone
-
+			appDetails.FlushStats()
 			stats = reporter.stats()
 			return len(stats) > 0
 		}, testutil.WaitLong, testutil.IntervalFast, "stats not reported")
@@ -1549,3 +1564,28 @@ func testReconnectingPTY(ctx context.Context, t *testing.T, client *codersdk.Cli
 	// Ensure the connection closes.
 	require.ErrorIs(t, tr.ReadUntil(ctx, nil), io.EOF)
 }
+
+// Accessing an app should update the workspace's LastUsedAt.
+// NOTE: Despite our efforts with the flush channel, this is inherently racy.
+func assertWorkspaceLastUsedAtUpdated(t testing.TB, details *Details) {
+	t.Helper()
+
+	// Wait for stats to fully flush.
+	require.Eventually(t, func() bool {
+		details.FlushStats()
+		ws, err := details.SDKClient.Workspace(context.Background(), details.Workspace.ID)
+		assert.NoError(t, err)
+		return ws.LastUsedAt.After(details.Workspace.LastUsedAt)
+	}, testutil.WaitShort, testutil.IntervalMedium, "workspace LastUsedAt not updated when it should have been")
+}
+
+// Except when it sometimes shouldn't (e.g. no access)
+// NOTE: Despite our efforts with the flush channel, this is inherently racy.
+func assertWorkspaceLastUsedAtNotUpdated(t testing.TB, details *Details) {
+	t.Helper()
+
+	details.FlushStats()
+	ws, err := details.SDKClient.Workspace(context.Background(), details.Workspace.ID)
+	require.NoError(t, err)
+	require.Equal(t, ws.LastUsedAt, details.Workspace.LastUsedAt, "workspace LastUsedAt updated when it should not have been")
+}
diff --git a/coderd/workspaceapps/apptest/setup.go b/coderd/workspaceapps/apptest/setup.go
@@ -71,6 +71,7 @@ type Deployment struct {
 	SDKClient      *codersdk.Client
 	FirstUser      codersdk.CreateFirstUserResponse
 	PathAppBaseURL *url.URL
+	FlushStats     func()
 }
 
 // DeploymentFactory generates a deployment with an API client, a path base URL,

diff --git a/coderd/workspaceapps/stats.go b/coderd/workspaceapps/stats.go
@@ -13,6 +13,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/util/slice"
 )
 
 const (
@@ -117,11 +118,25 @@ func (r *StatsDBReporter) Report(ctx context.Context, stats []StatsReport) error
 				batch.Requests = batch.Requests[:0]
 			}
 		}
-		if len(batch.UserID) > 0 {
-			err := tx.InsertWorkspaceAppStats(ctx, batch)
-			if err != nil {
-				return err
-			}
+		if len(batch.UserID) == 0 {
+			return nil
+		}
+
+		if err := tx.InsertWorkspaceAppStats(ctx, batch); err != nil {
+			return err
+		}
+
+		// TODO: We currently measure workspace usage based on when we get stats from it.
+		// There are currently two paths for this:
+		// 1) From SSH -> workspace agent stats POSTed from agent
+		// 2) From workspace apps / rpty -> workspace app stats (from coderd / wsproxy)
+		// Ideally we would have a single code path for this.
+		uniqueIDs := slice.Unique(batch.WorkspaceID)
+		if err := tx.BatchUpdateWorkspaceLastUsedAt(ctx, database.BatchUpdateWorkspaceLastUsedAtParams{
+			IDs:        uniqueIDs,
+			LastUsedAt: dbtime.Now(), // This isn't 100% accurate, but it's good enough.
+		}); err != nil {
+			return err
 		}
 
 		return nil
@@ -234,6 +249,7 @@ func (sc *StatsCollector) Collect(report StatsReport) {
 		}
 		delete(sc.statsBySessionID, report.SessionID)
 	}
+	sc.opts.Logger.Debug(sc.ctx, "collected workspace app stats", slog.F("report", report))
 }
 
 // rollup performs stats rollup for sessions that fall within the