44"context"
55"database/sql"
66"fmt"
7+ "sort"
78"sync"
89"testing"
910"time"
@@ -1429,6 +1430,245 @@ func TestTrackResourceReplacement(t *testing.T) {
14291430require .EqualValues (t ,1 ,metric .GetCounter ().GetValue ())
14301431}
14311432
1433+ func TestExpiredPrebuildsMultipleActions (t * testing.T ) {
1434+ t .Parallel ()
1435+
1436+ if ! dbtestutil .WillUsePostgres () {
1437+ t .Skip ("This test requires postgres" )
1438+ }
1439+
1440+ // Test cases verify the behavior of prebuild creation depending on configured failure limits.
1441+ testCases := []struct {
1442+ name string
1443+ running int
1444+ desired int32
1445+ expired int
1446+ extraneous int
1447+ created int
1448+ }{
1449+ // With 2 running prebuilds, none of which are expired, and the desired count is met,
1450+ // no deletions or creations should occur.
1451+ {
1452+ name :"no expired prebuilds - no actions taken" ,
1453+ running :2 ,
1454+ desired :2 ,
1455+ expired :0 ,
1456+ extraneous :0 ,
1457+ created :0 ,
1458+ },
1459+ // With 2 running prebuilds, 1 of which is expired, the expired prebuild should be deleted,
1460+ // and one new prebuild should be created to maintain the desired count.
1461+ {
1462+ name :"one expired prebuild – deleted and replaced" ,
1463+ running :2 ,
1464+ desired :2 ,
1465+ expired :1 ,
1466+ extraneous :0 ,
1467+ created :1 ,
1468+ },
1469+ // With 2 running prebuilds, both expired, both should be deleted,
1470+ // and 2 new prebuilds created to match the desired count.
1471+ {
1472+ name :"all prebuilds expired – all deleted and recreated" ,
1473+ running :2 ,
1474+ desired :2 ,
1475+ expired :2 ,
1476+ extraneous :0 ,
1477+ created :2 ,
1478+ },
1479+ // With 4 running prebuilds, 2 of which are expired, and the desired count is 2,
1480+ // the expired prebuilds should be deleted. No new creations are needed
1481+ // since removing the expired ones brings actual = desired.
1482+ {
1483+ name :"expired prebuilds deleted to reach desired count" ,
1484+ running :4 ,
1485+ desired :2 ,
1486+ expired :2 ,
1487+ extraneous :0 ,
1488+ created :0 ,
1489+ },
1490+ // With 4 running prebuilds (1 expired), and the desired count is 2,
1491+ // the first action should delete the expired one,
1492+ // and the second action should delete one additional (non-expired) prebuild
1493+ // to eliminate the remaining excess.
1494+ {
1495+ name :"expired prebuild deleted first, then extraneous" ,
1496+ running :4 ,
1497+ desired :2 ,
1498+ expired :1 ,
1499+ extraneous :1 ,
1500+ created :0 ,
1501+ },
1502+ }
1503+
1504+ for _ ,tc := range testCases {
1505+ t .Run (tc .name ,func (t * testing.T ) {
1506+ t .Parallel ()
1507+
1508+ clock := quartz .NewMock (t )
1509+ ctx := testutil .Context (t ,testutil .WaitLong )
1510+ cfg := codersdk.PrebuildsConfig {}
1511+ logger := slogtest .Make (
1512+ t ,& slogtest.Options {IgnoreErrors :true },
1513+ ).Leveled (slog .LevelDebug )
1514+ db ,pubSub := dbtestutil .NewDB (t )
1515+ fakeEnqueuer := newFakeEnqueuer ()
1516+ registry := prometheus .NewRegistry ()
1517+ controller := prebuilds .NewStoreReconciler (db ,pubSub ,cfg ,logger ,clock ,registry ,fakeEnqueuer )
1518+
1519+ // Set up test environment with a template, version, and preset
1520+ ownerID := uuid .New ()
1521+ dbgen .User (t ,db , database.User {
1522+ ID :ownerID ,
1523+ })
1524+ org ,template := setupTestDBTemplate (t ,db ,ownerID ,false )
1525+ templateVersionID := setupTestDBTemplateVersion (ctx ,t ,clock ,db ,pubSub ,org .ID ,ownerID ,template .ID )
1526+
1527+ ttlDuration := muchEarlier - time .Hour
1528+ ttl := int32 (- ttlDuration .Seconds ())
1529+ preset := setupTestDBPreset (t ,db ,templateVersionID ,tc .desired ,"b0rked" ,withTTL (ttl ))
1530+
1531+ // The implementation uses time.Since(prebuild.CreatedAt) > ttl to check a prebuild expiration.
1532+ // Since our mock clock defaults to a fixed time, we must align it with the current time
1533+ // to ensure time-based logic works correctly in tests.
1534+ clock .Set (time .Now ())
1535+
1536+ runningWorkspaces := make (map [string ]database.WorkspaceTable )
1537+ nonExpiredWorkspaces := make ([]database.WorkspaceTable ,0 ,tc .running - tc .expired )
1538+ expiredWorkspaces := make ([]database.WorkspaceTable ,0 ,tc .expired )
1539+ expiredCount := 0
1540+ for r := range tc .running {
1541+ // Space out createdAt timestamps by 1 second to ensure deterministic ordering.
1542+ // This lets the test verify that the correct (oldest) extraneous prebuilds are deleted.
1543+ createdAt := muchEarlier + time .Duration (r )* time .Second
1544+ isExpired := false
1545+ if tc .expired > expiredCount {
1546+ // Set createdAt far enough in the past so that time.Since(createdAt) > TTL,
1547+ // ensuring the prebuild is treated as expired in the test.
1548+ createdAt = ttlDuration - 1 * time .Minute
1549+ isExpired = true
1550+ expiredCount ++
1551+ }
1552+
1553+ workspace ,_ := setupTestDBPrebuild (
1554+ t ,
1555+ clock ,
1556+ db ,
1557+ pubSub ,
1558+ database .WorkspaceTransitionStart ,
1559+ database .ProvisionerJobStatusSucceeded ,
1560+ org .ID ,
1561+ preset ,
1562+ template .ID ,
1563+ templateVersionID ,
1564+ withCreatedAt (clock .Now ().Add (createdAt )),
1565+ )
1566+ if isExpired {
1567+ expiredWorkspaces = append (expiredWorkspaces ,workspace )
1568+ }else {
1569+ nonExpiredWorkspaces = append (nonExpiredWorkspaces ,workspace )
1570+ }
1571+ runningWorkspaces [workspace .ID .String ()]= workspace
1572+ }
1573+
1574+ getJobStatusMap := func (workspaces []database.WorkspaceTable )map [database.ProvisionerJobStatus ]int {
1575+ jobStatusMap := make (map [database.ProvisionerJobStatus ]int )
1576+ for _ ,workspace := range workspaces {
1577+ workspaceBuilds ,err := db .GetWorkspaceBuildsByWorkspaceID (ctx , database.GetWorkspaceBuildsByWorkspaceIDParams {
1578+ WorkspaceID :workspace .ID ,
1579+ })
1580+ require .NoError (t ,err )
1581+
1582+ for _ ,workspaceBuild := range workspaceBuilds {
1583+ job ,err := db .GetProvisionerJobByID (ctx ,workspaceBuild .JobID )
1584+ require .NoError (t ,err )
1585+ jobStatusMap [job .JobStatus ]++
1586+ }
1587+ }
1588+ return jobStatusMap
1589+ }
1590+
1591+ // Assert that the build associated with the given workspace has a 'start' transition status.
1592+ isWorkspaceStarted := func (workspace database.WorkspaceTable ) {
1593+ workspaceBuilds ,err := db .GetWorkspaceBuildsByWorkspaceID (ctx , database.GetWorkspaceBuildsByWorkspaceIDParams {
1594+ WorkspaceID :workspace .ID ,
1595+ })
1596+ require .NoError (t ,err )
1597+ require .Equal (t ,1 ,len (workspaceBuilds ))
1598+ require .Equal (t ,database .WorkspaceTransitionStart ,workspaceBuilds [0 ].Transition )
1599+ }
1600+
1601+ // Assert that the workspace build history includes a 'start' followed by a 'delete' transition status.
1602+ isWorkspaceDeleted := func (workspace database.WorkspaceTable ) {
1603+ workspaceBuilds ,err := db .GetWorkspaceBuildsByWorkspaceID (ctx , database.GetWorkspaceBuildsByWorkspaceIDParams {
1604+ WorkspaceID :workspace .ID ,
1605+ })
1606+ require .NoError (t ,err )
1607+ require .Equal (t ,2 ,len (workspaceBuilds ))
1608+ require .Equal (t ,database .WorkspaceTransitionDelete ,workspaceBuilds [0 ].Transition )
1609+ require .Equal (t ,database .WorkspaceTransitionStart ,workspaceBuilds [1 ].Transition )
1610+ }
1611+
1612+ // Verify that all running workspaces, whether expired or not, have successfully started.
1613+ workspaces ,err := db .GetWorkspacesByTemplateID (ctx ,template .ID )
1614+ require .NoError (t ,err )
1615+ require .Equal (t ,tc .running ,len (workspaces ))
1616+ jobStatusMap := getJobStatusMap (workspaces )
1617+ require .Len (t ,workspaces ,tc .running )
1618+ require .Len (t ,jobStatusMap ,1 )
1619+ require .Equal (t ,tc .running ,jobStatusMap [database .ProvisionerJobStatusSucceeded ])
1620+
1621+ // Assert that all running workspaces (expired and non-expired) have a 'start' transition state.
1622+ for _ ,workspace := range runningWorkspaces {
1623+ isWorkspaceStarted (workspace )
1624+ }
1625+
1626+ // Trigger reconciliation to process expired prebuilds and enforce desired state.
1627+ require .NoError (t ,controller .ReconcileAll (ctx ))
1628+
1629+ // Sort non-expired workspaces by CreatedAt in ascending order (oldest first)
1630+ sort .Slice (nonExpiredWorkspaces ,func (i ,j int )bool {
1631+ return nonExpiredWorkspaces [i ].CreatedAt .Before (nonExpiredWorkspaces [j ].CreatedAt )
1632+ })
1633+
1634+ // Verify the status of each non-expired workspace:
1635+ // - the oldest `tc.extraneous` should have been deleted (i.e., have a 'delete' transition),
1636+ // - while the remaining newer ones should still be running (i.e., have a 'start' transition).
1637+ extraneousCount := 0
1638+ for _ ,running := range nonExpiredWorkspaces {
1639+ if extraneousCount < tc .extraneous {
1640+ isWorkspaceDeleted (running )
1641+ extraneousCount ++
1642+ }else {
1643+ isWorkspaceStarted (running )
1644+ }
1645+ }
1646+ require .Equal (t ,tc .extraneous ,extraneousCount )
1647+
1648+ // Verify that each expired workspace has a 'delete' transition recorded,
1649+ // confirming it was properly marked for cleanup after reconciliation.
1650+ for _ ,expired := range expiredWorkspaces {
1651+ isWorkspaceDeleted (expired )
1652+ }
1653+
1654+ // After handling expired prebuilds, if running < desired, new prebuilds should be created.
1655+ // Verify that the correct number of new prebuild workspaces were created and started.
1656+ allWorkspaces ,err := db .GetWorkspacesByTemplateID (ctx ,template .ID )
1657+ require .NoError (t ,err )
1658+
1659+ createdCount := 0
1660+ for _ ,workspace := range allWorkspaces {
1661+ if _ ,ok := runningWorkspaces [workspace .ID .String ()];! ok {
1662+ // Count and verify only the newly created workspaces (i.e., not part of the original running set)
1663+ isWorkspaceStarted (workspace )
1664+ createdCount ++
1665+ }
1666+ }
1667+ require .Equal (t ,tc .created ,createdCount )
1668+ })
1669+ }
1670+ }
1671+
14321672func newNoopEnqueuer ()* notifications.NoopEnqueuer {
14331673return notifications .NewNoopEnqueuer ()
14341674}
@@ -1538,22 +1778,42 @@ func setupTestDBTemplateVersion(
15381778return templateVersion .ID
15391779}
15401780
1781+ // Preset optional parameters.
1782+ // presetOptions defines a function type for modifying InsertPresetParams.
1783+ type presetOptions func (* database.InsertPresetParams )
1784+
1785+ // withTTL returns a presetOptions function that sets the invalidate_after_secs (TTL) field in InsertPresetParams.
1786+ func withTTL (ttl int32 )presetOptions {
1787+ return func (p * database.InsertPresetParams ) {
1788+ p .InvalidateAfterSecs = sql.NullInt32 {Valid :true ,Int32 :ttl }
1789+ }
1790+ }
1791+
15411792func setupTestDBPreset (
15421793t * testing.T ,
15431794db database.Store ,
15441795templateVersionID uuid.UUID ,
15451796desiredInstances int32 ,
15461797presetName string ,
1798+ opts ... presetOptions ,
15471799) database.TemplateVersionPreset {
15481800t .Helper ()
1549- preset := dbgen . Preset ( t , db , database.InsertPresetParams {
1801+ insertPresetParams := database.InsertPresetParams {
15501802TemplateVersionID :templateVersionID ,
15511803Name :presetName ,
15521804DesiredInstances : sql.NullInt32 {
15531805Valid :true ,
15541806Int32 :desiredInstances ,
15551807},
1556- })
1808+ }
1809+
1810+ // Apply optional parameters to insertPresetParams (e.g., TTL).
1811+ for _ ,opt := range opts {
1812+ opt (& insertPresetParams )
1813+ }
1814+
1815+ preset := dbgen .Preset (t ,db ,insertPresetParams )
1816+
15571817dbgen .PresetParameter (t ,db , database.InsertPresetParametersParams {
15581818TemplateVersionPresetID :preset .ID ,
15591819Names : []string {"test" },
@@ -1562,6 +1822,21 @@ func setupTestDBPreset(
15621822return preset
15631823}
15641824
1825+ // prebuildOptions holds optional parameters for creating a prebuild workspace.
1826+ type prebuildOptions struct {
1827+ createdAt * time.Time
1828+ }
1829+
1830+ // prebuildOption defines a function type to apply optional settings to prebuildOptions.
1831+ type prebuildOption func (* prebuildOptions )
1832+
1833+ // withCreatedAt returns a prebuildOption that sets the CreatedAt timestamp.
1834+ func withCreatedAt (createdAt time.Time )prebuildOption {
1835+ return func (opts * prebuildOptions ) {
1836+ opts .createdAt = & createdAt
1837+ }
1838+ }
1839+
15651840func setupTestDBPrebuild (
15661841t * testing.T ,
15671842clock quartz.Clock ,
@@ -1573,9 +1848,10 @@ func setupTestDBPrebuild(
15731848preset database.TemplateVersionPreset ,
15741849templateID uuid.UUID ,
15751850templateVersionID uuid.UUID ,
1851+ opts ... prebuildOption ,
15761852) (database.WorkspaceTable , database.WorkspaceBuild ) {
15771853t .Helper ()
1578- return setupTestDBWorkspace (t ,clock ,db ,ps ,transition ,prebuildStatus ,orgID ,preset ,templateID ,templateVersionID ,agplprebuilds .SystemUserID ,agplprebuilds .SystemUserID )
1854+ return setupTestDBWorkspace (t ,clock ,db ,ps ,transition ,prebuildStatus ,orgID ,preset ,templateID ,templateVersionID ,agplprebuilds .SystemUserID ,agplprebuilds .SystemUserID , opts ... )
15791855}
15801856
15811857func setupTestDBWorkspace (
@@ -1591,6 +1867,7 @@ func setupTestDBWorkspace(
15911867templateVersionID uuid.UUID ,
15921868initiatorID uuid.UUID ,
15931869ownerID uuid.UUID ,
1870+ opts ... prebuildOption ,
15941871) (database.WorkspaceTable , database.WorkspaceBuild ) {
15951872t .Helper ()
15961873cancelledAt := sql.NullTime {}
@@ -1618,15 +1895,30 @@ func setupTestDBWorkspace(
16181895default :
16191896}
16201897
1898+ // Apply all provided prebuild options.
1899+ prebuiltOptions := & prebuildOptions {}
1900+ for _ ,opt := range opts {
1901+ opt (prebuiltOptions )
1902+ }
1903+
1904+ // Set createdAt to default value if not overridden by options.
1905+ createdAt := clock .Now ().Add (muchEarlier )
1906+ if prebuiltOptions .createdAt != nil {
1907+ createdAt = * prebuiltOptions .createdAt
1908+ // Ensure startedAt matches createdAt for consistency.
1909+ startedAt = sql.NullTime {Time :createdAt ,Valid :true }
1910+ }
1911+
16211912workspace := dbgen .Workspace (t ,db , database.WorkspaceTable {
16221913TemplateID :templateID ,
16231914OrganizationID :orgID ,
16241915OwnerID :ownerID ,
16251916Deleted :false ,
1917+ CreatedAt :createdAt ,
16261918})
16271919job := dbgen .ProvisionerJob (t ,db ,ps , database.ProvisionerJob {
16281920InitiatorID :initiatorID ,
1629- CreatedAt :clock . Now (). Add ( muchEarlier ) ,
1921+ CreatedAt :createdAt ,
16301922StartedAt :startedAt ,
16311923CompletedAt :completedAt ,
16321924CanceledAt :cancelledAt ,