mirror of
https://github.com/coder/coder.git
synced 2025-07-06 15:41:45 +00:00
fix(cli): port-forward: update workspace last_used_at (#12659)
This PR updates the coder port-forward command to periodically inform coderd that the workspace is being used: - Adds workspaceusage.Tracker which periodically batch-updates workspace LastUsedAt - Adds coderd endpoint to signal workspace usage - Updates coder port-forward to periodically hit this endpoint - Modifies BatchUpdateWorkspacesLastUsedAt to avoid overwriting with stale data Co-authored-by: Danny Kopping <danny@coder.com>
This commit is contained in:
Cian Johnston
GitHub
225
coderd/workspaceusage/tracker_test.go
Normal file
225
coderd/workspaceusage/tracker_test.go
Normal file
@ -0,0 +1,225 @@
|
||||
package workspaceusage_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"sort"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/goleak"
|
||||
"go.uber.org/mock/gomock"
|
||||
|
||||
"cdr.dev/slog"
|
||||
"cdr.dev/slog/sloggers/slogtest"
|
||||
"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/dbmock"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtime"
|
||||
"github.com/coder/coder/v2/coderd/database/pubsub"
|
||||
"github.com/coder/coder/v2/coderd/workspaceusage"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
func TestTracker(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mDB := dbmock.NewMockStore(ctrl)
|
||||
log := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
|
||||
|
||||
tickCh := make(chan time.Time)
|
||||
flushCh := make(chan int, 1)
|
||||
wut := workspaceusage.New(mDB,
|
||||
workspaceusage.WithLogger(log),
|
||||
workspaceusage.WithTickFlush(tickCh, flushCh),
|
||||
)
|
||||
defer wut.Close()
|
||||
|
||||
// 1. No marked workspaces should imply no flush.
|
||||
now := dbtime.Now()
|
||||
tickCh <- now
|
||||
count := <-flushCh
|
||||
require.Equal(t, 0, count, "expected zero flushes")
|
||||
|
||||
// 2. One marked workspace should cause a flush.
|
||||
ids := []uuid.UUID{uuid.New()}
|
||||
now = dbtime.Now()
|
||||
wut.Add(ids[0])
|
||||
mDB.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
|
||||
LastUsedAt: now,
|
||||
IDs: ids,
|
||||
}).Times(1)
|
||||
tickCh <- now
|
||||
count = <-flushCh
|
||||
require.Equal(t, 1, count, "expected one flush with one id")
|
||||
|
||||
// 3. Lots of marked workspaces should also cause a flush.
|
||||
for i := 0; i < 31; i++ {
|
||||
ids = append(ids, uuid.New())
|
||||
}
|
||||
|
||||
// Sort ids so mDB know what to expect.
|
||||
sort.Slice(ids, func(i, j int) bool {
|
||||
return bytes.Compare(ids[i][:], ids[j][:]) < 0
|
||||
})
|
||||
|
||||
now = dbtime.Now()
|
||||
mDB.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
|
||||
LastUsedAt: now,
|
||||
IDs: ids,
|
||||
})
|
||||
for _, id := range ids {
|
||||
wut.Add(id)
|
||||
}
|
||||
tickCh <- now
|
||||
count = <-flushCh
|
||||
require.Equal(t, len(ids), count, "incorrect number of ids flushed")
|
||||
|
||||
// 4. Try to cause a race condition!
|
||||
now = dbtime.Now()
|
||||
// Difficult to know what to EXPECT here, so we won't check strictly here.
|
||||
mDB.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), gomock.Any()).MinTimes(1).MaxTimes(len(ids))
|
||||
// Try to force a race condition.
|
||||
var wg sync.WaitGroup
|
||||
count = 0
|
||||
for i := 0; i < len(ids); i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
tickCh <- now
|
||||
}()
|
||||
wut.Add(ids[i])
|
||||
}
|
||||
|
||||
for i := 0; i < len(ids); i++ {
|
||||
count += <-flushCh
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
require.Equal(t, len(ids), count, "incorrect number of ids flushed")
|
||||
|
||||
// 5. Closing multiple times should not be a problem.
|
||||
wut.Close()
|
||||
wut.Close()
|
||||
}
|
||||
|
||||
// This test performs a more 'integration-style' test with multiple instances.
|
||||
func TestTracker_MultipleInstances(t *testing.T) {
|
||||
t.Parallel()
|
||||
if !dbtestutil.WillUsePostgres() {
|
||||
t.Skip("this test only makes sense with postgres")
|
||||
}
|
||||
|
||||
// Given we have two coderd instances connected to the same database
|
||||
var (
|
||||
ctx = testutil.Context(t, testutil.WaitLong)
|
||||
db, _ = dbtestutil.NewDB(t)
|
||||
// real pubsub is not safe for concurrent use, and this test currently
|
||||
// does not depend on pubsub
|
||||
ps = pubsub.NewInMemory()
|
||||
wuTickA = make(chan time.Time)
|
||||
wuFlushA = make(chan int, 1)
|
||||
wuTickB = make(chan time.Time)
|
||||
wuFlushB = make(chan int, 1)
|
||||
clientA = coderdtest.New(t, &coderdtest.Options{
|
||||
WorkspaceUsageTrackerTick: wuTickA,
|
||||
WorkspaceUsageTrackerFlush: wuFlushA,
|
||||
Database: db,
|
||||
Pubsub: ps,
|
||||
})
|
||||
clientB = coderdtest.New(t, &coderdtest.Options{
|
||||
WorkspaceUsageTrackerTick: wuTickB,
|
||||
WorkspaceUsageTrackerFlush: wuFlushB,
|
||||
Database: db,
|
||||
Pubsub: ps,
|
||||
})
|
||||
owner = coderdtest.CreateFirstUser(t, clientA)
|
||||
now = dbtime.Now()
|
||||
)
|
||||
|
||||
clientB.SetSessionToken(clientA.SessionToken())
|
||||
|
||||
// Create a number of workspaces
|
||||
numWorkspaces := 10
|
||||
w := make([]dbfake.WorkspaceResponse, numWorkspaces)
|
||||
for i := 0; i < numWorkspaces; i++ {
|
||||
wr := dbfake.WorkspaceBuild(t, db, database.Workspace{
|
||||
OwnerID: owner.UserID,
|
||||
OrganizationID: owner.OrganizationID,
|
||||
LastUsedAt: now,
|
||||
}).WithAgent().Do()
|
||||
w[i] = wr
|
||||
}
|
||||
|
||||
// Use client A to update LastUsedAt of the first three
|
||||
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[0].Workspace.ID))
|
||||
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[1].Workspace.ID))
|
||||
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[2].Workspace.ID))
|
||||
// Use client B to update LastUsedAt of the next three
|
||||
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[3].Workspace.ID))
|
||||
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[4].Workspace.ID))
|
||||
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[5].Workspace.ID))
|
||||
// The next two will have updated from both instances
|
||||
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[6].Workspace.ID))
|
||||
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[6].Workspace.ID))
|
||||
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[7].Workspace.ID))
|
||||
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[7].Workspace.ID))
|
||||
// The last two will not report any usage.
|
||||
|
||||
// Tick both with different times and wait for both flushes to complete
|
||||
nowA := now.Add(time.Minute)
|
||||
nowB := now.Add(2 * time.Minute)
|
||||
var wg sync.WaitGroup
|
||||
var flushedA, flushedB int
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
wuTickA <- nowA
|
||||
flushedA = <-wuFlushA
|
||||
}()
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
wuTickB <- nowB
|
||||
flushedB = <-wuFlushB
|
||||
}()
|
||||
wg.Wait()
|
||||
|
||||
// We expect 5 flushed IDs each
|
||||
require.Equal(t, 5, flushedA)
|
||||
require.Equal(t, 5, flushedB)
|
||||
|
||||
// Fetch updated workspaces
|
||||
updated := make([]codersdk.Workspace, numWorkspaces)
|
||||
for i := 0; i < numWorkspaces; i++ {
|
||||
ws, err := clientA.Workspace(ctx, w[i].Workspace.ID)
|
||||
require.NoError(t, err)
|
||||
updated[i] = ws
|
||||
}
|
||||
// We expect the first three to have the timestamp of flushA
|
||||
require.Equal(t, nowA.UTC(), updated[0].LastUsedAt.UTC())
|
||||
require.Equal(t, nowA.UTC(), updated[1].LastUsedAt.UTC())
|
||||
require.Equal(t, nowA.UTC(), updated[2].LastUsedAt.UTC())
|
||||
// We expect the next three to have the timestamp of flushB
|
||||
require.Equal(t, nowB.UTC(), updated[3].LastUsedAt.UTC())
|
||||
require.Equal(t, nowB.UTC(), updated[4].LastUsedAt.UTC())
|
||||
require.Equal(t, nowB.UTC(), updated[5].LastUsedAt.UTC())
|
||||
// The next two should have the timestamp of flushB as it is newer than flushA
|
||||
require.Equal(t, nowB.UTC(), updated[6].LastUsedAt.UTC())
|
||||
require.Equal(t, nowB.UTC(), updated[7].LastUsedAt.UTC())
|
||||
// And the last two should be untouched
|
||||
require.Equal(t, w[8].Workspace.LastUsedAt.UTC(), updated[8].LastUsedAt.UTC())
|
||||
require.Equal(t, w[8].Workspace.LastUsedAt.UTC(), updated[8].LastUsedAt.UTC())
|
||||
require.Equal(t, w[9].Workspace.LastUsedAt.UTC(), updated[9].LastUsedAt.UTC())
|
||||
require.Equal(t, w[9].Workspace.LastUsedAt.UTC(), updated[9].LastUsedAt.UTC())
|
||||
}
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
goleak.VerifyTestMain(m)
|
||||
}
|
||||
Reference in New Issue
Block a user