coder · johnstcn · Mar 20, 2024 · Mar 18, 2024 · Mar 18, 2024 · Mar 18, 2024
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -136,6 +136,8 @@ func (r *RootCmd) portForward() *serpent.Command {
 				listeners[i] = l
 			}
 
+			stopUpdating := client.UpdateWorkspaceUsageContext(ctx, workspace.ID)
+
 			// Wait for the context to be canceled or for a signal and close
 			// all listeners.
 			var closeErr error
@@ -156,6 +158,7 @@ func (r *RootCmd) portForward() *serpent.Command {
 				}
 
 				cancel()
+				stopUpdating()
 				closeAllListeners()
 			}()
 

diff --git a/cli/portforward_test.go b/cli/portforward_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/pty/ptytest"
 	"github.com/coder/coder/v2/testutil"
@@ -96,7 +97,12 @@ func TestPortForward(t *testing.T) {
 	// Setup agent once to be shared between test-cases (avoid expensive
 	// non-parallel setup).
 	var (
-		client, db         = coderdtest.NewWithDatabase(t, nil)
+		wuTick     = make(chan time.Time)
+		wuFlush    = make(chan int, 1)
+		client, db = coderdtest.NewWithDatabase(t, &coderdtest.Options{
+			WorkspaceUsageTrackerTick:  wuTick,
+			WorkspaceUsageTrackerFlush: wuFlush,
+		})
 		admin              = coderdtest.CreateFirstUser(t, client)
 		member, memberUser = coderdtest.CreateAnotherUser(t, client, admin.OrganizationID)
 		workspace          = runAgent(t, client, memberUser.ID, db)
@@ -148,6 +154,13 @@ func TestPortForward(t *testing.T) {
 			cancel()
 			err = <-errC
 			require.ErrorIs(t, err, context.Canceled)
+
+			flushCtx := testutil.Context(t, testutil.WaitShort)
+			testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
+			_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
+			updated, err := client.Workspace(context.Background(), workspace.ID)
+			require.NoError(t, err)
+			require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
 		})
 
 		t.Run(c.name+"_TwoPorts", func(t *testing.T) {
@@ -196,6 +209,13 @@ func TestPortForward(t *testing.T) {
 			cancel()
 			err = <-errC
 			require.ErrorIs(t, err, context.Canceled)
+
+			flushCtx := testutil.Context(t, testutil.WaitShort)
+			testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
+			_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
+			updated, err := client.Workspace(context.Background(), workspace.ID)
+			require.NoError(t, err)
+			require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
 		})
 	}
 
@@ -257,6 +277,13 @@ func TestPortForward(t *testing.T) {
 		cancel()
 		err := <-errC
 		require.ErrorIs(t, err, context.Canceled)
+
+		flushCtx := testutil.Context(t, testutil.WaitShort)
+		testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
+		_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
+		updated, err := client.Workspace(context.Background(), workspace.ID)
+		require.NoError(t, err)
+		require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
 	})
 }
 

diff --git a/cli/server.go b/cli/server.go
@@ -86,6 +86,7 @@ import (
 	stringutil "github.com/coder/coder/v2/coderd/util/strings"
 	"github.com/coder/coder/v2/coderd/workspaceapps"
 	"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
+	"github.com/coder/coder/v2/coderd/workspaceusage"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/drpc"
 	"github.com/coder/coder/v2/cryptorand"
@@ -968,6 +969,13 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 			purger := dbpurge.New(ctx, logger, options.Database)
 			defer purger.Close()
 
+			// Updates workspace usage
+			tracker := workspaceusage.New(options.Database,
+				workspaceusage.WithLogger(logger.Named("workspace_usage_tracker")),
+			)
+			options.WorkspaceUsageTracker = tracker
+			defer tracker.Close()
+
 			// Wrap the server in middleware that redirects to the access URL if
 			// the request is not to a local IP.
 			var handler http.Handler = coderAPI.RootHandler

diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -66,6 +66,7 @@ import (
 	"github.com/coder/coder/v2/coderd/updatecheck"
 	"github.com/coder/coder/v2/coderd/util/slice"
 	"github.com/coder/coder/v2/coderd/workspaceapps"
+	"github.com/coder/coder/v2/coderd/workspaceusage"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/drpc"
 	"github.com/coder/coder/v2/provisionerd/proto"
@@ -190,6 +191,9 @@ type Options struct {
 
 	// NewTicker is used for unit tests to replace "time.NewTicker".
 	NewTicker func(duration time.Duration) (tick <-chan time.Time, done func())
+
+	// WorkspaceUsageTracker tracks workspace usage by the CLI.
+	WorkspaceUsageTracker *workspaceusage.Tracker
 }
 
 // @title Coder API
@@ -362,6 +366,12 @@ func New(options *Options) *API {
 		OIDC:   options.OIDCConfig,
 	}
 
+	if options.WorkspaceUsageTracker == nil {
+		options.WorkspaceUsageTracker = workspaceusage.New(options.Database,
+			workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
+		)
+	}
+
 	ctx, cancel := context.WithCancel(context.Background())
 	r := chi.NewRouter()
 
@@ -405,6 +415,7 @@ func New(options *Options) *API {
 			options.Logger.Named("acquirer"),
 			options.Database,
 			options.Pubsub),
+		workspaceUsageTracker: options.WorkspaceUsageTracker,
 	}
 
 	api.AppearanceFetcher.Store(&appearance.DefaultFetcher)
@@ -972,6 +983,7 @@ func New(options *Options) *API {
 				})
 				r.Get("/watch", api.watchWorkspace)
 				r.Put("/extend", api.putExtendWorkspace)
+				r.Post("/usage", api.postWorkspaceUsage)
 				r.Put("/dormant", api.putWorkspaceDormant)
 				r.Put("/favorite", api.putFavoriteWorkspace)
 				r.Delete("/favorite", api.deleteFavoriteWorkspace)
@@ -1179,6 +1191,8 @@ type API struct {
 	statsBatcher *batchstats.Batcher
 
 	Acquirer *provisionerdserver.Acquirer
+
+	workspaceUsageTracker *workspaceusage.Tracker
 }
 
 // Close waits for all WebSocket connections to drain before returning.
@@ -1200,6 +1214,7 @@ func (api *API) Close() error {
 		_ = (*coordinator).Close()
 	}
 	_ = api.agentProvider.Close()
+	api.workspaceUsageTracker.Close()
 	return nil
 }
 

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -70,6 +70,7 @@ import (
 	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/coderd/workspaceapps"
 	"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
+	"github.com/coder/coder/v2/coderd/workspaceusage"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/codersdk/drpc"
@@ -146,6 +147,8 @@ type Options struct {
 	WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
 	AllowWorkspaceRenames              bool
 	NewTicker                          func(duration time.Duration) (<-chan time.Time, func())
+	WorkspaceUsageTrackerFlush         chan int
+	WorkspaceUsageTrackerTick          chan time.Time
 }
 
 // New constructs a codersdk client connected to an in-memory API instance.
@@ -306,6 +309,36 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 	hangDetector.Start()
 	t.Cleanup(hangDetector.Close)
 
+	// Did last_used_at not update? Scratching your noggin? Here's why.
+	// Workspace usage tracking must be triggered manually in tests.
+	// The vast majority of existing tests do not depend on last_used_at
+	// and adding an extra time-based background goroutine to all existing
+	// tests may lead to future flakes and goleak complaints.
+	// Instead, pass in your own flush and ticker like so:
+	//
+	//   tickCh = make(chan time.Time)
+	//   flushCh = make(chan int, 1)
+	//   client  = coderdtest.New(t, &coderdtest.Options{
+	//     WorkspaceUsageTrackerFlush: flushCh,
+	//     WorkspaceUsageTrackerTick: tickCh
+	//   })
+	//
+	// Now to trigger a tick, just write to `tickCh`.
+	// Reading from `flushCh` will ensure that workspaceusage.Tracker flushed.
+	// See TestPortForward or TestTracker_MultipleInstances for how this works in practice.
+	if options.WorkspaceUsageTrackerFlush == nil {
+		options.WorkspaceUsageTrackerFlush = make(chan int, 1) // buffering just in case
+	}
+	if options.WorkspaceUsageTrackerTick == nil {
+		options.WorkspaceUsageTrackerTick = make(chan time.Time, 1) // buffering just in case
+	}
+	// Close is called by API.Close()
+	wuTracker := workspaceusage.New(
+		options.Database,
+		workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
+		workspaceusage.WithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush),
+	)
+
 	var mutex sync.RWMutex
 	var handler http.Handler
 	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -454,6 +487,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			WorkspaceAppsStatsCollectorOptions: options.WorkspaceAppsStatsCollectorOptions,
 			AllowWorkspaceRenames:              options.AllowWorkspaceRenames,
 			NewTicker:                          options.NewTicker,
+			WorkspaceUsageTracker:              wuTracker,
 		}
 }
 

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -1046,6 +1046,10 @@ func (q *FakeQuerier) BatchUpdateWorkspaceLastUsedAt(_ context.Context, arg data
 		if _, found := m[q.workspaces[i].ID]; !found {
 			continue
 		}
+		// WHERE last_used_at < @last_used_at
+		if !q.workspaces[i].LastUsedAt.Before(arg.LastUsedAt) {
+			continue
+		}
 		q.workspaces[i].LastUsedAt = arg.LastUsedAt
 		n++
 	}

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
@@ -433,7 +433,10 @@ UPDATE
 SET
 	last_used_at = @last_used_at
 WHERE
-	id = ANY(@ids :: uuid[]);
+	id = ANY(@ids :: uuid[])
+AND
+  -- Do not overwrite with older data
+  last_used_at < @last_used_at;
 
 -- name: GetDeploymentWorkspaceStats :one
 WITH workspaces_with_jobs AS (

diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -1096,6 +1096,24 @@ func (api *API) putExtendWorkspace(rw http.ResponseWriter, r *http.Request) {
 	httpapi.Write(ctx, rw, code, resp)
 }
 
+// @Summary Post Workspace Usage by ID
+// @ID post-workspace-usage-by-id
+// @Security CoderSessionToken
+// @Tags Workspaces
+// @Param workspace path string true "Workspace ID" format(uuid)
+// @Success 204
+// @Router /workspaces/{workspace}/usage [post]
+func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
+	workspace := httpmw.WorkspaceParam(r)
+	if !api.Authorize(r, rbac.ActionUpdate, workspace) {
+		httpapi.Forbidden(rw)
+		return
+	}
+
+	api.workspaceUsageTracker.Add(workspace.ID)
+	rw.WriteHeader(http.StatusNoContent)
+}
+
 // @Summary Favorite workspace by ID.
 // @ID favorite-workspace-by-id
 // @Security CoderSessionToken