Fix reconnecting PTY flake

code-asher · code-asher · commit b054a12cc784 · 2023-08-14T14:34:30.000-08:00
I forgot that waiting on the cond releases the lock so it was possible
to get pty output after writing the buffer but before adding the pty to
the map.  To fix, add the pty to the map while under the same lock where
we write to the buffer.

The rest does not need to be behind the lock so I moved it out of
doAttach, and that also means we no longer need
waitForStateOrContextLocked.

Also, this can hit a logger error saying the attach failed which fails
the tests however it is not that the attach failed, just that the
process already ran and exited, so remove the error there for now.
diff --git a/agent/reconnectingpty/buffered.go b/agent/reconnectingpty/buffered.go
@@ -178,7 +178,7 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	err := rpty.doAttach(ctx, connID, conn, height, width, logger)
+	err := rpty.doAttach(connID, conn)
 	if err != nil {
 		return err
 	}
@@ -189,15 +189,34 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
 		delete(rpty.activeConns, connID)
 	}()
 
+	state, _ := rpty.state.waitForStateOrContext(ctx, StateReady)
+	if state != StateReady {
+		// This could be real failure if the command failed to start but it could
+		// also be perfectly fine if it was a short-lived command that exited before
+		// the attach completed (happens frequently in tests).  If we need to
+		// distinguish then we will have to check the process's exit code.
+		return nil
+	}
+
+	go heartbeat(ctx, rpty.timer, rpty.timeout)
+
+	// Resize the PTY to initial height + width.
+	err = rpty.ptty.Resize(height, width)
+	if err != nil {
+		// We can continue after this, it's not fatal!
+		logger.Warn(ctx, "reconnecting PTY initial resize failed, but will continue", slog.Error(err))
+		rpty.metrics.WithLabelValues("resize").Add(1)
+	}
+
 	// Pipe conn -> pty and block.  pty -> conn is handled in newBuffered().
 	readConnLoop(ctx, conn, rpty.ptty, rpty.metrics, logger)
 	return nil
 }
 
-// doAttach adds the connection to the map, replays the buffer, and starts the
-// heartbeat.  It exists separately only so we can defer the mutex unlock which
-// is not possible in Attach since it blocks.
-func (rpty *bufferedReconnectingPTY) doAttach(ctx context.Context, connID string, conn net.Conn, height, width uint16, logger slog.Logger) error {
+// doAttach adds the connection to the map and replays the buffer.  It exists
+// separately only for convenience to defer the mutex unlock which is not
+// possible in Attach since it blocks.
+func (rpty *bufferedReconnectingPTY) doAttach(connID string, conn net.Conn) error {
 	rpty.state.cond.L.Lock()
 	defer rpty.state.cond.L.Unlock()
 
@@ -211,21 +230,6 @@ func (rpty *bufferedReconnectingPTY) doAttach(ctx context.Context, connID string
 		return xerrors.Errorf("write buffer to conn: %w", err)
 	}
 
-	state, err := rpty.state.waitForStateOrContextLocked(ctx, StateReady)
-	if state != StateReady {
-		return xerrors.Errorf("reconnecting pty ready wait: %w", err)
-	}
-
-	go heartbeat(ctx, rpty.timer, rpty.timeout)
-
-	// Resize the PTY to initial height + width.
-	err = rpty.ptty.Resize(height, width)
-	if err != nil {
-		// We can continue after this, it's not fatal!
-		logger.Warn(ctx, "reconnecting PTY initial resize failed, but will continue", slog.Error(err))
-		rpty.metrics.WithLabelValues("resize").Add(1)
-	}
-
 	rpty.activeConns[connID] = conn
 
 	return nil
diff --git a/agent/reconnectingpty/reconnectingpty.go b/agent/reconnectingpty/reconnectingpty.go
@@ -171,12 +171,7 @@ func (s *ptyState) waitForState(state State) (State, error) {
 func (s *ptyState) waitForStateOrContext(ctx context.Context, state State) (State, error) {
 	s.cond.L.Lock()
 	defer s.cond.L.Unlock()
-	return s.waitForStateOrContextLocked(ctx, state)
-}
 
-// waitForStateOrContextLocked is the same as waitForStateOrContext except it
-// assumes the caller has already locked cond.
-func (s *ptyState) waitForStateOrContextLocked(ctx context.Context, state State) (State, error) {
 	nevermind := make(chan struct{})
 	defer close(nevermind)
 	go func() {
