Skip to content

Commit aeb24d4

Browse files
committed
Do not hold mutex while waiting for state
Since we might add to the map after it closes, tweak the delete to also run when the pty closes. I think this also fixes that issue where if you canceled the context you pass into Attach it would not do anything since readConnLoop does not use the context for anything but logging.
1 parent a6922f8 commit aeb24d4

File tree

1 file changed

+20
-27
lines changed

1 file changed

+20
-27
lines changed

agent/reconnectingpty/buffered.go

Lines changed: 20 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ func newBuffered(ctx context.Context, cmd *pty.Cmd, options *Options, logger slo
125125
}
126126

127127
// lifecycle manages the lifecycle of the reconnecting pty. If the context ends
128-
// the reconnecting pty and every connection will be closed.
128+
// or the reconnecting pty closes the pty will be shut down.
129129
func (rpty *bufferedReconnectingPTY) lifecycle(ctx context.Context, logger slog.Logger) {
130130
rpty.timer = time.AfterFunc(attachTimeout, func() {
131131
rpty.Close("reconnecting pty timeout")
@@ -142,22 +142,8 @@ func (rpty *bufferedReconnectingPTY) lifecycle(ctx context.Context, logger slog.
142142
}
143143
rpty.timer.Stop()
144144

145-
rpty.mutex.Lock()
146-
defer rpty.mutex.Unlock()
147-
148-
// Log these closes only for debugging since the connections or processes
149-
// might have already closed on their own.
150-
for _, conn := range rpty.activeConns {
151-
err := conn.Close()
152-
if err != nil {
153-
logger.Debug(ctx, "closed conn with error", slog.Error(err))
154-
}
155-
}
156-
// Connections get removed once the pty closes but it is possible there is
157-
// still some data that needs to be written so clear the map now to avoid
158-
// writing to closed connections.
159-
rpty.activeConns = map[string]net.Conn{}
160-
145+
// Log close/kill only for debugging since the process might have already
146+
// closed on its own.
161147
err := rpty.ptty.Close()
162148
if err != nil {
163149
logger.Debug(ctx, "closed ptty with error", slog.Error(err))
@@ -179,15 +165,29 @@ func (rpty *bufferedReconnectingPTY) Attach(ctx context.Context, connID string,
179165
ctx, cancel := context.WithCancel(ctx)
180166
defer cancel()
181167

182-
err := rpty.doAttach(ctx, connID, conn, height, width, logger)
168+
state, err := rpty.state.waitForStateOrContext(ctx, StateReady)
169+
if state != StateReady {
170+
return xerrors.Errorf("reconnecting pty ready wait: %w", err)
171+
}
172+
173+
go heartbeat(ctx, rpty.timer, rpty.timeout)
174+
175+
err = rpty.doAttach(ctx, connID, conn, height, width, logger)
183176
if err != nil {
184177
return err
185178
}
186179

187-
defer func() {
180+
go func() {
181+
_, _ = rpty.state.waitForStateOrContext(ctx, StateClosing)
188182
rpty.mutex.Lock()
189183
defer rpty.mutex.Unlock()
190184
delete(rpty.activeConns, connID)
185+
// Log closes only for debugging since the connection might have already
186+
// closed on its own.
187+
err := conn.Close()
188+
if err != nil {
189+
logger.Debug(ctx, "closed conn with error", slog.Error(err))
190+
}
191191
}()
192192

193193
// Pipe conn -> pty and block. pty -> conn is handled in newBuffered().
@@ -203,15 +203,8 @@ func (rpty *bufferedReconnectingPTY) doAttach(ctx context.Context, connID string
203203
rpty.mutex.Lock()
204204
defer rpty.mutex.Unlock()
205205

206-
state, err := rpty.state.waitForStateOrContext(ctx, StateReady)
207-
if state != StateReady {
208-
return xerrors.Errorf("reconnecting pty ready wait: %w", err)
209-
}
210-
211-
go heartbeat(ctx, rpty.timer, rpty.timeout)
212-
213206
// Resize the PTY to initial height + width.
214-
err = rpty.ptty.Resize(height, width)
207+
err := rpty.ptty.Resize(height, width)
215208
if err != nil {
216209
// We can continue after this, it's not fatal!
217210
logger.Warn(ctx, "reconnecting PTY initial resize failed, but will continue", slog.Error(err))

0 commit comments

Comments
 (0)