mess around

ammario · ammario · commit a1b68e55996b · 2023-06-08T04:29:54.000Z
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -76,7 +76,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			if logDirPath != "" {
 				nonce, err := cryptorand.StringCharset(cryptorand.Lower, 5)
 				if err != nil {
-					panic(err)
+					return xerrors.Errorf("generate nonce: %w", err)
 				}
 				logFilePath := filepath.Join(
 					logDirPath,
@@ -109,6 +109,19 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 				client.Logger = logger
 			}
 
+			// This WaitGroup solves for a race condition where we were logging
+			// while closing the log file in in a defer. It probably solves
+			// others too.
+			//
+			// Its position in this function is important. It must be after
+			// the logger is created but before any goroutines or wind-down
+			// defers (e.g. context cancels) are declared.
+			var wg sync.WaitGroup
+			defer func() {
+				cancel()
+				wg.Wait()
+			}()
+
 			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, codersdk.Me, inv.Args[0])
 			if err != nil {
 				return err
@@ -169,12 +182,6 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			stopPolling := tryPollWorkspaceAutostop(ctx, client, workspace)
 			defer stopPolling()
 
-			// This WaitGroup solves for a race condition where we were logging
-			// while closing the log file in in a defer. It probably solves
-			// others too.
-			var wg sync.WaitGroup
-			defer wg.Wait()
-
 			if stdio {
 				rawSSH, err := conn.SSH(ctx)
 				if err != nil {
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -714,7 +714,7 @@ type stdioConn struct {
 	io.Writer
 }
 
-func (*stdioConn) Close() (err error) {
+func (c *stdioConn) Close() (err error) {
 	return nil
 }
 

Original file line number	Diff line number	Diff line change
`@@ -714,7 +714,7 @@ type stdioConn struct {`
`714`	`714`	`io.Writer`
`715`	`715`	`}`
`716`	`716`
`717`		`-func (*stdioConn) Close() (err error) {`
	`717`	`+func (c *stdioConn) Close() (err error) {`
`718`	`718`	`return nil`
`719`	`719`	`}`
`720`	`720`