Avoid wait callback when context expires

code-asher · code-asher · commit 34c5c1a6a71d · 2023-08-11T10:28:30.000-08:00
diff --git a/agent/reconnectingpty/buffered.go b/agent/reconnectingpty/buffered.go
@@ -191,7 +191,7 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
 		}
 
 		if state != StateReady {
-			return xerrors.Errorf("reconnecting pty ready wait: %w", err)
+			return err
 		}
 
 		go heartbeat(ctx, rpty.timer, rpty.timeout)
@@ -210,7 +210,7 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
 		return nil
 	})
 	if err != nil {
-		return err
+		return xerrors.Errorf("reconnecting pty ready wait: %w", err)
 	}
 
 	defer func() {
diff --git a/agent/reconnectingpty/reconnectingpty.go b/agent/reconnectingpty/reconnectingpty.go
@@ -168,7 +168,8 @@ 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 and fn's error will replace waitForStateOrContext's error.
+// 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) {
 	nevermind := make(chan struct{})
 	defer close(nevermind)
@@ -186,14 +187,13 @@ func (s *ptyState) waitForStateOrContext(ctx context.Context, state State, fn fu
 	for ctx.Err() == nil && state > s.state {
 		s.cond.Wait()
 	}
-	err := s.error
 	if ctx.Err() != nil {
-		err = ctx.Err()
+		return s.state, ctx.Err()
 	}
 	if fn != nil {
-		return s.state, fn(s.state, err)
+		return s.state, fn(s.state, s.error)
 	}
-	return s.state, err
+	return s.state, s.error
 }
 
 // readConnLoop reads messages from conn and writes to ptty as needed.  Blocks

Original file line number	Diff line number	Diff line change
`@@ -191,7 +191,7 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,`
`191`	`191`	`}`
`192`	`192`
`193`	`193`	`if state != StateReady {`
`194`		`- return xerrors.Errorf("reconnecting pty ready wait: %w", err)`
	`194`	`+ return err`
`195`	`195`	`}`
`196`	`196`
`197`	`197`	`go heartbeat(ctx, rpty.timer, rpty.timeout)`
`@@ -210,7 +210,7 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,`
`210`	`210`	`return nil`
`211`	`211`	`})`
`212`	`212`	`if err != nil {`
`213`		`- return err`
	`213`	`+ return xerrors.Errorf("reconnecting pty ready wait: %w", err)`
`214`	`214`	`}`
`215`	`215`
`216`	`216`	`defer func() {`