feat: make default autobuild poll intervals configurable (#1618)

johnstcn · web-flow · commit 52230fab566f · 2022-05-20T10:57:02.000Z
* feat: make default poll intervals for autobuild and ssh ttl polling configurable
diff --git a/cli/cliflag/cliflag.go b/cli/cliflag/cliflag.go
@@ -15,6 +15,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/spf13/pflag"
 )
@@ -83,6 +84,23 @@ func BoolVarP(flagset *pflag.FlagSet, ptr *bool, name string, shorthand string,
 	flagset.BoolVarP(ptr, name, shorthand, valb, fmtUsage(usage, env))
 }
 
+// DurationVarP sets a time.Duration flag on the given flag set.
+func DurationVarP(flagset *pflag.FlagSet, ptr *time.Duration, name string, shorthand string, env string, def time.Duration, usage string) {
+	val, ok := os.LookupEnv(env)
+	if !ok || val == "" {
+		flagset.DurationVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
+		return
+	}
+
+	valb, err := time.ParseDuration(val)
+	if err != nil {
+		flagset.DurationVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
+		return
+	}
+
+	flagset.DurationVarP(ptr, name, shorthand, valb, fmtUsage(usage, env))
+}
+
 func fmtUsage(u string, env string) string {
 	if env == "" {
 		return fmt.Sprintf("%s.", u)
diff --git a/cli/cliflag/cliflag_test.go b/cli/cliflag/cliflag_test.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/spf13/pflag"
 	"github.com/stretchr/testify/require"
@@ -183,6 +184,45 @@ func TestCliflag(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, def, got)
 	})
+
+	t.Run("DurationDefault", func(t *testing.T) {
+		var ptr time.Duration
+		flagset, name, shorthand, env, usage := randomFlag()
+		def, _ := cryptorand.Duration()
+
+		cliflag.DurationVarP(flagset, &ptr, name, shorthand, env, def, usage)
+		got, err := flagset.GetDuration(name)
+		require.NoError(t, err)
+		require.Equal(t, def, got)
+		require.Contains(t, flagset.FlagUsages(), usage)
+		require.Contains(t, flagset.FlagUsages(), fmt.Sprintf(" - consumes $%s", env))
+	})
+
+	t.Run("DurationEnvVar", func(t *testing.T) {
+		var ptr time.Duration
+		flagset, name, shorthand, env, usage := randomFlag()
+		envValue, _ := cryptorand.Duration()
+		t.Setenv(env, envValue.String())
+		def, _ := cryptorand.Duration()
+
+		cliflag.DurationVarP(flagset, &ptr, name, shorthand, env, def, usage)
+		got, err := flagset.GetDuration(name)
+		require.NoError(t, err)
+		require.Equal(t, envValue, got)
+	})
+
+	t.Run("DurationFailParse", func(t *testing.T) {
+		var ptr time.Duration
+		flagset, name, shorthand, env, usage := randomFlag()
+		envValue, _ := cryptorand.String(10)
+		t.Setenv(env, envValue)
+		def, _ := cryptorand.Duration()
+
+		cliflag.DurationVarP(flagset, &ptr, name, shorthand, env, def, usage)
+		got, err := flagset.GetDuration(name)
+		require.NoError(t, err)
+		require.Equal(t, def, got)
+	})
 }
 
 func randomFlag() (*pflag.FlagSet, string, string, string, string) {
diff --git a/cli/server.go b/cli/server.go
@@ -60,17 +60,18 @@ import (
 // nolint:gocyclo
 func server() *cobra.Command {
 	var (
-		accessURL       string
-		address         string
-		promEnabled     bool
-		promAddress     string
-		pprofEnabled    bool
-		pprofAddress    string
-		cacheDir        string
-		dev             bool
-		devUserEmail    string
-		devUserPassword string
-		postgresURL     string
+		accessURL             string
+		address               string
+		autobuildPollInterval time.Duration
+		promEnabled           bool
+		promAddress           string
+		pprofEnabled          bool
+		pprofAddress          string
+		cacheDir              string
+		dev                   bool
+		devUserEmail          string
+		devUserPassword       string
+		postgresURL           string
 		// provisionerDaemonCount is a uint8 to ensure a number > 0.
 		provisionerDaemonCount           uint8
 		oauth2GithubClientID             string
@@ -361,10 +362,10 @@ func server() *cobra.Command {
 				return xerrors.Errorf("notify systemd: %w", err)
 			}
 
-			lifecyclePoller := time.NewTicker(time.Minute)
-			defer lifecyclePoller.Stop()
-			lifecycleExecutor := executor.New(cmd.Context(), options.Database, logger, lifecyclePoller.C)
-			lifecycleExecutor.Run()
+			autobuildPoller := time.NewTicker(autobuildPollInterval)
+			defer autobuildPoller.Stop()
+			autobuildExecutor := executor.New(cmd.Context(), options.Database, logger, autobuildPoller.C)
+			autobuildExecutor.Run()
 
 			// Because the graceful shutdown includes cleaning up workspaces in dev mode, we're
 			// going to make it harder to accidentally skip the graceful shutdown by hitting ctrl+c
@@ -454,6 +455,7 @@ func server() *cobra.Command {
 		},
 	}
 
+	cliflag.DurationVarP(root.Flags(), &autobuildPollInterval, "autobuild-poll-interval", "", "CODER_AUTOBUILD_POLL_INTERVAL", time.Minute, "Specifies the interval at which to poll for and execute automated workspace build operations.")
 	cliflag.StringVarP(root.Flags(), &accessURL, "access-url", "", "CODER_ACCESS_URL", "", "Specifies the external URL to access Coder.")
 	cliflag.StringVarP(root.Flags(), &address, "address", "a", "CODER_ADDRESS", "127.0.0.1:3000", "The address to serve the API and dashboard.")
 	cliflag.BoolVarP(root.Flags(), &promEnabled, "prometheus-enable", "", "CODER_PROMETHEUS_ENABLE", false, "Enable serving prometheus metrics on the addressdefined by --prometheus-address.")
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -25,13 +25,14 @@ import (
 	"github.com/coder/coder/cryptorand"
 )
 
-var autostopPollInterval = 30 * time.Second
+var workspacePollInterval = time.Minute
 var autostopNotifyCountdown = []time.Duration{30 * time.Minute}
 
 func ssh() *cobra.Command {
 	var (
-		stdio   bool
-		shuffle bool
+		stdio          bool
+		shuffle        bool
+		wsPollInterval time.Duration
 	)
 	cmd := &cobra.Command{
 		Annotations: workspaceCommand,
@@ -155,6 +156,7 @@ func ssh() *cobra.Command {
 	}
 	cliflag.BoolVarP(cmd.Flags(), &stdio, "stdio", "", "CODER_SSH_STDIO", false, "Specifies whether to emit SSH output over stdin/stdout.")
 	cliflag.BoolVarP(cmd.Flags(), &shuffle, "shuffle", "", "CODER_SSH_SHUFFLE", false, "Specifies whether to choose a random workspace")
+	cliflag.DurationVarP(cmd.Flags(), &wsPollInterval, "workspace-poll-interval", "", "CODER_WORKSPACE_POLL_INTERVAL", workspacePollInterval, "Specifies how often to poll for workspace automated shutdown.")
 	_ = cmd.Flags().MarkHidden("shuffle")
 
 	return cmd
@@ -252,14 +254,14 @@ func getWorkspaceAndAgent(cmd *cobra.Command, client *codersdk.Client, orgID uui
 func tryPollWorkspaceAutostop(ctx context.Context, client *codersdk.Client, workspace codersdk.Workspace) (stop func()) {
 	lock := flock.New(filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String()))
 	condition := notifyCondition(ctx, client, workspace.ID, lock)
-	return notify.Notify(condition, autostopPollInterval, autostopNotifyCountdown...)
+	return notify.Notify(condition, workspacePollInterval, autostopNotifyCountdown...)
 }
 
 // Notify the user if the workspace is due to shutdown.
 func notifyCondition(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID, lock *flock.Flock) notify.Condition {
 	return func(now time.Time) (deadline time.Time, callback func()) {
 		// Keep trying to regain the lock.
-		locked, err := lock.TryLockContext(ctx, autostopPollInterval)
+		locked, err := lock.TryLockContext(ctx, workspacePollInterval)
 		if err != nil || !locked {
 			return time.Time{}, nil
 		}
diff --git a/cryptorand/numbers.go b/cryptorand/numbers.go
@@ -3,6 +3,7 @@ package cryptorand
 import (
 	"crypto/rand"
 	"encoding/binary"
+	"time"
 
 	"golang.org/x/xerrors"
 )
@@ -193,3 +194,13 @@ func Bool() (bool, error) {
 	// True if the least significant bit is 1
 	return i&1 == 1, nil
 }
+
+// Duration returns a random time.Duration value
+func Duration() (time.Duration, error) {
+	i, err := Int63()
+	if err != nil {
+		return time.Duration(0), err
+	}
+
+	return time.Duration(i), nil
+}
diff --git a/cryptorand/numbers_test.go b/cryptorand/numbers_test.go
@@ -175,3 +175,14 @@ func TestBool(t *testing.T) {
 	require.True(t, percentage > 48, "expected more than 48 percent of values to be true")
 	require.True(t, percentage < 52, "expected less than 52 percent of values to be true")
 }
+
+func TestDuration(t *testing.T) {
+	t.Parallel()
+
+	for i := 0; i < 20; i++ {
+		v, err := cryptorand.Duration()
+		require.NoError(t, err, "unexpected error from Duration")
+		t.Logf("value: %v <- random?", v)
+		require.True(t, v >= 0.0, "values must be positive")
+	}
+}
