coder · johnstcn · Jun 9, 2022 · Jun 7, 2022 · Jun 7, 2022 · Jun 8, 2022
diff --git a/cli/ttl.go b/cli/ttl.go
@@ -1,12 +1,14 @@
 package cli
 
 import (
+	"errors"
 	"fmt"
 	"time"
 
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/codersdk"
 )
 
@@ -89,6 +91,30 @@ func ttlset() *cobra.Command {
 				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "warning: ttl rounded down to %s\n", truncated)
 			}
 
+			if changed, newDeadline := changedNewDeadline(workspace, truncated); changed {
+				// For the purposes of the user, "less than a minute" is essentially the same as "immediately".
+				timeRemaining := time.Until(newDeadline).Truncate(time.Minute)
+				humanRemaining := "in " + timeRemaining.String()
+				if timeRemaining <= 0 {
+					humanRemaining = "immediately"
+				}
+				_, err = cliui.Prompt(cmd, cliui.PromptOptions{
+					Text: fmt.Sprintf(
+						"Workspace %q will be stopped %s. Are you sure?",
+						workspace.Name,
+						humanRemaining,
+					),
+					Default:   "yes",
+					IsConfirm: true,
+				})
+				if err != nil {
+					if errors.Is(err, cliui.Canceled) {
+						return nil
+					}
+					return err
+				}
+			}
+
 			millis := truncated.Milliseconds()
 			if err = client.UpdateWorkspaceTTL(cmd.Context(), workspace.ID, codersdk.UpdateWorkspaceTTLRequest{
 				TTLMillis: &millis,
@@ -131,3 +157,18 @@ func ttlunset() *cobra.Command {
 		},
 	}
 }
+
+func changedNewDeadline(ws codersdk.Workspace, newTTL time.Duration) (changed bool, newDeadline time.Time) {
+	if ws.LatestBuild.Transition != codersdk.WorkspaceTransitionStart {
+		// not running
+		return false, newDeadline
+	}
+
+	if ws.LatestBuild.Job.CompletedAt == nil {
+		// still building
+		return false, newDeadline
+	}
+
+	newDeadline = ws.LatestBuild.Job.CompletedAt.Add(newTTL)
+	return true, newDeadline
+}
diff --git a/cli/ttl_test.go b/cli/ttl_test.go
@@ -3,16 +3,19 @@ package cli_test
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/cli/clitest"
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/coderd/util/ptr"
 	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/pty/ptytest"
 )
 
 func TestTTL(t *testing.T) {
@@ -22,33 +25,29 @@ func TestTTL(t *testing.T) {
 		t.Parallel()
 
 		var (
-			ctx       = context.Background()
 			client    = coderdtest.New(t, &coderdtest.Options{IncludeProvisionerD: true})
 			user      = coderdtest.CreateFirstUser(t, client)
 			version   = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
 			_         = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
-			project   = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
-			workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, project.ID)
+			template  = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+			ttl       = 7*time.Hour + 30*time.Minute + 30*time.Second
+			workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
+				cwr.TTLMillis = ptr.Ref(ttl.Milliseconds())
+			})
 			cmdArgs   = []string{"ttl", "show", workspace.Name}
-			ttl       = 8*time.Hour + 30*time.Minute + 30*time.Second
 			stdoutBuf = &bytes.Buffer{}
 		)
 
-		err := client.UpdateWorkspaceTTL(ctx, workspace.ID, codersdk.UpdateWorkspaceTTLRequest{
-			TTLMillis: ptr.Ref(ttl.Milliseconds()),
-		})
-		require.NoError(t, err)
-
 		cmd, root := clitest.New(t, cmdArgs...)
 		clitest.SetupConfig(t, client, root)
 		cmd.SetOut(stdoutBuf)
 
-		err = cmd.Execute()
+		err := cmd.Execute()
 		require.NoError(t, err, "unexpected error")
 		require.Equal(t, ttl.Truncate(time.Minute).String(), strings.TrimSpace(stdoutBuf.String()))
 	})
 
-	t.Run("SetUnsetOK", func(t *testing.T) {
+	t.Run("UnsetOK", func(t *testing.T) {
 		t.Parallel()
 
 		var (
@@ -58,9 +57,11 @@ func TestTTL(t *testing.T) {
 			version   = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
 			_         = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
 			project   = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
-			workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, project.ID)
 			ttl       = 8*time.Hour + 30*time.Minute + 30*time.Second
-			cmdArgs   = []string{"ttl", "set", workspace.Name, ttl.String()}
+			workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, project.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
+				cwr.TTLMillis = ptr.Ref(ttl.Milliseconds())
+			})
+			cmdArgs   = []string{"ttl", "unset", workspace.Name}
 			stdoutBuf = &bytes.Buffer{}
 		)
 
@@ -71,24 +72,52 @@ func TestTTL(t *testing.T) {
 		err := cmd.Execute()
 		require.NoError(t, err, "unexpected error")
 
-		// Ensure ttl updated
+		// Ensure ttl unset
 		updated, err := client.Workspace(ctx, workspace.ID)
 		require.NoError(t, err, "fetch updated workspace")
-		require.Equal(t, ttl.Truncate(time.Minute), time.Duration(*updated.TTLMillis)*time.Millisecond)
-		require.Contains(t, stdoutBuf.String(), "warning: ttl rounded down")
+		require.Nil(t, updated.TTLMillis, "expected ttl to not be set")
+	})
 
-		// unset schedule
-		cmd, root = clitest.New(t, "ttl", "unset", workspace.Name)
-		clitest.SetupConfig(t, client, root)
-		cmd.SetOut(stdoutBuf)
+	t.Run("SetOK", func(t *testing.T) {
+		t.Parallel()
 
-		err = cmd.Execute()
-		require.NoError(t, err, "unexpected error")
+		var (
+			ctx       = context.Background()
+			client    = coderdtest.New(t, &coderdtest.Options{IncludeProvisionerD: true})
+			user      = coderdtest.CreateFirstUser(t, client)
+			version   = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
+			_         = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+			project   = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+			ttl       = 8*time.Hour + 30*time.Minute + 30*time.Second
+			workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, project.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
+				cwr.TTLMillis = ptr.Ref(ttl.Milliseconds())
+			})
+			_       = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+			cmdArgs = []string{"ttl", "set", workspace.Name, ttl.String()}
+			done    = make(chan struct{})
+		)
 
+		cmd, root := clitest.New(t, cmdArgs...)
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetOut(pty.Output())
+
+		go func() {
+			defer close(done)
+			err := cmd.Execute()
+			assert.NoError(t, err, "unexpected error")
+		}()
+
+		pty.ExpectMatch(fmt.Sprintf("warning: ttl rounded down to %s", ttl.Truncate(time.Minute)))
+		pty.ExpectMatch(fmt.Sprintf("Workspace %q will be stopped in 8h29m0s. Are you sure?", workspace.Name))
+		pty.WriteLine("yes")
 		// Ensure ttl updated
-		updated, err = client.Workspace(ctx, workspace.ID)
+		updated, err := client.Workspace(ctx, workspace.ID)
 		require.NoError(t, err, "fetch updated workspace")
-		require.Nil(t, updated.TTLMillis, "expected ttl to not be set")
+		require.Equal(t, ttl.Truncate(time.Minute), time.Duration(*updated.TTLMillis)*time.Millisecond)
+
+		<-done
 	})
 
 	t.Run("ZeroInvalid", func(t *testing.T) {

diff --git a/coderd/autobuild/executor/lifecycle_executor_test.go b/coderd/autobuild/executor/lifecycle_executor_test.go
@@ -440,18 +440,41 @@ func TestExecutorWorkspaceAutostopNoWaitChangedMyMind(t *testing.T) {
 	err := client.UpdateWorkspaceTTL(ctx, workspace.ID, codersdk.UpdateWorkspaceTTLRequest{TTLMillis: nil})
 	require.NoError(t, err)
 
-	// When: the autobuild executor ticks after the deadline
+	// Then: the deadline should be the zero value
+	updated := coderdtest.MustWorkspace(t, client, workspace.ID)
+	assert.Zero(t, updated.LatestBuild.Deadline)
+
+	// When: the autobuild executor ticks after the original deadline
 	go func() {
 		tickCh <- workspace.LatestBuild.Deadline.Add(time.Minute)
-		close(tickCh)
 	}()
 
-	// Then: the workspace should still stop - sorry!
+	// Then: the workspace should not stop
 	stats := <-statsCh
 	assert.NoError(t, stats.Error)
+	assert.Len(t, stats.Transitions, 0)
+
+	// Given: the user changes their mind again and wants to enable auto-stop
+	newTTL := 8 * time.Hour
+	expectedDeadline := workspace.LatestBuild.UpdatedAt.Add(newTTL)
+	err = client.UpdateWorkspaceTTL(ctx, workspace.ID, codersdk.UpdateWorkspaceTTLRequest{TTLMillis: ptr.Ref(newTTL.Milliseconds())})
+	require.NoError(t, err)
+
+	// Then: the deadline should be updated based on the TTL
+	updated = coderdtest.MustWorkspace(t, client, workspace.ID)
+	assert.WithinDuration(t, expectedDeadline, updated.LatestBuild.Deadline, time.Minute)
+
+	// When: the relentless onward march of time continues
+	go func() {
+		tickCh <- workspace.LatestBuild.Deadline.Add(newTTL + time.Minute)
+		close(tickCh)
+	}()
+
+	// Then: the workspace should stop
+	stats = <-statsCh
+	assert.NoError(t, stats.Error)
 	assert.Len(t, stats.Transitions, 1)
-	assert.Contains(t, stats.Transitions, workspace.ID)
-	assert.Equal(t, database.WorkspaceTransitionStop, stats.Transitions[workspace.ID])
+	assert.Equal(t, stats.Transitions[workspace.ID], database.WorkspaceTransitionStop)
 }
 
 func TestExecutorAutostartMultipleOK(t *testing.T) {

diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -566,17 +566,57 @@ func (api *API) putWorkspaceTTL(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	err = api.Database.UpdateWorkspaceTTL(r.Context(), database.UpdateWorkspaceTTLParams{
-		ID:  workspace.ID,
-		Ttl: dbTTL,
+	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 := s.GetLatestWorkspaceBuildByWorkspaceID(r.Context(), workspace.ID)
+		if err != nil {
+			return xerrors.Errorf("get latest workspace build: %w", err)
+		}
+
+		if latestBuild.Transition != database.WorkspaceTransitionStart {
+			return nil // nothing to do
+		}
+
+		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))
+		}
+
+		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 updating workspace TTL.",
+			Message: "Error updating workspace time until shutdown!",
 			Detail:  err.Error(),
 		})
 		return
 	}
+
+	httpapi.Write(rw, http.StatusOK, nil)
 }
 
 func (api *API) putExtendWorkspace(rw http.ResponseWriter, r *http.Request) {