Merge branch 'main' into matifali/devcontainer

matifali · web-flow · commit 7fc87bf7323e · 2023-06-27T16:01:43.000+03:00
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -35,6 +35,7 @@ import (
 	"go.uber.org/goleak"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/exp/maps"
+	"golang.org/x/exp/slices"
 	"golang.org/x/xerrors"
 	"tailscale.com/net/speedtest"
 	"tailscale.com/tailcfg"
@@ -1099,7 +1100,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
-			StartupScript:        "sleep 5",
+			StartupScript:        "sleep 3",
 			StartupScriptTimeout: time.Nanosecond,
 		}, 0)
 
@@ -1111,10 +1112,10 @@ func TestAgent_Lifecycle(t *testing.T) {
 		var got []codersdk.WorkspaceAgentLifecycle
 		assert.Eventually(t, func() bool {
 			got = client.getLifecycleStates()
-			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+			return slices.Contains(got, want[len(want)-1])
 		}, testutil.WaitShort, testutil.IntervalMedium)
 
-		require.Equal(t, want, got)
+		require.Equal(t, want, got[:len(want)])
 	})
 
 	t.Run("StartError", func(t *testing.T) {
@@ -1133,10 +1134,10 @@ func TestAgent_Lifecycle(t *testing.T) {
 		var got []codersdk.WorkspaceAgentLifecycle
 		assert.Eventually(t, func() bool {
 			got = client.getLifecycleStates()
-			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+			return slices.Contains(got, want[len(want)-1])
 		}, testutil.WaitShort, testutil.IntervalMedium)
 
-		require.Equal(t, want, got)
+		require.Equal(t, want, got[:len(want)])
 	})
 
 	t.Run("Ready", func(t *testing.T) {
@@ -1165,14 +1166,12 @@ func TestAgent_Lifecycle(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
-			ShutdownScript:       "sleep 5",
+			ShutdownScript:       "sleep 3",
 			StartupScriptTimeout: 30 * time.Second,
 		}, 0)
 
-		var ready []codersdk.WorkspaceAgentLifecycle
 		assert.Eventually(t, func() bool {
-			ready = client.getLifecycleStates()
-			return len(ready) > 0 && ready[len(ready)-1] == codersdk.WorkspaceAgentLifecycleReady
+			return slices.Contains(client.getLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
 		}, testutil.WaitShort, testutil.IntervalMedium)
 
 		// Start close asynchronously so that we an inspect the state.
@@ -1187,30 +1186,30 @@ func TestAgent_Lifecycle(t *testing.T) {
 		})
 
 		want := []codersdk.WorkspaceAgentLifecycle{
+			codersdk.WorkspaceAgentLifecycleStarting,
+			codersdk.WorkspaceAgentLifecycleReady,
 			codersdk.WorkspaceAgentLifecycleShuttingDown,
 		}
 
 		var got []codersdk.WorkspaceAgentLifecycle
 		assert.Eventually(t, func() bool {
-			got = client.getLifecycleStates()[len(ready):]
-			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+			got = client.getLifecycleStates()
+			return slices.Contains(got, want[len(want)-1])
 		}, testutil.WaitShort, testutil.IntervalMedium)
 
-		require.Equal(t, want, got)
+		require.Equal(t, want, got[:len(want)])
 	})
 
 	t.Run("ShutdownTimeout", func(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
-			ShutdownScript:        "sleep 5",
+			ShutdownScript:        "sleep 3",
 			ShutdownScriptTimeout: time.Nanosecond,
 		}, 0)
 
-		var ready []codersdk.WorkspaceAgentLifecycle
 		assert.Eventually(t, func() bool {
-			ready = client.getLifecycleStates()
-			return len(ready) > 0 && ready[len(ready)-1] == codersdk.WorkspaceAgentLifecycleReady
+			return slices.Contains(client.getLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
 		}, testutil.WaitShort, testutil.IntervalMedium)
 
 		// Start close asynchronously so that we an inspect the state.
@@ -1225,17 +1224,19 @@ func TestAgent_Lifecycle(t *testing.T) {
 		})
 
 		want := []codersdk.WorkspaceAgentLifecycle{
+			codersdk.WorkspaceAgentLifecycleStarting,
+			codersdk.WorkspaceAgentLifecycleReady,
 			codersdk.WorkspaceAgentLifecycleShuttingDown,
 			codersdk.WorkspaceAgentLifecycleShutdownTimeout,
 		}
 
 		var got []codersdk.WorkspaceAgentLifecycle
 		assert.Eventually(t, func() bool {
-			got = client.getLifecycleStates()[len(ready):]
-			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+			got = client.getLifecycleStates()
+			return slices.Contains(got, want[len(want)-1])
 		}, testutil.WaitShort, testutil.IntervalMedium)
 
-		require.Equal(t, want, got)
+		require.Equal(t, want, got[:len(want)])
 	})
 
 	t.Run("ShutdownError", func(t *testing.T) {
@@ -1246,10 +1247,8 @@ func TestAgent_Lifecycle(t *testing.T) {
 			ShutdownScriptTimeout: 30 * time.Second,
 		}, 0)
 
-		var ready []codersdk.WorkspaceAgentLifecycle
 		assert.Eventually(t, func() bool {
-			ready = client.getLifecycleStates()
-			return len(ready) > 0 && ready[len(ready)-1] == codersdk.WorkspaceAgentLifecycleReady
+			return slices.Contains(client.getLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
 		}, testutil.WaitShort, testutil.IntervalMedium)
 
 		// Start close asynchronously so that we an inspect the state.
@@ -1264,17 +1263,19 @@ func TestAgent_Lifecycle(t *testing.T) {
 		})
 
 		want := []codersdk.WorkspaceAgentLifecycle{
+			codersdk.WorkspaceAgentLifecycleStarting,
+			codersdk.WorkspaceAgentLifecycleReady,
 			codersdk.WorkspaceAgentLifecycleShuttingDown,
 			codersdk.WorkspaceAgentLifecycleShutdownError,
 		}
 
 		var got []codersdk.WorkspaceAgentLifecycle
 		assert.Eventually(t, func() bool {
-			got = client.getLifecycleStates()[len(ready):]
-			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+			got = client.getLifecycleStates()
+			return slices.Contains(got, want[len(want)-1])
 		}, testutil.WaitShort, testutil.IntervalMedium)
 
-		require.Equal(t, want, got)
+		require.Equal(t, want, got[:len(want)])
 	})
 
 	t.Run("ShutdownScriptOnce", func(t *testing.T) {
diff --git a/cli/agent.go b/cli/agent.go
@@ -106,12 +106,12 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 			// Spawn a reaper so that we don't accumulate a ton
 			// of zombie processes.
 			if reaper.IsInitProcess() && !noReap && isLinux {
-				logWriter := &lumberjack.Logger{
+				logWriter := &lumberjackWriteCloseFixer{w: &lumberjack.Logger{
 					Filename: filepath.Join(logDir, "coder-agent-init.log"),
 					MaxSize:  5, // MB
 					// Without this, rotated logs will never be deleted.
 					MaxBackups: 1,
-				}
+				}}
 				defer logWriter.Close()
 
 				sinks = append(sinks, sloghuman.Sink(logWriter))
@@ -149,14 +149,12 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 			// reaper.
 			go DumpHandler(ctx)
 
-			ljLogger := &lumberjack.Logger{
+			logWriter := &lumberjackWriteCloseFixer{w: &lumberjack.Logger{
 				Filename: filepath.Join(logDir, "coder-agent.log"),
 				MaxSize:  5, // MB
 				// Without this, rotated logs will never be deleted.
 				MaxBackups: 1,
-			}
-			defer ljLogger.Close()
-			logWriter := &closeWriter{w: ljLogger}
+			}}
 			defer logWriter.Close()
 
 			sinks = append(sinks, sloghuman.Sink(logWriter))
@@ -403,24 +401,24 @@ func ServeHandler(ctx context.Context, logger slog.Logger, handler http.Handler,
 	}
 }
 
-// closeWriter is a wrapper around an io.WriteCloser that prevents
-// writes after Close. This is necessary because lumberjack will
-// re-open the file on write.
-type closeWriter struct {
+// lumberjackWriteCloseFixer is a wrapper around an io.WriteCloser that
+// prevents writes after Close. This is necessary because lumberjack
+// re-opens the file on Write.
+type lumberjackWriteCloseFixer struct {
 	w      io.WriteCloser
 	mu     sync.Mutex // Protects following.
 	closed bool
 }
 
-func (c *closeWriter) Close() error {
+func (c *lumberjackWriteCloseFixer) Close() error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
 	c.closed = true
 	return c.w.Close()
 }
 
-func (c *closeWriter) Write(p []byte) (int, error) {
+func (c *lumberjackWriteCloseFixer) Write(p []byte) (int, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
