synced/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2025-07-09 11:45:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

More claim tests

Browse Source 
Signed-off-by: Danny Kopping <dannykopping@gmail.com>
This commit is contained in:
Danny Kopping
2025-03-04 20:22:51 +00:00
parent 431ceceee1
commit e261aee191
2 changed files with 203 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
enterprise/coderd/coderd.go
View File

@ -638,7 +638,7 @@ type API struct {
licenseMetricsCollector *license.MetricsCollector
tailnetService *tailnet.ClientService
prebuildsReconciler agplprebuilds.Reconciler
PrebuildsReconciler agplprebuilds.Reconciler
prebuildsMetricsCollector *prebuilds.MetricsCollector
}
@ -671,10 +671,10 @@ func (api *API) Close() error {
api.Options.CheckInactiveUsersCancelFunc()
}
if api.prebuildsReconciler != nil {
if api.PrebuildsReconciler != nil {
ctx, giveUp := context.WithTimeoutCause(context.Background(), time.Second*30, errors.New("gave up waiting for reconciler to stop"))
defer giveUp()
api.prebuildsReconciler.Stop(ctx, xerrors.New("api closed")) // TODO: determine root cause (requires changes up the stack, though).
api.PrebuildsReconciler.Stop(ctx, xerrors.New("api closed")) // TODO: determine root cause (requires changes up the stack, though).
}
return api.AGPL.Close()
@ -879,12 +879,12 @@ func (api *API) updateEntitlements(ctx context.Context) error {
api.AGPL.PortSharer.Store(&ps)
}
if initial, changed, enabled := featureChanged(codersdk.FeatureWorkspacePrebuilds); shouldUpdate(initial, changed, enabled) {
if initial, changed, enabled := featureChanged(codersdk.FeatureWorkspacePrebuilds); shouldUpdate(initial, changed, enabled) || api.PrebuildsReconciler == nil {
reconciler, claimer, metrics := api.setupPrebuilds(enabled)
if api.prebuildsReconciler != nil {
if api.PrebuildsReconciler != nil {
stopCtx, giveUp := context.WithTimeoutCause(context.Background(), time.Second*30, errors.New("gave up waiting for reconciler to stop"))
defer giveUp()
api.prebuildsReconciler.Stop(stopCtx, errors.New("entitlements change"))
api.PrebuildsReconciler.Stop(stopCtx, errors.New("entitlements change"))
}
// Only register metrics once.
@ -892,7 +892,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
api.prebuildsMetricsCollector = metrics
}
api.prebuildsReconciler = reconciler
api.PrebuildsReconciler = reconciler
go reconciler.RunLoop(context.Background())
api.AGPL.PrebuildsClaimer.Store(&claimer)

152
enterprise/coderd/prebuilds/claim_test.go
View File

@ -8,6 +8,7 @@ import (
"testing"
"time"
"github.com/coder/serpent"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
@ -19,7 +20,6 @@ import (
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
"github.com/coder/coder/v2/enterprise/coderd/license"
"github.com/coder/coder/v2/enterprise/coderd/prebuilds"
"github.com/coder/coder/v2/provisioner/echo"
"github.com/coder/coder/v2/provisionersdk/proto"
"github.com/coder/coder/v2/testutil"
@ -58,40 +58,102 @@ func (m *storeSpy) ClaimPrebuild(ctx context.Context, arg database.ClaimPrebuild
m.claims.Add(1)
m.claimParams.Store(&arg)
result, err := m.Store.ClaimPrebuild(ctx, arg)
if err == nil {
m.claimedWorkspace.Store(&result)
}
return result, err
}
func TestClaimPrebuild(t *testing.T) {
t.Parallel()
// Setup. // TODO: abstract?
ctx := testutil.Context(t, testutil.WaitSuperLong)
db, pubsub := dbtestutil.NewDB(t)
spy := newStoreSpy(db)
client, _, _, owner := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
Options: &coderdtest.Options{
IncludeProvisionerDaemon: true,
Database: spy,
Pubsub: pubsub,
},
LicenseOptions: &coderdenttest.LicenseOptions{
Features: license.Features{
codersdk.FeatureWorkspacePrebuilds: 1,
},
},
})
controller := prebuilds.NewStoreReconciler(spy, pubsub, codersdk.PrebuildsConfig{}, testutil.Logger(t))
const (
desiredInstances = 1
presetCount = 2
)
cases := map[string]struct {
entitlementEnabled bool
experimentEnabled bool
attemptPrebuildClaim bool
expectPrebuildClaimed bool
markPrebuildsClaimable bool
expectedPrebuildsCount int
}{
"without the experiment enabled, prebuilds will not provisioned": {
experimentEnabled: false,
entitlementEnabled: true,
attemptPrebuildClaim: false,
expectedPrebuildsCount: 0,
},
"without the entitlement, prebuilds will not provisioned": {
experimentEnabled: true,
entitlementEnabled: false,
attemptPrebuildClaim: false,
expectedPrebuildsCount: 0,
},
"with everything enabled, but no eligible prebuilds to claim": {
entitlementEnabled: true,
experimentEnabled: true,
attemptPrebuildClaim: true,
expectPrebuildClaimed: false,
markPrebuildsClaimable: false,
expectedPrebuildsCount: desiredInstances * presetCount,
},
"with everything enabled, claiming an eligible prebuild should succeed": {
entitlementEnabled: true,
experimentEnabled: true,
attemptPrebuildClaim: true,
expectPrebuildClaimed: true,
markPrebuildsClaimable: true,
expectedPrebuildsCount: desiredInstances * presetCount,
},
}
for name, tc := range cases {
tc := tc
t.Run(name, func(t *testing.T) {
t.Parallel()
// Setup. // TODO: abstract?
ctx := testutil.Context(t, testutil.WaitMedium)
db, pubsub := dbtestutil.NewDB(t)
spy := newStoreSpy(db)
var prebuildsEntitled int64
if tc.entitlementEnabled {
prebuildsEntitled = 1
}
client, _, api, owner := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
Options: &coderdtest.Options{
IncludeProvisionerDaemon: true,
Database: spy,
Pubsub: pubsub,
DeploymentValues: coderdtest.DeploymentValues(t, func(values *codersdk.DeploymentValues) {
values.Prebuilds.ReconciliationInterval = serpent.Duration(time.Hour) // We will kick off a reconciliation manually.
if tc.experimentEnabled {
values.Experiments = serpent.StringArray{string(codersdk.ExperimentWorkspacePrebuilds)}
}
}),
},
EntitlementsUpdateInterval: time.Second,
LicenseOptions: &coderdenttest.LicenseOptions{
Features: license.Features{
codersdk.FeatureWorkspacePrebuilds: prebuildsEntitled,
},
},
})
// The entitlements will need to refresh before the reconciler is set.
require.Eventually(t, func() bool {
return api.PrebuildsReconciler != nil
}, testutil.WaitSuperLong, testutil.IntervalFast)
version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, templateWithAgentAndPresetsWithPrebuilds(desiredInstances))
_ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID)
@ -104,7 +166,7 @@ func TestClaimPrebuild(t *testing.T) {
ctx = dbauthz.AsSystemRestricted(ctx)
// Given: a reconciliation completes.
require.NoError(t, controller.ReconcileAll(ctx))
require.NoError(t, api.PrebuildsReconciler.ReconcileAll(ctx))
// Given: a set of running, eligible prebuilds eventually starts up.
runningPrebuilds := make(map[uuid.UUID]database.GetRunningPrebuildsRow, desiredInstances*presetCount)
@ -115,6 +177,10 @@ func TestClaimPrebuild(t *testing.T) {
for _, row := range rows {
runningPrebuilds[row.CurrentPresetID.UUID] = row
if !tc.markPrebuildsClaimable {
continue
}
agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, row.WorkspaceID)
require.NoError(t, err)
@ -128,9 +194,9 @@ func TestClaimPrebuild(t *testing.T) {
}
}
t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), desiredInstances*presetCount)
t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), tc.expectedPrebuildsCount)
return len(runningPrebuilds) == (desiredInstances * presetCount)
return len(runningPrebuilds) == tc.expectedPrebuildsCount
}, testutil.WaitSuperLong, testutil.IntervalSlow)
// When: a user creates a new workspace with a preset for which prebuilds are configured.
@ -149,11 +215,28 @@ func TestClaimPrebuild(t *testing.T) {
require.NoError(t, err)
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, userWorkspace.LatestBuild.ID)
// TODO: this feels... wrong; we should probably be injecting an implementation of prebuilds.Claimer.
// Then: if we're not expecting any prebuild claims to succeed, handle this specifically.
if !tc.attemptPrebuildClaim {
require.EqualValues(t, spy.claims.Load(), 0)
require.Nil(t, spy.claimedWorkspace.Load())
currentPrebuilds, err := spy.GetRunningPrebuilds(ctx)
require.NoError(t, err)
// The number of prebuilds should NOT change.
require.Equal(t, len(currentPrebuilds), len(runningPrebuilds))
return
}
// Then: a prebuild should have been claimed.
require.EqualValues(t, spy.claims.Load(), 1)
require.NotNil(t, spy.claims.Load())
require.EqualValues(t, *spy.claimParams.Load(), params)
if !tc.expectPrebuildClaimed {
require.Nil(t, spy.claimedWorkspace.Load())
return
}
require.NotNil(t, spy.claimedWorkspace.Load())
claimed := *spy.claimedWorkspace.Load()
require.NotEqual(t, claimed, uuid.Nil)
@ -180,6 +263,20 @@ func TestClaimPrebuild(t *testing.T) {
}
}
require.False(t, found, "claimed prebuild should not still be considered a running prebuild")
// Then: reconciling at this point will provision a new prebuild to replace the claimed one.
require.NoError(t, api.PrebuildsReconciler.ReconcileAll(ctx))
require.Eventually(t, func() bool {
rows, err := spy.GetRunningPrebuilds(ctx)
require.NoError(t, err)
t.Logf("found %d running prebuilds so far, want %d", len(rows), tc.expectedPrebuildsCount)
return len(runningPrebuilds) == tc.expectedPrebuildsCount
}, testutil.WaitSuperLong, testutil.IntervalSlow)
})
}
}
func templateWithAgentAndPresetsWithPrebuilds(desiredInstances int32) *echo.Responses {
@ -256,5 +353,4 @@ func templateWithAgentAndPresetsWithPrebuilds(desiredInstances int32) *echo.Resp
}
}
// TODO(dannyk): test claiming a prebuild causes a replacement to be provisioned.
// TODO(dannyk): test that prebuilds are only attempted to be claimed for net-new workspace builds