Remove err from wait callback

code-asher · code-asher · commit b30bc200a716 · 2023-08-11T10:42:18.000-08:00
We do not really do anything with it other than just return it, and if
we do need to do something with it we can just do it on the return from
`waitForStateOrContext`, we probably would not need to use it behind the
mutex.

Also we should check the state outside, not err, since there is
technically no guarantee an error is set on a state change (although in
practice for close/done there always is, maybe this should be enforced
in some way).
diff --git a/agent/reconnectingpty/buffered.go b/agent/reconnectingpty/buffered.go
@@ -179,19 +179,19 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
 	defer cancel()
 
 	// Once we are ready, attach the active connection while we hold the mutex.
-	_, err := rpty.state.waitForStateOrContext(ctx, StateReady, func(state State, err error) error {
+	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())
-		_, writeErr := conn.Write(prevBuf)
-		if writeErr != nil {
+		_, err := conn.Write(prevBuf)
+		if err != nil {
 			rpty.metrics.WithLabelValues("write").Add(1)
-			return xerrors.Errorf("write buffer to conn: %w", writeErr)
+			return xerrors.Errorf("write buffer to conn: %w", err)
 		}
 
 		if state != StateReady {
-			return err
+			return nil
 		}
 
 		go heartbeat(ctx, rpty.timer, rpty.timeout)
@@ -209,7 +209,7 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
 
 		return nil
 	})
-	if err != nil {
+	if state != StateReady {
 		return xerrors.Errorf("reconnecting pty ready wait: %w", err)
 	}
 
diff --git a/agent/reconnectingpty/reconnectingpty.go b/agent/reconnectingpty/reconnectingpty.go
@@ -168,9 +168,9 @@ 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, and fn's error will replace
-// waitForStateOrContext's error.
-func (s *ptyState) waitForStateOrContext(ctx context.Context, state State, fn func(state State, err error) error) (State, error) {
+// held unless the context ends.  If fn returns an error then fn's error will
+// replace waitForStateOrContext's error.
+func (s *ptyState) waitForStateOrContext(ctx context.Context, state State, fn func(state State) error) (State, error) {
 	nevermind := make(chan struct{})
 	defer close(nevermind)
 	go func() {
@@ -191,7 +191,10 @@ func (s *ptyState) waitForStateOrContext(ctx context.Context, state State, fn fu
 		return s.state, ctx.Err()
 	}
 	if fn != nil {
-		return s.state, fn(s.state, s.error)
+		err := fn(s.state)
+		if err != nil {
+			return s.state, err
+		}
 	}
 	return s.state, s.error
 }
