coder · johnstcn · Sep 14, 2023 · Sep 7, 2023 · Sep 8, 2023 · Sep 12, 2023
diff --git a/coderd/activitybump.go b/coderd/activitybump.go
@@ -2,16 +2,13 @@ package coderd
 
 import (
 	"context"
-	"database/sql"
-	"errors"
 	"time"
 
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/coderd/database"
-	"github.com/coder/coder/v2/coderd/database/dbtime"
 )
 
 // activityBumpWorkspace automatically bumps the workspace's auto-off timer
@@ -21,72 +18,7 @@ func activityBumpWorkspace(ctx context.Context, log slog.Logger, db database.Sto
 	// low priority operations fail first.
 	ctx, cancel := context.WithTimeout(ctx, time.Second*15)
 	defer cancel()
-
-	err := db.InTx(func(s database.Store) error {
-		build, err := s.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspaceID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return nil
-		} else if err != nil {
-			return xerrors.Errorf("get latest workspace build: %w", err)
-		}
-
-		job, err := s.GetProvisionerJobByID(ctx, build.JobID)
-		if err != nil {
-			return xerrors.Errorf("get provisioner job: %w", err)
-		}
-
-		if build.Transition != database.WorkspaceTransitionStart || !job.CompletedAt.Valid {
-			return nil
-		}
-
-		if build.Deadline.IsZero() {
-			// Workspace shutdown is manual
-			return nil
-		}
-
-		workspace, err := s.GetWorkspaceByID(ctx, workspaceID)
-		if err != nil {
-			return xerrors.Errorf("get workspace: %w", err)
-		}
-
-		var (
-			// We bump by the original TTL to prevent counter-intuitive behavior
-			// as the TTL wraps. For example, if I set the TTL to 12 hours, sign off
-			// work at midnight, come back at 10am, I would want another full day
-			// of uptime. In the prior implementation, the workspace would enter
-			// a state of always expiring 1 hour in the future
-			bumpAmount = time.Duration(workspace.Ttl.Int64)
-			// DB writes are expensive so we only bump when 5% of the deadline
-			// has elapsed.
-			bumpEvery         = bumpAmount / 20
-			timeSinceLastBump = bumpAmount - time.Until(build.Deadline)
-		)
-
-		if timeSinceLastBump < bumpEvery {
-			return nil
-		}
-
-		if bumpAmount == 0 {
-			return nil
-		}
-
-		newDeadline := dbtime.Now().Add(bumpAmount)
-		if !build.MaxDeadline.IsZero() && newDeadline.After(build.MaxDeadline) {
-			newDeadline = build.MaxDeadline
-		}
-
-		if err := s.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
-			ID:               build.ID,
-			UpdatedAt:        dbtime.Now(),
-			ProvisionerState: build.ProvisionerState,
-			Deadline:         newDeadline,
-			MaxDeadline:      build.MaxDeadline,
-		}); err != nil {
-			return xerrors.Errorf("update workspace build: %w", err)
-		}
-		return nil
-	}, nil)
-	if err != nil {
+	if err := db.ActivityBumpWorkspace(ctx, workspaceID); err != nil {
 		if !xerrors.Is(err, context.Canceled) && !database.IsQueryCanceledError(err) {
 			// Bump will fail if the context is canceled, but this is ok.
 			log.Error(ctx, "bump failed", slog.Error(err),

diff --git a/coderd/activitybump_internal_test.go b/coderd/activitybump_internal_test.go
@@ -0,0 +1,175 @@
+package coderd
+
+import (
+	"database/sql"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbgen"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/testutil"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_ActivityBumpWorkspace(t *testing.T) {
+	t.Parallel()
+
+	for _, tt := range []struct {
+		name           string
+		transition     database.WorkspaceTransition
+		jobCompletedAt sql.NullTime
+		buildDeadline  time.Time
+		maxDeadline    time.Time
+		workspaceTTL   time.Duration
+		expectedBump   time.Duration
+	}{
+		{
+			name:           "NotFinishedYet",
+			transition:     database.WorkspaceTransitionStart,
+			jobCompletedAt: sql.NullTime{},
+			buildDeadline:  dbtime.Now().Add(8 * time.Hour),
+			workspaceTTL:   8 * time.Hour,
+			expectedBump:   0,
+		},
+		{
+			name:           "ManualShutdown",
+			transition:     database.WorkspaceTransitionStart,
+			jobCompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
+			buildDeadline:  time.Time{},
+			expectedBump:   0,
+		},
+		{
+			name:           "NotTimeToBumpYet",
+			transition:     database.WorkspaceTransitionStart,
+			jobCompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
+			buildDeadline:  dbtime.Now().Add(8 * time.Hour),
+			workspaceTTL:   8 * time.Hour,
+			expectedBump:   0,
+		},
+		{
+			name:           "TimeToBump",
+			transition:     database.WorkspaceTransitionStart,
+			jobCompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now().Add(-24 * time.Minute)},
+			buildDeadline:  dbtime.Now().Add(8*time.Hour - 24*time.Minute),
+			workspaceTTL:   8 * time.Hour,
+			expectedBump:   8 * time.Hour,
+		},
+		{
+			name:           "MaxDeadline",
+			transition:     database.WorkspaceTransitionStart,
+			jobCompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now().Add(-24 * time.Minute)},
+			buildDeadline:  dbtime.Now().Add(time.Minute), // last chance to bump!
+			maxDeadline:    dbtime.Now().Add(time.Hour),
+			workspaceTTL:   8 * time.Hour,
+			expectedBump:   1 * time.Hour,
+		},
+		{
+			// A workspace that is still running, has passed its deadline, but has not
+			// yet been auto-stopped should still bump the deadline.
+			name:           "PastDeadlineStillBumps",
+			transition:     database.WorkspaceTransitionStart,
+			jobCompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now().Add(-24 * time.Minute)},
+			buildDeadline:  dbtime.Now().Add(-time.Minute),
+			workspaceTTL:   8 * time.Hour,
+			expectedBump:   8 * time.Hour,
+		},
+		{
+			// A stopped workspace should never bump.
+			name:           "StoppedWorkspace",
+			transition:     database.WorkspaceTransitionStop,
+			jobCompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now().Add(-time.Minute)},
+			buildDeadline:  dbtime.Now().Add(-time.Minute),
+			workspaceTTL:   8 * time.Hour,
+			expectedBump:   0,
+		},
+	} {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			var (
+				ctx   = testutil.Context(t, testutil.WaitShort)
+				log   = slogtest.Make(t, nil)
+				db, _ = dbtestutil.NewDB(t)
+				org   = dbgen.Organization(t, db, database.Organization{})
+				user  = dbgen.User(t, db, database.User{
+					Status: database.UserStatusActive,
+				})
+				_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
+					UserID:         user.ID,
+					OrganizationID: org.ID,
+				})
+				templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
+					OrganizationID: org.ID,
+					CreatedBy:      user.ID,
+				})
+				template = dbgen.Template(t, db, database.Template{
+					OrganizationID:  org.ID,
+					ActiveVersionID: templateVersion.ID,
+					CreatedBy:       user.ID,
+				})
+				ws = dbgen.Workspace(t, db, database.Workspace{
+					OwnerID:        user.ID,
+					OrganizationID: org.ID,
+					TemplateID:     template.ID,
+					Ttl:            sql.NullInt64{Valid: true, Int64: int64(tt.workspaceTTL)},
+				})
+				job = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
+					OrganizationID: org.ID,
+					CompletedAt:    tt.jobCompletedAt,
+				})
+				_ = dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+					JobID: job.ID,
+				})
+				buildID = uuid.New()
+			)
+			// dbgen.WorkspaceBuild automatically sets deadline to now+1 hour if not set
+			err := db.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{
+				ID:                buildID,
+				CreatedAt:         dbtime.Now(),
+				UpdatedAt:         dbtime.Now(),
+				BuildNumber:       1,
+				InitiatorID:       user.ID,
+				Reason:            database.BuildReasonInitiator,
+				WorkspaceID:       ws.ID,
+				JobID:             job.ID,
+				TemplateVersionID: templateVersion.ID,
+				Transition:        tt.transition,
+				Deadline:          tt.buildDeadline,
+				MaxDeadline:       tt.maxDeadline,
+			})
+			require.NoError(t, err, "unexpected error inserting workspace build")
+			bld, err := db.GetWorkspaceBuildByID(ctx, buildID)
+			require.NoError(t, err, "unexpected error fetching inserted workspace build")
+
+			// Validate our initial state before bump
+			require.Equal(t, tt.transition, bld.Transition, "unexpected transition before bump")
+			require.Equal(t, tt.jobCompletedAt.Time.UTC(), job.CompletedAt.Time.UTC(), "unexpected job completed at before bump")
+			require.Equal(t, tt.buildDeadline.UTC(), bld.Deadline.UTC(), "unexpected build deadline before bump")
+			require.Equal(t, tt.maxDeadline.UTC(), bld.MaxDeadline.UTC(), "unexpected max deadline before bump")
+			require.Equal(t, tt.workspaceTTL, time.Duration(ws.Ttl.Int64), "unexpected workspace TTL before bump")
+
+			// new deadline is calculated from the time of the bump
+			approxBumpTime := dbtime.Now()
+			activityBumpWorkspace(ctx, log, db, bld.WorkspaceID)
+
+			// Validate our state after bump
+			updatedBuild, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, bld.WorkspaceID)
+			require.NoError(t, err, "unexpected error getting latest workspace build")
+			if tt.expectedBump == 0 {
+				require.Equal(t, bld.UpdatedAt.UTC(), updatedBuild.UpdatedAt.UTC(), "should not have bumped updated_at")
+				require.Equal(t, bld.Deadline.UTC(), updatedBuild.Deadline.UTC(), "should not have bumped deadline")
+			} else {
+				require.NotEqual(t, bld.UpdatedAt.UTC(), updatedBuild.UpdatedAt.UTC(), "should have bumped updated_at")
+				expectedDeadline := approxBumpTime.Add(tt.expectedBump).UTC()
+				require.WithinDuration(t, expectedDeadline, updatedBuild.Deadline.UTC(), 15*time.Second, "unexpected deadline after bump")
+			}
+		})
+	}
+}
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -657,6 +657,13 @@ func (q *querier) AcquireProvisionerJob(ctx context.Context, arg database.Acquir
 	return q.db.AcquireProvisionerJob(ctx, arg)
 }
 
+func (q *querier) ActivityBumpWorkspace(ctx context.Context, arg uuid.UUID) error {
+	fetch := func(ctx context.Context, arg uuid.UUID) (database.Workspace, error) {
+		return q.db.GetWorkspaceByID(ctx, arg)
+	}
+	return update(q.log, q.auth, fetch, q.db.ActivityBumpWorkspace)(ctx, arg)
+}
+
 func (q *querier) CleanTailnetCoordinators(ctx context.Context) error {
 	if err := q.authorizeContext(ctx, rbac.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
 		return err

diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -685,6 +685,13 @@ func (q *FakeQuerier) GetActiveDBCryptKeys(_ context.Context) ([]database.DBCryp
 	return ks, nil
 }
 
+func minTime(t, u time.Time) time.Time {
+	if t.Before(u) {
+		return t
+	}
+	return u
+}
+
 func (*FakeQuerier) AcquireLock(_ context.Context, _ int64) error {
 	return xerrors.New("AcquireLock must only be called within a transaction")
 }
@@ -744,6 +751,37 @@ func (q *FakeQuerier) AcquireProvisionerJob(_ context.Context, arg database.Acqu
 	return database.ProvisionerJob{}, sql.ErrNoRows
 }
 
+func (q *FakeQuerier) ActivityBumpWorkspace(ctx context.Context, workspaceID uuid.UUID) error {
+	err := validateDatabaseType(workspaceID)
+	if err != nil {
+		return err
+	}
+
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	workspace, err := q.getWorkspaceByIDNoLock(ctx, workspaceID)
+	if err != nil {
+		return err
+	}
+	latestBuild, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, workspaceID)
+	if err != nil {
+		return err
+	}
+
+	for i := range q.workspaceBuilds {
+		if q.workspaceBuilds[i].BuildNumber != latestBuild.BuildNumber {
+			continue
+		}
+		newDeadline := q.workspaceBuilds[i].Deadline.Add(time.Duration(workspace.Ttl.Int64))
+		q.workspaceBuilds[i].UpdatedAt = latestBuild.UpdatedAt
+		q.workspaceBuilds[i].Deadline = minTime(newDeadline, q.workspaceBuilds[i].MaxDeadline)
+		return nil
+	}
+
+	return sql.ErrNoRows
+}
+
 func (*FakeQuerier) CleanTailnetCoordinators(_ context.Context) error {
 	return ErrUnimplemented
 }

diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go