Do not hold mutex while waiting for state in screen

code-asher · code-asher · commit 7409e7a2a682 · 2023-08-09T10:41:11.000-08:00
Holding off on changing this for the buffered version until a potential
issue with the pty closing before we add to the map is sorted.

I moved the heartbeat start out as well since while we want to wait to
start that until after ready, it does not need to be behind the mutex.
diff --git a/agent/reconnectingpty/screen.go b/agent/reconnectingpty/screen.go
@@ -155,6 +155,13 @@ func (rpty *screenReconnectingPTY) Attach(ctx context.Context, _ string, conn ne
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
+	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)
+
 	ptty, process, err := rpty.doAttach(ctx, height, width, logger)
 	if err != nil {
 		return err
@@ -233,13 +240,6 @@ func (rpty *screenReconnectingPTY) doAttach(ctx context.Context, height, width u
 	rpty.mutex.Lock()
 	defer rpty.mutex.Unlock()
 
-	state, err := rpty.state.waitForStateOrContext(ctx, StateReady)
-	if state != StateReady {
-		return nil, nil, xerrors.Errorf("reconnecting pty ready wait: %w", err)
-	}
-
-	go heartbeat(ctx, rpty.timer, rpty.timeout)
-
 	logger.Debug(ctx, "spawning screen client", slog.F("screen_id", rpty.id))
 
 	// Wrap the command with screen and tie it to the connection's context.
