Fix reconnecting PTY flake

code-asher · code-asher · commit 16c9105aa630 · 2023-08-14T14:20:24.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.
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,30 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
 		delete(rpty.activeConns, connID)
 	}()
 
+	state, err := rpty.state.waitForStateOrContext(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)
+	}
+
 	// 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 +226,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() {
