Mar 20, 2024 · Mar 18, 2024 · Mar 18, 2024 · Mar 18, 2024 · Mar 19, 2024 · Mar 19, 2024
diff --git a/cli/portforward.go b/cli/portforward.go
 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
 }

 cancel()
 stopUpdating()
 closeAllListeners()
 }()

diff --git a/cli/portforward_test.go b/cli/portforward_test.go
 "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"
 // 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)
 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) {
 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)
 })
 }

 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
 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"
 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
 "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"

 // 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
 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()

 options.Logger.Named("acquirer"),
 options.Database,
 options.Pubsub),
 workspaceUsageTracker: options.WorkspaceUsageTracker,
 }

 api.AppearanceFetcher.Store(&appearance.DefaultFetcher)
 })
 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)
 statsBatcher *batchstats.Batcher

 Acquirer *provisionerdserver.Acquirer

 workspaceUsageTracker *workspaceusage.Tracker
 }

 // Close waits for all WebSocket connections to drain before returning.
 _ = (*coordinator).Close()
 }
 _ = api.agentProvider.Close()
 api.workspaceUsageTracker.Close()
 return nil
 }

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
 "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"
 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.
 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) {
 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
 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
 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
 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
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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
Expand All		@@ -156,6 +158,7 @@ func (r RootCmd) portForward() serpent.Command {
		}

		cancel()
		stopUpdating()
		closeAllListeners()
		}()

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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"
Expand DownExpand Up		@@ -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)
Expand DownExpand Up		@@ -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) {
Expand DownExpand Up		@@ -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)
		})
		}

Expand DownExpand Up		@@ -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)
		})
		}

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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"
Expand DownExpand Up		@@ -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
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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"
Expand DownExpand Up		@@ -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
Expand DownExpand Up		@@ -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()

Expand DownExpand Up		@@ -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)
Expand DownExpand Up		@@ -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)
Expand DownExpand Up		@@ -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.
Expand All		@@ -1200,6 +1214,7 @@ func (api *API) Close() error {
		_ = (*coordinator).Close()
		}
		_ = api.agentProvider.Close()
		api.workspaceUsageTracker.Close()
		return nil
		}

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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"
Expand DownExpand Up		@@ -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.
Expand DownExpand Up		@@ -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) {
Expand DownExpand Up		@@ -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,
		}
		}

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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++
		}
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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 (
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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
Expand Down