fix: avoid vscodessh exit when server restarts (#13970)

deansheather · web-flow · commit 38c7dcda942c · 2024-07-22T17:19:44.000Z
Mitigates an issue where vscodessh would restart when the control plane
restarts, causing the entire SSH session to be reestablished.
diff --git a/cli/vscodessh.go b/cli/vscodessh.go
@@ -151,7 +151,11 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
 			// command via the ProxyCommand SSH option.
 			pid := os.Getppid()
 
-			logger := inv.Logger
+			// Use a stripped down writer that doesn't sync, otherwise you get
+			// "failed to sync sloghuman: sync /dev/stderr: The handle is
+			// invalid" on Windows. Syncing isn't required for stdout/stderr
+			// anyways.
+			logger := inv.Logger.AppendSinks(sloghuman.Sink(slogWriter{w: inv.Stderr})).Leveled(slog.LevelDebug)
 			if logDir != "" {
 				logFilePath := filepath.Join(logDir, fmt.Sprintf("%d.log", pid))
 				logFile, err := fs.OpenFile(logFilePath, os.O_CREATE|os.O_WRONLY, 0o600)
@@ -160,7 +164,7 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
 				}
 				dc := cliutil.DiscardAfterClose(logFile)
 				defer dc.Close()
-				logger = logger.AppendSinks(sloghuman.Sink(dc)).Leveled(slog.LevelDebug)
+				logger = logger.AppendSinks(sloghuman.Sink(dc))
 			}
 			if r.disableDirect {
 				logger.Info(ctx, "direct connections disabled")
@@ -204,31 +208,48 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
 			// command via the ProxyCommand SSH option.
 			networkInfoFilePath := filepath.Join(networkInfoDir, fmt.Sprintf("%d.json", pid))
 
-			statsErrChan := make(chan error, 1)
+			var (
+				firstErrTime time.Time
+				errCh        = make(chan error, 1)
+			)
 			cb := func(start, end time.Time, virtual, _ map[netlogtype.Connection]netlogtype.Counts) {
-				sendErr := func(err error) {
+				sendErr := func(tolerate bool, err error) {
+					logger.Error(ctx, "collect network stats", slog.Error(err))
+					// Tolerate up to 1 minute of errors.
+					if tolerate {
+						if firstErrTime.IsZero() {
+							logger.Info(ctx, "tolerating network stats errors for up to 1 minute")
+							firstErrTime = time.Now()
+						}
+						if time.Since(firstErrTime) < time.Minute {
+							return
+						}
+					}
+
 					select {
-					case statsErrChan <- err:
+					case errCh <- err:
 					default:
 					}
 				}
 
 				stats, err := collectNetworkStats(ctx, agentConn, start, end, virtual)
 				if err != nil {
-					sendErr(err)
+					sendErr(true, err)
 					return
 				}
 
 				rawStats, err := json.Marshal(stats)
 				if err != nil {
-					sendErr(err)
+					sendErr(false, err)
 					return
 				}
 				err = afero.WriteFile(fs, networkInfoFilePath, rawStats, 0o600)
 				if err != nil {
-					sendErr(err)
+					sendErr(false, err)
 					return
 				}
+
+				firstErrTime = time.Time{}
 			}
 
 			now := time.Now()
@@ -238,7 +259,7 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
 			select {
 			case <-ctx.Done():
 				return nil
-			case err := <-statsErrChan:
+			case err := <-errCh:
 				return err
 			}
 		},
@@ -280,6 +301,18 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
 	return cmd
 }
 
+// slogWriter wraps an io.Writer and removes all other methods (such as Sync),
+// which may cause undesired/broken behavior.
+type slogWriter struct {
+	w io.Writer
+}
+
+var _ io.Writer = slogWriter{}
+
+func (s slogWriter) Write(p []byte) (n int, err error) {
+	return s.w.Write(p)
+}
+
 type sshNetworkStats struct {
 	P2P              bool               `json:"p2p"`
 	Latency          float64            `json:"latency"`
