Skip to content

Commit 3be6220

Browse files
johnstcnjohnstcn
committed
add workspace usage tracking to cli/portforward, fix tests
1 parent ca51af6 commit 3be6220

File tree

3 files changed

+71
-1
lines changed

3 files changed

+71
-1
lines changed

‎cli/portforward.go

+3
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,8 @@ func (r *RootCmd) portForward() *serpent.Command {
136136
listeners[i] = l
137137
}
138138

139+
stopUpdating := client.UpdateWorkspaceUsageContext(ctx, workspace.ID)
140+
139141
// Wait for the context to be canceled or for a signal and close
140142
// all listeners.
141143
var closeErr error
@@ -156,6 +158,7 @@ func (r *RootCmd) portForward() *serpent.Command {
156158
}
157159

158160
cancel()
161+
stopUpdating()
159162
closeAllListeners()
160163
}()
161164

‎cli/portforward_test.go

+18-1
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ import (
2121
"github.com/coder/coder/v2/coderd/coderdtest"
2222
"github.com/coder/coder/v2/coderd/database"
2323
"github.com/coder/coder/v2/coderd/database/dbfake"
24+
"github.com/coder/coder/v2/coderd/database/dbtestutil"
25+
"github.com/coder/coder/v2/coderd/database/dbtime"
26+
"github.com/coder/coder/v2/coderd/workspaceusage"
2427
"github.com/coder/coder/v2/codersdk"
2528
"github.com/coder/coder/v2/pty/ptytest"
2629
"github.com/coder/coder/v2/testutil"
@@ -96,7 +99,15 @@ func TestPortForward(t *testing.T) {
9699
// Setup agent once to be shared between test-cases (avoid expensive
97100
// non-parallel setup).
98101
var (
99-
client, db = coderdtest.NewWithDatabase(t, nil)
102+
db, ps = dbtestutil.NewDB(t)
103+
wuTick = make(chan time.Time)
104+
wuFlush = make(chan int, 1)
105+
wut = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlush), workspaceusage.WithTickChannel(wuTick))
106+
client = coderdtest.New(t, &coderdtest.Options{
107+
WorkspaceUsageTracker: wut,
108+
Database: db,
109+
Pubsub: ps,
110+
})
100111
admin = coderdtest.CreateFirstUser(t, client)
101112
member, memberUser = coderdtest.CreateAnotherUser(t, client, admin.OrganizationID)
102113
workspace = runAgent(t, client, memberUser.ID, db)
@@ -149,6 +160,8 @@ func TestPortForward(t *testing.T) {
149160
err = <-errC
150161
require.ErrorIs(t, err, context.Canceled)
151162

163+
wuTick <- dbtime.Now()
164+
<-wuFlush
152165
updated, err := client.Workspace(context.Background(), workspace.ID)
153166
require.NoError(t, err)
154167
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
@@ -201,6 +214,8 @@ func TestPortForward(t *testing.T) {
201214
err = <-errC
202215
require.ErrorIs(t, err, context.Canceled)
203216

217+
wuTick <- dbtime.Now()
218+
<-wuFlush
204219
updated, err := client.Workspace(context.Background(), workspace.ID)
205220
require.NoError(t, err)
206221
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
@@ -266,6 +281,8 @@ func TestPortForward(t *testing.T) {
266281
err := <-errC
267282
require.ErrorIs(t, err, context.Canceled)
268283

284+
wuTick <- dbtime.Now()
285+
<-wuFlush
269286
updated, err := client.Workspace(context.Background(), workspace.ID)
270287
require.NoError(t, err)
271288
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)

‎codersdk/workspaces.go

+50
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import (
1111
"github.com/google/uuid"
1212
"golang.org/x/xerrors"
1313

14+
"cdr.dev/slog"
15+
1416
"github.com/coder/coder/v2/coderd/tracing"
1517
)
1618

@@ -314,6 +316,54 @@ func (c *Client) PutExtendWorkspace(ctx context.Context, id uuid.UUID, req PutEx
314316
return nil
315317
}
316318

319+
// PostWorkspaceUsage marks the workspace as having been used recently.
320+
func (c *Client) PostWorkspaceUsage(ctx context.Context, id uuid.UUID) error {
321+
path := fmt.Sprintf("/api/v2/workspaces/%s/usage", id.String())
322+
res, err := c.Request(ctx, http.MethodPost, path, nil)
323+
if err != nil {
324+
return xerrors.Errorf("post workspace usage: %w", err)
325+
}
326+
defer res.Body.Close()
327+
if res.StatusCode != http.StatusNoContent {
328+
return ReadBodyAsError(res)
329+
}
330+
return nil
331+
}
332+
333+
// UpdateWorkspaceUsageContext periodically posts workspace usage for the workspace
334+
// with the given id in the background.
335+
// The caller is responsible for calling the returned function to stop the background
336+
// process.
337+
func (c *Client) UpdateWorkspaceUsageContext(ctx context.Context, id uuid.UUID) func() {
338+
hbCtx, hbCancel := context.WithCancel(ctx)
339+
// Perform one initial heartbeat
340+
if err := c.PostWorkspaceUsage(hbCtx, id); err != nil {
341+
c.logger.Warn(ctx, "failed to post workspace usage", slog.Error(err))
342+
}
343+
ticker := time.NewTicker(time.Minute)
344+
doneCh := make(chan struct{})
345+
go func() {
346+
defer func() {
347+
ticker.Stop()
348+
close(doneCh)
349+
}()
350+
for {
351+
select {
352+
case <-ticker.C:
353+
if err := c.PostWorkspaceUsage(hbCtx, id); err != nil {
354+
c.logger.Warn(ctx, "failed to post workspace usage in background", slog.Error(err))
355+
}
356+
case <-hbCtx.Done():
357+
return
358+
}
359+
}
360+
}()
361+
return func() {
362+
hbCancel()
363+
<-doneCh
364+
}
365+
}
366+
317367
// UpdateWorkspaceDormancy is a request to activate or make a workspace dormant.
318368
// A value of false will activate a dormant workspace.
319369
type UpdateWorkspaceDormancy struct {

0 commit comments

Comments
 (0)