@@ -3,6 +3,7 @@ package prebuilds_test
33import (
44"context"
55"database/sql"
6+ "errors"
67"strings"
78"sync/atomic"
89"testing"
@@ -66,6 +67,32 @@ func (m *storeSpy) ClaimPrebuiltWorkspace(ctx context.Context, arg database.Clai
6667return result ,err
6768}
6869
70+ type errorStore struct {
71+ claimingErr error
72+
73+ database.Store
74+ }
75+
76+ func newErrorStore (db database.Store ,claimingErr error )* errorStore {
77+ return & errorStore {
78+ Store :db ,
79+ claimingErr :claimingErr ,
80+ }
81+ }
82+
83+ func (es * errorStore )InTx (fn func (store database.Store )error ,opts * database.TxOptions )error {
84+ // Pass failure store down into transaction store.
85+ return es .Store .InTx (func (store database.Store )error {
86+ newES := newErrorStore (store ,es .claimingErr )
87+
88+ return fn (newES )
89+ },opts )
90+ }
91+
92+ func (es * errorStore )ClaimPrebuiltWorkspace (ctx context.Context ,arg database.ClaimPrebuiltWorkspaceParams ) (database.ClaimPrebuiltWorkspaceRow ,error ) {
93+ return database.ClaimPrebuiltWorkspaceRow {},es .claimingErr
94+ }
95+
6996func TestClaimPrebuild (t * testing.T ) {
7097t .Parallel ()
7198
@@ -284,6 +311,176 @@ func TestClaimPrebuild(t *testing.T) {
284311}
285312}
286313
314+ func TestClaimPrebuild_CheckDifferentErrors (t * testing.T ) {
315+ t .Parallel ()
316+
317+ if ! dbtestutil .WillUsePostgres () {
318+ t .Skip ("This test requires postgres" )
319+ }
320+
321+ const (
322+ desiredInstances = 1
323+ presetCount = 2
324+
325+ expectedPrebuildsCount = desiredInstances * presetCount
326+ )
327+
328+ cases := map [string ]struct {
329+ claimingErr error
330+ checkFn func (
331+ t * testing.T ,
332+ ctx context.Context ,
333+ store database.Store ,
334+ userClient * codersdk.Client ,
335+ user codersdk.User ,
336+ templateVersionID uuid.UUID ,
337+ presetID uuid.UUID ,
338+ )
339+ }{
340+ "ErrNoClaimablePrebuiltWorkspaces is returned" : {
341+ claimingErr :agplprebuilds .ErrNoClaimablePrebuiltWorkspaces ,
342+ checkFn :func (
343+ t * testing.T ,
344+ ctx context.Context ,
345+ store database.Store ,
346+ userClient * codersdk.Client ,
347+ user codersdk.User ,
348+ templateVersionID uuid.UUID ,
349+ presetID uuid.UUID ,
350+ ) {
351+ // When: a user creates a new workspace with a preset for which prebuilds are configured.
352+ workspaceName := strings .ReplaceAll (testutil .GetRandomName (t ),"_" ,"-" )
353+ userWorkspace ,err := userClient .CreateUserWorkspace (ctx ,user .Username , codersdk.CreateWorkspaceRequest {
354+ TemplateVersionID :templateVersionID ,
355+ Name :workspaceName ,
356+ TemplateVersionPresetID :presetID ,
357+ })
358+
359+ require .NoError (t ,err )
360+ coderdtest .AwaitWorkspaceBuildJobCompleted (t ,userClient ,userWorkspace .LatestBuild .ID )
361+
362+ // Then: the number of running prebuilds hasn't changed because claiming prebuild is failed and we fallback to creating new workspace.
363+ currentPrebuilds ,err := store .GetRunningPrebuiltWorkspaces (ctx )
364+ require .NoError (t ,err )
365+ require .Equal (t ,expectedPrebuildsCount ,len (currentPrebuilds ))
366+ },
367+ },
368+ "unexpected error during claim is returned" : {
369+ claimingErr :errors .New ("unexpected error during claim" ),
370+ checkFn :func (
371+ t * testing.T ,
372+ ctx context.Context ,
373+ store database.Store ,
374+ userClient * codersdk.Client ,
375+ user codersdk.User ,
376+ templateVersionID uuid.UUID ,
377+ presetID uuid.UUID ,
378+ ) {
379+ // When: a user creates a new workspace with a preset for which prebuilds are configured.
380+ workspaceName := strings .ReplaceAll (testutil .GetRandomName (t ),"_" ,"-" )
381+ _ ,err := userClient .CreateUserWorkspace (ctx ,user .Username , codersdk.CreateWorkspaceRequest {
382+ TemplateVersionID :templateVersionID ,
383+ Name :workspaceName ,
384+ TemplateVersionPresetID :presetID ,
385+ })
386+
387+ // Then: unexpected error happened and was propagated all the way to the caller
388+ require .Error (t ,err )
389+ require .ErrorContains (t ,err ,"unexpected error during claim" )
390+
391+ // Then: the number of running prebuilds hasn't changed because claiming prebuild is failed.
392+ currentPrebuilds ,err := store .GetRunningPrebuiltWorkspaces (ctx )
393+ require .NoError (t ,err )
394+ require .Equal (t ,expectedPrebuildsCount ,len (currentPrebuilds ))
395+ },
396+ },
397+ }
398+
399+ for name ,tc := range cases {
400+ t .Run (name ,func (t * testing.T ) {
401+ t .Parallel ()
402+
403+ // Setup.
404+ ctx := testutil .Context (t ,testutil .WaitMedium )
405+ db ,pubsub := dbtestutil .NewDB (t )
406+ errorStore := newErrorStore (db ,tc .claimingErr )
407+
408+ logger := testutil .Logger (t )
409+ client ,_ ,api ,owner := coderdenttest .NewWithAPI (t ,& coderdenttest.Options {
410+ Options :& coderdtest.Options {
411+ IncludeProvisionerDaemon :true ,
412+ Database :errorStore ,
413+ Pubsub :pubsub ,
414+ },
415+
416+ EntitlementsUpdateInterval :time .Second ,
417+ })
418+
419+ reconciler := prebuilds .NewStoreReconciler (errorStore ,pubsub , codersdk.PrebuildsConfig {},logger ,quartz .NewMock (t ))
420+ var claimer agplprebuilds.Claimer = & prebuilds.EnterpriseClaimer {}
421+ api .AGPL .PrebuildsClaimer .Store (& claimer )
422+
423+ version := coderdtest .CreateTemplateVersion (t ,client ,owner .OrganizationID ,templateWithAgentAndPresetsWithPrebuilds (desiredInstances ))
424+ _ = coderdtest .AwaitTemplateVersionJobCompleted (t ,client ,version .ID )
425+ coderdtest .CreateTemplate (t ,client ,owner .OrganizationID ,version .ID )
426+ presets ,err := client .TemplateVersionPresets (ctx ,version .ID )
427+ require .NoError (t ,err )
428+ require .Len (t ,presets ,presetCount )
429+
430+ userClient ,user := coderdtest .CreateAnotherUser (t ,client ,owner .OrganizationID ,rbac .RoleMember ())
431+
432+ ctx = dbauthz .AsPrebuildsOrchestrator (ctx )
433+
434+ // Given: the reconciliation state is snapshot.
435+ state ,err := reconciler .SnapshotState (ctx ,errorStore )
436+ require .NoError (t ,err )
437+ require .Len (t ,state .Presets ,presetCount )
438+
439+ // When: a reconciliation is setup for each preset.
440+ for _ ,preset := range presets {
441+ ps ,err := state .FilterByPreset (preset .ID )
442+ require .NoError (t ,err )
443+ require .NotNil (t ,ps )
444+ actions ,err := reconciler .CalculateActions (ctx ,* ps )
445+ require .NoError (t ,err )
446+ require .NotNil (t ,actions )
447+
448+ require .NoError (t ,reconciler .ReconcilePreset (ctx ,* ps ))
449+ }
450+
451+ // Given: a set of running, eligible prebuilds eventually starts up.
452+ runningPrebuilds := make (map [uuid.UUID ]database.GetRunningPrebuiltWorkspacesRow ,desiredInstances * presetCount )
453+ require .Eventually (t ,func ()bool {
454+ rows ,err := errorStore .GetRunningPrebuiltWorkspaces (ctx )
455+ require .NoError (t ,err )
456+
457+ for _ ,row := range rows {
458+ runningPrebuilds [row .CurrentPresetID .UUID ]= row
459+
460+ agents ,err := db .GetWorkspaceAgentsInLatestBuildByWorkspaceID (ctx ,row .ID )
461+ require .NoError (t ,err )
462+
463+ // Workspaces are eligible once its agent is marked "ready".
464+ for _ ,agent := range agents {
465+ require .NoError (t ,db .UpdateWorkspaceAgentLifecycleStateByID (ctx , database.UpdateWorkspaceAgentLifecycleStateByIDParams {
466+ ID :agent .ID ,
467+ LifecycleState :database .WorkspaceAgentLifecycleStateReady ,
468+ StartedAt : sql.NullTime {Time :time .Now ().Add (time .Hour ),Valid :true },
469+ ReadyAt : sql.NullTime {Time :time .Now ().Add (- 1 * time .Hour ),Valid :true },
470+ }))
471+ }
472+ }
473+
474+ t .Logf ("found %d running prebuilds so far, want %d" ,len (runningPrebuilds ),expectedPrebuildsCount )
475+
476+ return len (runningPrebuilds )== expectedPrebuildsCount
477+ },testutil .WaitSuperLong ,testutil .IntervalSlow )
478+
479+ tc .checkFn (t ,ctx ,errorStore ,userClient ,user ,version .ID ,presets [0 ].ID )
480+ })
481+ }
482+ }
483+
287484func templateWithAgentAndPresetsWithPrebuilds (desiredInstances int32 )* echo.Responses {
288485return & echo.Responses {
289486Parse :echo .ParseComplete ,