refactor(coderd): collapse activityBumpWorkspace into a single query (#9652)

* Adds unit-style tests for activityBumpWorkspace
* Ports logic of activityBumpWorkspace to a SQL query
* Updates activityBumpWorkspace to call above query

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

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,67 @@ 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
+	}
+
+	now := dbtime.Now()
+	for i := range q.workspaceBuilds {
+		if q.workspaceBuilds[i].BuildNumber != latestBuild.BuildNumber {
+			continue
+		}
+		// If the build is not active, do not bump.
+		if q.workspaceBuilds[i].Transition != database.WorkspaceTransitionStart {
+			return nil
+		}
+		// If the provisioner job is not completed, do not bump.
+		pj, err := q.getProvisionerJobByIDNoLock(ctx, q.workspaceBuilds[i].JobID)
+		if err != nil {
+			return err
+		}
+		if !pj.CompletedAt.Valid {
+			return nil
+		}
+		// Do not bump if the deadline is not set.
+		if q.workspaceBuilds[i].Deadline.IsZero() {
+			return nil
+		}
+		// Only bump if 5% of the deadline has passed.
+		ttlDur := time.Duration(workspace.Ttl.Int64)
+		ttlDur95 := ttlDur - (ttlDur / 20)
+		minBumpDeadline := q.workspaceBuilds[i].Deadline.Add(-ttlDur95)
+		if now.Before(minBumpDeadline) {
+			return nil
+		}
+
+		// Bump.
+		newDeadline := now.Add(ttlDur)
+		q.workspaceBuilds[i].UpdatedAt = now
+		if !q.workspaceBuilds[i].MaxDeadline.IsZero() {
+			q.workspaceBuilds[i].Deadline = minTime(newDeadline, q.workspaceBuilds[i].MaxDeadline)
+		} else {
+			q.workspaceBuilds[i].Deadline = newDeadline
+		}
+		return nil
+	}
+
+	return sql.ErrNoRows
+}
+
 func (*FakeQuerier) CleanTailnetCoordinators(_ context.Context) error {
 	return ErrUnimplemented
 }
@@ -4741,6 +4809,7 @@ func (q *FakeQuerier) InsertWorkspaceBuild(_ context.Context, arg database.Inser
 		JobID:             arg.JobID,
 		ProvisionerState:  arg.ProvisionerState,
 		Deadline:          arg.Deadline,
+		MaxDeadline:       arg.MaxDeadline,
 		Reason:            arg.Reason,
 	}
 	q.workspaceBuilds = append(q.workspaceBuilds, workspaceBuild)

coderd/database/dbmetrics/dbmetrics.go

@@ -93,6 +93,13 @@ func (m metricsStore) AcquireProvisionerJob(ctx context.Context, arg database.Ac
 	return provisionerJob, err
 }
 
+func (m metricsStore) ActivityBumpWorkspace(ctx context.Context, arg uuid.UUID) error {
+	start := time.Now()
+	r0 := m.s.ActivityBumpWorkspace(ctx, arg)
+	m.queryLatencies.WithLabelValues("ActivityBumpWorkspace").Observe(time.Since(start).Seconds())
+	return r0
+}
+
 func (m metricsStore) CleanTailnetCoordinators(ctx context.Context) error {
 	start := time.Now()
 	err := m.s.CleanTailnetCoordinators(ctx)

coderd/database/dbmock/dbmock.go

@@ -68,6 +68,20 @@ func (mr *MockStoreMockRecorder) AcquireProvisionerJob(arg0, arg1 interface{}) *
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcquireProvisionerJob", reflect.TypeOf((*MockStore)(nil).AcquireProvisionerJob), arg0, arg1)
 }
 
+// ActivityBumpWorkspace mocks base method.
+func (m *MockStore) ActivityBumpWorkspace(arg0 context.Context, arg1 uuid.UUID) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ActivityBumpWorkspace", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// ActivityBumpWorkspace indicates an expected call of ActivityBumpWorkspace.
+func (mr *MockStoreMockRecorder) ActivityBumpWorkspace(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ActivityBumpWorkspace", reflect.TypeOf((*MockStore)(nil).ActivityBumpWorkspace), arg0, arg1)
+}
+
 // CleanTailnetCoordinators mocks base method.
 func (m *MockStore) CleanTailnetCoordinators(arg0 context.Context) error {
 	m.ctrl.T.Helper()

coderd/database/querier.go

@@ -24,6 +24,13 @@ type sqlcQuerier interface {
 	// multiple provisioners from acquiring the same jobs. See:
 	// https://www.postgresql.org/docs/9.5/sql-select.html#SQL-FOR-UPDATE-SHARE
 	AcquireProvisionerJob(ctx context.Context, arg AcquireProvisionerJobParams) (ProvisionerJob, error)
+	// 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.
+	// We only bump if workspace shutdown is manual.
+	// We only bump when 5% of the deadline has elapsed.
+	ActivityBumpWorkspace(ctx context.Context, workspaceID uuid.UUID) error
 	CleanTailnetCoordinators(ctx context.Context) error
 	DeleteAPIKeyByID(ctx context.Context, id string) error
 	DeleteAPIKeysByUserID(ctx context.Context, userID uuid.UUID) error

coderd/database/queries.sql.go

@@ -15,6 +15,52 @@ import (
 	"github.com/sqlc-dev/pqtype"
 )
 
+const activityBumpWorkspace = `-- name: ActivityBumpWorkspace :exec
+WITH latest AS (
+	SELECT
+		workspace_builds.id::uuid AS build_id,
+		workspace_builds.deadline::timestamp AS build_deadline,
+		workspace_builds.max_deadline::timestamp AS build_max_deadline,
+		workspace_builds.transition AS build_transition,
+		provisioner_jobs.completed_at::timestamp AS job_completed_at,
+		(workspaces.ttl / 1000 / 1000 / 1000 || ' seconds')::interval AS ttl_interval
+	FROM workspace_builds
+	JOIN provisioner_jobs
+		ON provisioner_jobs.id = workspace_builds.job_id
+	JOIN workspaces
+		ON workspaces.id = workspace_builds.workspace_id
+	WHERE workspace_builds.workspace_id = $1::uuid
+	ORDER BY workspace_builds.build_number DESC
+	LIMIT 1
+)
+UPDATE
+	workspace_builds wb
+SET
+	updated_at = NOW(),
+	deadline = CASE
+		WHEN l.build_max_deadline = '0001-01-01 00:00:00+00'
+		THEN NOW() + l.ttl_interval
+		ELSE LEAST(NOW() + l.ttl_interval, l.build_max_deadline)
+	END
+FROM latest l
+WHERE wb.id = l.build_id
+AND l.job_completed_at IS NOT NULL
+AND l.build_transition = 'start'
+AND l.build_deadline != '0001-01-01 00:00:00+00'
+AND l.build_deadline - (l.ttl_interval * 0.95) < NOW()
+`
+
+// 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.
+// We only bump if workspace shutdown is manual.
+// We only bump when 5% of the deadline has elapsed.
+func (q *sqlQuerier) ActivityBumpWorkspace(ctx context.Context, workspaceID uuid.UUID) error {
+	_, err := q.db.ExecContext(ctx, activityBumpWorkspace, workspaceID)
+	return err
+}
+
 const deleteAPIKeyByID = `-- name: DeleteAPIKeyByID :exec
 DELETE FROM
 	api_keys

coderd/database/queries/activitybump.sql

@@ -0,0 +1,40 @@
+-- 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.
+-- name: ActivityBumpWorkspace :exec
+WITH latest AS (
+	SELECT
+		workspace_builds.id::uuid AS build_id,
+		workspace_builds.deadline::timestamp AS build_deadline,
+		workspace_builds.max_deadline::timestamp AS build_max_deadline,
+		workspace_builds.transition AS build_transition,
+		provisioner_jobs.completed_at::timestamp AS job_completed_at,
+		(workspaces.ttl / 1000 / 1000 / 1000 || ' seconds')::interval AS ttl_interval
+	FROM workspace_builds
+	JOIN provisioner_jobs
+		ON provisioner_jobs.id = workspace_builds.job_id
+	JOIN workspaces
+		ON workspaces.id = workspace_builds.workspace_id
+	WHERE workspace_builds.workspace_id = @workspace_id::uuid
+	ORDER BY workspace_builds.build_number DESC
+	LIMIT 1
+)
+UPDATE
+	workspace_builds wb
+SET
+	updated_at = NOW(),
+	deadline = CASE
+		WHEN l.build_max_deadline = '0001-01-01 00:00:00+00'
+		THEN NOW() + l.ttl_interval
+		ELSE LEAST(NOW() + l.ttl_interval, l.build_max_deadline)
+	END
+FROM latest l
+WHERE wb.id = l.build_id
+AND l.job_completed_at IS NOT NULL
+AND l.build_transition = 'start'
+-- We only bump if workspace shutdown is manual.
+AND l.build_deadline != '0001-01-01 00:00:00+00'
+-- We only bump when 5% of the deadline has elapsed.
+AND l.build_deadline - (l.ttl_interval * 0.95) < NOW()
+;