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,9 @@ 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/dbtestutil"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/workspaceusage"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/pty/ptytest"
 	"github.com/coder/coder/v2/testutil"
@@ -96,7 +99,15 @@ 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)
+		db, ps  = dbtestutil.NewDB(t)
+		wuTick  = make(chan time.Time)
+		wuFlush = make(chan int, 1)
+		wut     = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlush), workspaceusage.WithTickChannel(wuTick))
+		client  = coderdtest.New(t, &coderdtest.Options{
+			WorkspaceUsageTracker: wut,
+			Database:              db,
+			Pubsub:                ps,
+		})
 		admin              = coderdtest.CreateFirstUser(t, client)
 		member, memberUser = coderdtest.CreateAnotherUser(t, client, admin.OrganizationID)
 		workspace          = runAgent(t, client, memberUser.ID, db)
@@ -148,6 +159,12 @@ func TestPortForward(t *testing.T) {
 			cancel()
 			err = <-errC
 			require.ErrorIs(t, err, context.Canceled)
+
+			wuTick <- dbtime.Now()
+			<-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 +213,12 @@ func TestPortForward(t *testing.T) {
 			cancel()
 			err = <-errC
 			require.ErrorIs(t, err, context.Canceled)
+
+			wuTick <- dbtime.Now()
+			<-wuFlush
+			updated, err := client.Workspace(context.Background(), workspace.ID)
+			require.NoError(t, err)
+			require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
 		})
 	}
 
@@ -257,6 +280,12 @@ func TestPortForward(t *testing.T) {
 		cancel()
 		err := <-errC
 		require.ErrorIs(t, err, context.Canceled)
+
+		wuTick <- dbtime.Now()
+		<-wuFlush
+		updated, err := client.Workspace(context.Background(), workspace.ID)
+		require.NoError(t, err)
+		require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
 	})
 }
 

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,14 @@ func New(options *Options) *API {
 		OIDC:   options.OIDCConfig,
 	}
 
+	if options.WorkspaceUsageTracker == nil {
+		options.WorkspaceUsageTracker = workspaceusage.New(options.Database,
+			workspaceusage.WithFlushInterval(workspaceusage.DefaultFlushInterval),
+			workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
+		)
+	}
+	go options.WorkspaceUsageTracker.Loop()
+
 	ctx, cancel := context.WithCancel(context.Background())
 	r := chi.NewRouter()
 
@@ -405,6 +417,7 @@ func New(options *Options) *API {
 			options.Logger.Named("acquirer"),
 			options.Database,
 			options.Pubsub),
+		workspaceUsageTracker: options.WorkspaceUsageTracker,
 	}
 
 	api.AppearanceFetcher.Store(&appearance.DefaultFetcher)
@@ -972,6 +985,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 +1193,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 +1216,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,7 @@ type Options struct {
 	WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
 	AllowWorkspaceRenames              bool
 	NewTicker                          func(duration time.Duration) (<-chan time.Time, func())
+	WorkspaceUsageTracker              *workspaceusage.Tracker
 }
 
 // New constructs a codersdk client connected to an in-memory API instance.
@@ -306,6 +308,36 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 	hangDetector.Start()
 	t.Cleanup(hangDetector.Close)
 
+	if options.WorkspaceUsageTracker == nil {
+		// 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 background goroutine to all existing tests may
+		// lead to future flakes and goleak complaints.
+		// To do this, pass in your own WorkspaceUsageTracker like so:
+		//
+		// 	db, ps  = dbtestutil.NewDB(t)
+		//   wuTick  = make(chan time.Time)
+		//   wuFlush = make(chan int, 1)
+		//   wut     = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlush), workspaceusage.WithTickChannel(wuTick))
+		//   client  = coderdtest.New(t, &coderdtest.Options{
+		//     WorkspaceUsageTracker: wut,
+		//     Database:              db,
+		//     Pubsub:                ps,
+		//   })
+		//
+		// See TestPortForward for how this works in practice.
+		wutFlush := make(chan int)
+		wutTick := make(chan time.Time)
+		options.WorkspaceUsageTracker = workspaceusage.New(
+			options.Database,
+			workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
+			workspaceusage.WithFlushChannel(wutFlush),
+			workspaceusage.WithTickChannel(wutTick),
+		)
+	}
+	t.Cleanup(options.WorkspaceUsageTracker.Close)
+
 	var mutex sync.RWMutex
 	var handler http.Handler
 	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -454,6 +486,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			WorkspaceAppsStatsCollectorOptions: options.WorkspaceAppsStatsCollectorOptions,
 			AllowWorkspaceRenames:              options.AllowWorkspaceRenames,
 			NewTicker:                          options.NewTicker,
+			WorkspaceUsageTracker:              options.WorkspaceUsageTracker,
 		}
 }
 

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