Add state wait func where caller holds the lock

code-asher · code-asher · commit 1fd4e9ab50b3 · 2023-08-14T08:15:29.000-08:00
diff --git a/agent/reconnectingpty/buffered.go b/agent/reconnectingpty/buffered.go
@@ -178,39 +178,9 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	// Once we are ready, attach the active connection while we hold the mutex.
-	state, err := rpty.state.waitForStateOrContext(ctx, StateReady, func(state State) error {
-		// Write any previously stored data for the TTY.  Since the command might be
-		// short-lived and have already exited, make sure we always at least output
-		// the buffer before returning.
-		prevBuf := slices.Clone(rpty.circularBuffer.Bytes())
-		_, err := conn.Write(prevBuf)
-		if err != nil {
-			rpty.metrics.WithLabelValues("write").Add(1)
-			return xerrors.Errorf("write buffer to conn: %w", err)
-		}
-
-		if state != StateReady {
-			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)
-		}
-
-		// Store the connection for future writes.
-		rpty.activeConns[connID] = conn
-
-		return nil
-	})
-	if state != StateReady || err != nil {
-		return xerrors.Errorf("reconnecting pty ready wait: %w", err)
+	err := rpty.doAttach(ctx, connID, conn, height, width, logger)
+	if err != nil {
+		return err
 	}
 
 	defer func() {
@@ -224,6 +194,43 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
 	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 {
+	rpty.state.cond.L.Lock()
+	defer rpty.state.cond.L.Unlock()
+
+	// Write any previously stored data for the TTY.  Since the command might be
+	// short-lived and have already exited, make sure we always at least output
+	// the buffer before returning, mostly just so tests pass.
+	prevBuf := slices.Clone(rpty.circularBuffer.Bytes())
+	_, err := conn.Write(prevBuf)
+	if err != nil {
+		rpty.metrics.WithLabelValues("write").Add(1)
+		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
+}
+
 func (rpty *bufferedReconnectingPTY) Wait() {
 	_, _ = rpty.state.waitForState(StateClosing)
 }
diff --git a/agent/reconnectingpty/reconnectingpty.go b/agent/reconnectingpty/reconnectingpty.go
@@ -167,10 +167,16 @@ func (s *ptyState) waitForState(state State) (State, error) {
 }
 
 // waitForStateOrContext blocks until the state or a greater one is reached or
-// the provided context ends.  If fn is non-nil it will be ran while the lock is
-// held unless the context ends.  If fn returns an error then fn's error will
-// replace waitForStateOrContext's error.
+// the provided context ends.
 func (s *ptyState) waitForStateOrContext(ctx context.Context, state State, fn func(state State) error) (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() {
@@ -182,20 +188,12 @@ func (s *ptyState) waitForStateOrContext(ctx context.Context, state State, fn fu
 		}
 	}()
 
-	s.cond.L.Lock()
-	defer s.cond.L.Unlock()
 	for ctx.Err() == nil && state > s.state {
 		s.cond.Wait()
 	}
 	if ctx.Err() != nil {
 		return s.state, ctx.Err()
 	}
-	if fn != nil {
-		err := fn(s.state)
-		if err != nil {
-			return s.state, err
-		}
-	}
 	return s.state, s.error
 }
 
