REFACTOR: coderd: putWorkspaceTTL: wrap in tx

coder · johnstcn · Jun 9, 2022 · Jun 7, 2022 · Jun 7, 2022 · Jun 8, 2022
commit 7846dd8640c1b28af6959c4cfe178e8bf292756d
diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -533,57 +533,52 @@ func (api *API) putWorkspaceTTL(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	err = api.Database.UpdateWorkspaceTTL(r.Context(), database.UpdateWorkspaceTTLParams{
-		ID:  workspace.ID,
-		Ttl: dbTTL,
-	})
-	if err != nil {
-		httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
-			Message: "Error updating workspace TTL!",
-			Detail:  err.Error(),
-		})
-		return
-	}
+	err = api.Database.InTx(func(s database.Store) error {
+		if err := s.UpdateWorkspaceTTL(r.Context(), database.UpdateWorkspaceTTLParams{
+			ID:  workspace.ID,
+			Ttl: dbTTL,
+		}); err != nil {
+			return xerrors.Errorf("update workspace TTL: %w", err)
+		}
 
-	// Also extend the workspace deadline if the workspace is running
-	latestBuild, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(r.Context(), workspace.ID)
-	if err != nil {
-		httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
-			Message: "Error fetching latest workspace build!",
-			Detail:  err.Error(),
-		})
-		return
-	}
+		// Also extend the workspace deadline if the workspace is running
+		latestBuild, err := s.GetLatestWorkspaceBuildByWorkspaceID(r.Context(), workspace.ID)
+		if err != nil {
+			return xerrors.Errorf("get latest workspace build: %w", err)
+		}
 
-	if latestBuild.Transition != database.WorkspaceTransitionStart {
-		httpapi.Write(rw, http.StatusOK, nil)
-		return
-	}
+		if latestBuild.Transition != database.WorkspaceTransitionStart {
+			return nil // nothing to do
+		}
 
-	if latestBuild.UpdatedAt.IsZero() {
-		// Build in progress; provisionerd should update with the new TTL.
-		httpapi.Write(rw, http.StatusOK, nil)
-		return
-	}
+		if latestBuild.UpdatedAt.IsZero() {
+			// Build in progress; provisionerd should update with the new TTL.
+			return nil
 workspaceDeadline = now.Add(time.Duration(workspace.Ttl.Int64)) 
 workspaceDeadline = now.Add(time.Duration(workspace.Ttl.Int64)) 
+		}
 
-	var newDeadline time.Time
-	if dbTTL.Valid {
-		newDeadline = latestBuild.UpdatedAt.Add(time.Duration(dbTTL.Int64))
-	}
+		var newDeadline time.Time
+		if dbTTL.Valid {
+			newDeadline = latestBuild.UpdatedAt.Add(time.Duration(dbTTL.Int64))
+		}
 
-	err = api.Database.UpdateWorkspaceBuildByID(
-		r.Context(),
-		database.UpdateWorkspaceBuildByIDParams{
-			ID:               latestBuild.ID,
-			UpdatedAt:        latestBuild.UpdatedAt,
-			ProvisionerState: latestBuild.ProvisionerState,
-			Deadline:         newDeadline,
-		},
-	)
+		if err := s.UpdateWorkspaceBuildByID(
+			r.Context(),
+			database.UpdateWorkspaceBuildByIDParams{
+				ID:               latestBuild.ID,
+				UpdatedAt:        latestBuild.UpdatedAt,
+				ProvisionerState: latestBuild.ProvisionerState,
+				Deadline:         newDeadline,
+			},
+		); err != nil {
+			return xerrors.Errorf("update workspace deadline: %w", err)
+		}
+
+		return nil
+	})
 
 	if err != nil {
 		httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
-			Message: "Internal error extending workspace deadline.",
+			Message: "Error updating workspace time until shutdown!",
 			Detail:  err.Error(),
 		})
 		return

diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go
@@ -611,34 +611,37 @@ func TestWorkspaceUpdateTTL(t *testing.T) {
 		name             string
 		ttlMillis        *int64
 		expectedError    string
-		expectedDeadline time.Time
+		expectedDeadline *time.Time
 		modifyTemplate   func(*codersdk.CreateTemplateRequest)
 	}{
 		{
-			name:          "disable ttl",
-			ttlMillis:     nil,
-			expectedError: "",
+			name:             "disable ttl",
+			ttlMillis:        nil,
+			expectedError:    "",
+			expectedDeadline: ptr.Ref(time.Time{}),
 		},
 		{
 			name:             "update ttl",
 			ttlMillis:        ptr.Ref(12 * time.Hour.Milliseconds()),
 			expectedError:    "",
-			expectedDeadline: time.Now().Add(12 * time.Hour),
+			expectedDeadline: ptr.Ref(time.Now().Add(12 * time.Hour)),
 		},
 		{
 			name:          "below minimum ttl",
 			ttlMillis:     ptr.Ref((30 * time.Second).Milliseconds()),
 			expectedError: "ttl must be at least one minute",
 		},
 		{
-			name:          "minimum ttl",
-			ttlMillis:     ptr.Ref(time.Minute.Milliseconds()),
-			expectedError: "",
+			name:             "minimum ttl",
+			ttlMillis:        ptr.Ref(time.Minute.Milliseconds()),
+			expectedError:    "",
+			expectedDeadline: ptr.Ref(time.Now().Add(time.Minute)),
 		},
 		{
-			name:          "maximum ttl",
-			ttlMillis:     ptr.Ref((24 * 7 * time.Hour).Milliseconds()),
-			expectedError: "",
+			name:             "maximum ttl",
+			ttlMillis:        ptr.Ref((24 * 7 * time.Hour).Milliseconds()),
+			expectedError:    "",
+			expectedDeadline: ptr.Ref(time.Now().Add(24 * 7 * time.Hour)),
 		},
 		{
 			name:          "above maximum ttl",
@@ -690,8 +693,8 @@ func TestWorkspaceUpdateTTL(t *testing.T) {
 			require.NoError(t, err, "fetch updated workspace")
 
 			require.Equal(t, testCase.ttlMillis, updated.TTLMillis, "expected autostop ttl to equal requested")
-			if !testCase.expectedDeadline.IsZero() {
-				require.WithinDuration(t, testCase.expectedDeadline, updated.LatestBuild.Deadline, time.Minute, "expected autostop deadline to be equal expected")
+			if testCase.expectedDeadline != nil {
+				require.WithinDuration(t, *testCase.expectedDeadline, updated.LatestBuild.Deadline, time.Minute, "expected autostop deadline to be equal expected")
 			}
 		})
 	}