feat: connection-level keepalive + ConnectionClosingCallback - #4
Merged
Conversation
Hoists the keep-alive ticker loop out of session.handleRequests into a new connection-level goroutine in Server.HandleConn. It runs for the lifetime of the transport (using sshConn.SendRequest with the "keepalive@openssh.com" global request) rather than per-session, so a dead transport is detected even when no session is active — e.g., an idle ControlMaster after the outer ssh -O exit but before EOF propagates through the proxy chain. After ClientAliveCountMax consecutive intervals with no successful reply, sshConn is closed so HandleConn's channel loop unblocks. Also adds ConnectionClosingCallback, invoked synchronously the moment the inbound channels stream closes (before any defers, before sshConn.Wait()). Unlike ConnectionCompleteCallback this fires reliably even when the transport is permanently stuck. The callback receives the ssh.Context so per-connection state stashed via ctx.SetValue is reachable directly, removing the need for downstream callers to maintain a sync.Map keyed by *gossh.ServerConn. The per-session keepalive ticker driver in session.handleRequests is removed (the connection-level loop now drives the shared ctx.KeepAlive()); request-handler reset via KeepAliveRequestHandler still functions unchanged.
|
Warning Review limit reached
More reviews will be available in 42 minutes and 8 seconds. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughKeep-alive bookkeeping is moved to the transport level with open-channel tracking. A synchronous ConnectionClosingCallback is added and invoked when channel acceptance ends. Session-level keep-alive ticks are removed; connectionKeepAlive probes prefer channel requests, close stalled transports, and tests cover callback ordering and probe routing. ChangesKeep-Alive Migration
🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…ility Addresses review feedback on the connection-level keep-alive change: - Add `defer keepAlive.Close()` in connectionKeepAlive so the underlying time.Ticker is stopped deterministically when the goroutine exits (previously leaked until GC). - Race sshConn.SendRequest against time.After(interval) so a stalled transport can't pin the inner goroutine forever; the buffered replyCh ensures the inner goroutine never blocks on send after the parent moves on. - Restore log output that the per-session driver previously emitted: failed SendRequest, per-request timeout, the keep-alive timeout-close decision, and sshConn.Close errors. - Tighten ConnectionClosingCallback godoc: it only fires reliably on a stuck transport when ClientAliveInterval is set. Without keep-alive there is nothing to break the gossh channel range loop. - Replace the field-level comment that duplicated the type doc with a pointer, and drop the HandleConn inline comment that named another callback's internals (rot risk). Tests: - TestConnectionClosingCallback now uses a single ordered event channel to assert closing fires strictly before complete. - New TestConnectionKeepAliveClosesStalledConn dials with raw gossh.NewClientConn and drops incoming requests without replying, asserting ConnectionClosingCallback fires after roughly ClientAliveInterval * ClientAliveCountMax. This is the regression test for the original bug — without the connection-level keep-alive it would hang.
…ction OpenSSH's sshd sends keepalive@openssh.com as a per-channel request when at least one channel is open and as a global request otherwise (see serverloop.c client_alive_check). The previous connection-level implementation always sent a global request, which works with every real client (OpenSSH, gossh, libssh, paramiko all handle both forms) but diverged from sshd's wire behavior. This change tracks open channels per connection via a small openChannelSet stashed on the ssh.Context. Channel registration is transparent: each gossh.NewChannel is wrapped before being passed to the user-supplied ChannelHandler so Accept() registers the resulting Channel without any handler changes. Channel removal is lazy — when a keepalive SendRequest on a registered channel fails, that channel is pruned and the call falls through to the global request path. gossh exposes no Channel-close notification, so prune-on-failure is the cleanest approach. The keepalive loop now: - picks an open channel and calls ch.SendRequest if one exists, - prunes and falls back to sshConn.SendRequest on failure or empty set, - still bounds the send with the existing replyCh / time.After race, - still drives ctx.KeepAlive() for metrics and Reset on success, - still closes sshConn on TimeIsUp. Test added: TestConnectionKeepAliveUsesChannelRequestWhenSessionOpen opens a raw gossh channel, observes channel-typed keepalives arriving on the per-channel request stream (and zero on the global stream) for ~1s, then closes the channel and observes global-typed keepalives arriving for ~1s.
Addresses five divergences from OpenSSH sshd surfaced during review. Channel-close hook (closes #1 of review) Previously openChannelSet.remove fired only on a SendRequest failure, so a connection that opened and cleanly closed many channels accumulated dead entries forever. any() always returned slot 0 — the oldest, most-likely-dead channel — degrading the OpenSSH-mirror mode into "always probe a dead channel, then fall back to global" within one tick. trackingNewChannel.Accept now proxies the per-channel request stream through a buffered (16, matching gossh's chanSize) forwarder goroutine; when the upstream reqs channel closes, an onClose callback prunes the set. Remove the time.After race (closes #2 of review) connectionKeepAlive used a replyCh + time.After(interval) race around SendRequest. A reply arriving at interval+epsilon left keepAlive.Reset() un-called, ticking the false-disconnect counter for live-but-slow clients. Removed the outer timeout entirely. The inFlight semaphore already prevents overlapping probes; TimeIsUp enforces the deadline at the next tick; sshConn.Close on TimeIsUp unblocks any hung SendRequest. Reset on any inbound traffic (closes #3 of review) OpenSSH resets its client-alive counter on every received packet. Added SessionKeepAlive.NotePeerActivity — like Reset but without bumping the KeepAliveReplyReceived metric. Called from Server.handleRequests (global requests), session.handleRequests (per-session requests), and the per-channel request forwarder in trackingNewChannel. Resets the ticker so the next probe fires ClientAliveInterval after the most recent activity, which also closes #4 (probe suppression on recent traffic) for free. Unwrap helper + ChannelHandlers doc (closes #5 of review) Added Unwrap() gossh.NewChannel on trackingNewChannel so downstream handlers that need the underlying type can recover it. Documented the wrapping on Server.ChannelHandlers. Test added: TestConnectionKeepAlivePrunesClosedChannels opens three channels, closes them, then asserts subsequent keepalives over a 1s window are global-only (0 channel-typed, >=1 global). Without the close hook, the dead channels in the set would force channel-typed probes that fail and prune one-at-a-time over many intervals.
- Forwarder goroutine in trackingNewChannel.Accept now selects on ctx.Done(); on cancel it drains upstream reqs (replying false to WantReply) so the deferred onClose runs and gossh's sender doesn't leak. Previously a handler that abandoned wrapped would leak the forwarder forever once its 16-slot buffer filled. - NotePeerActivity no longer bumps lastReceived. It still resets the ticker (suppressing redundant probes when traffic flows), but the dead-peer deadline used by TimeIsUp is now cleared only by replies to our own probes. A chatty-but-wedged peer (send-side stuck, recv-side streaming) is no longer falsely considered alive — this matches OpenSSH's client_alive_check, which counts unanswered server-initiated probes, not arbitrary inbound traffic. - TestConnectionKeepAlivePrunesClosedChannels now actually exercises the close hook: interval bumped to 1s so we can poll the openChannelSet via context BEFORE the next keepalive tick. The previous timings allowed the old prune-on-failure path to pass the test without the hook. - Exported NewChannelUnwrapper named interface so Unwrap() is discoverable via go doc. ChannelHandlers godoc updated.
Round-2 review surfaced an internal contradiction: NotePeerActivity reset the ticker but did not bump lastReceived. A chatty inbound peer would keep deferring the probe (via ticker.Reset) so TimeIsUp never got consulted, defeating dead-peer detection. The doc comment also misrepresented OpenSSH semantics. OpenSSH's ssh_packet_read_poll_seqnr (packet.c) zeros keep_alive_timeouts on every successfully-received packet of any type, and serverloop.c defers the next probe on inbound traffic. NotePeerActivity now matches: bumps lastReceived AND resets the ticker on any inbound traffic. The Reset metric counter is still reserved for actual probe replies. Also: - Removed unreachable if t.ctx != nil guard in trackingNewChannel. Construction site (HandleConn) always passes a non-nil ctx, and the prior fallback comment was wrong (a nil select case deadlocks on full buffer, doesn't degrade). - Tightened TestConnectionKeepAlivePrunesClosedChannels: interval 1s -> 3s, poll deadline 800ms -> 500ms. Added invariant assertion that no keepalive tick fires inside the assertion window so the test can't silently pass via the prune-on-failure fallback.
- Removed defer keepAlive.Close() in connectionKeepAlive. The SessionKeepAlive is referenced by other goroutines (handleRequests, per-channel forwarders, in-flight probes) that may still call NotePeerActivity/Reset after connectionKeepAlive returns. The closed guards prevented an immediate panic but the pattern was fragile. The ticker is unreferenced once HandleConn returns and GC reclaims it. - ConnectionClosingCallback godoc now warns that channel handler goroutines may still be live when the callback fires; points to ConnectionCompleteCallback for all-work-done semantics. - Unexported the context key for openChannelSet (was ContextKeyOpenChannels, now contextKeyOpenChannels) — the value type is unexported so the exported key was misleading. - Trimmed NotePeerActivity godoc to drop the file-specific OpenSSH source citation while preserving the high-level semantic claim.
openChans is unconditionally created and stashed on ctx in HandleConn before connectionKeepAlive can run, so the comma-ok and nil-check were unreachable. Direct type assertion + dropping the guard simplifies the probe-send sequence.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
keepalive.go (1)
74-100:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftKeep generic peer activity separate from keepalive replies.
Line 84 updates
lastReceived, and Line 99 uses that same field to enforceClientAliveCountMax. Any inbound request therefore erases the missed-probe budget, so a peer can ignore every server keepalive and still stay connected by sending unrelated channel/global traffic.NotePeerActivity()should defer the next probe without resetting the state thatTimeIsUp()uses to count unanswered probes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@keepalive.go` around lines 74 - 100, NotePeerActivity currently updates lastReceived and therefore wipes the unanswered-probe count used by TimeIsUp; change the logic so general inbound traffic only defers the next probe (call ska.ticker.Reset) without touching the probe-response timestamp. Add a separate field on SessionKeepAlive (e.g., lastKeepAliveResponse or lastProbeAck) that TimeIsUp uses instead of lastReceived when computing clientAliveCountMax * clientAliveInterval, and ensure the code path that processes actual keepalive replies sets that new field. Update NotePeerActivity to stop modifying lastReceived (or rename lastReceived to reflect probe responses) and update TimeIsUp to reference the new probe-response field while leaving ticker.Reset behavior unchanged.server.go (1)
454-468:⚠️ Potential issue | 🟠 Major | ⚡ Quick winStop the keep-alive ticker during connection teardown.
This leaves
SessionKeepAliverunning afterHandleConnreturns, but the samectxis still handed to channel handlers andConnectionClosingCallback, so the ticker stays reachable and can keep firing indefinitely for stuck handlers. Please close it as part ofHandleConncleanup instead of relying on it becoming unreachable.Proposed fix
applyKeepAlive(ctx, srv.ClientAliveInterval, srv.ClientAliveCountMax) + defer ctx.KeepAlive().Close() openChans := &openChannelSet{} ctx.SetValue(contextKeyOpenChannels, openChans)Also applies to: 523-531
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server.go` around lines 454 - 468, The connection keep-alive goroutine's done channel (keepAliveDone) is currently closed with a defer here which may not run at the correct time, leaving SessionKeepAlive tickers running; change the teardown so the keep-alive ticker is explicitly stopped as part of HandleConn cleanup: create keepAliveDone and start the goroutine via srv.connectionKeepAlive(…, keepAliveDone) as now, but remove the defer close(keepAliveDone) and instead close keepAliveDone during the connection teardown path (e.g., at the end of HandleConn or inside the ConnectionClosingCallback invoked there) so the goroutine stops deterministically and SessionKeepAlive tickers are released. Ensure references to keepAliveDone live long enough for handlers that need it and that closing it is idempotent/safe for the goroutine reading it.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@keepalive.go`:
- Around line 74-100: NotePeerActivity currently updates lastReceived and
therefore wipes the unanswered-probe count used by TimeIsUp; change the logic so
general inbound traffic only defers the next probe (call ska.ticker.Reset)
without touching the probe-response timestamp. Add a separate field on
SessionKeepAlive (e.g., lastKeepAliveResponse or lastProbeAck) that TimeIsUp
uses instead of lastReceived when computing clientAliveCountMax *
clientAliveInterval, and ensure the code path that processes actual keepalive
replies sets that new field. Update NotePeerActivity to stop modifying
lastReceived (or rename lastReceived to reflect probe responses) and update
TimeIsUp to reference the new probe-response field while leaving ticker.Reset
behavior unchanged.
In `@server.go`:
- Around line 454-468: The connection keep-alive goroutine's done channel
(keepAliveDone) is currently closed with a defer here which may not run at the
correct time, leaving SessionKeepAlive tickers running; change the teardown so
the keep-alive ticker is explicitly stopped as part of HandleConn cleanup:
create keepAliveDone and start the goroutine via srv.connectionKeepAlive(…,
keepAliveDone) as now, but remove the defer close(keepAliveDone) and instead
close keepAliveDone during the connection teardown path (e.g., at the end of
HandleConn or inside the ConnectionClosingCallback invoked there) so the
goroutine stops deterministically and SessionKeepAlive tickers are released.
Ensure references to keepAliveDone live long enough for handlers that need it
and that closing it is idempotent/safe for the goroutine reading it.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 5c1b0980-48ab-4c39-9b9d-a65897dc8f91
📒 Files selected for processing (6)
context.gokeepalive.goserver.goserver_test.gosession.gossh.go
- New TestConnectionKeepAliveClosesStalledConnWithOpenSession exercises the channel-request path of dead-peer detection. Opens a session and has the client drain all request streams without replying; the server should still detect the stall via TimeIsUp and tear the transport down. Complements the existing no-session variant which exercises the global-request path. - Unexported NotePeerActivity -> notePeerActivity. It is an internal hook driven by this package's request loops; external callers have no reason to invoke it. Doing this before v1.2.0 is cut avoids a v2 break later.
skevetter
force-pushed
the
feat/connection-keepalive-closing-callback
branch
from
May 25, 2026 21:28
f80e595 to
1bd4e78
Compare
Per independent review: rely on GC was option A; explicit Stop in HandleConn's defer chain is option B. The ticker stays reachable after HandleConn returns because ctx is held by spawned handler goroutines and the per-channel request forwarder, so it kept firing until those exited. LIFO ensures close(keepAliveDone) runs first (signaling connectionKeepAlive to return), then ctx.KeepAlive().Close() Stops the ticker. All SessionKeepAlive methods are mutex-protected and check closed before touching the ticker, so a straggler caller is a safe no-op.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Two additions to fix a connection-cleanup bug in downstream consumers (devsy
pkg/ssh/server) where stdin EOF doesn't propagate promptly through proxy chains (outer ssh → ProxyCommand → docker exec → in-container helper), leavingHandleConnstuck infor ch := range chansand alldefers — includingConnectionCompleteCallback— never running.1. Connection-level keep-alive (OpenSSH-faithful)
The keep-alive previously lived in
session.handleRequestsand only ran while a session was active. Between sessions (idleControlMaster) nothing pinged the peer, so a dead transport went undetected.It's now hoisted to a goroutine started by
Server.HandleConnthat runs for the lifetime of the transport. The probe-selection mirrors OpenSSHsshd'sclient_alive_check():keepalive@openssh.comas a channel request on the oldest open channel.After
ClientAliveCountMaxconsecutive intervals with no inbound traffic from the peer,sshConn.Close()is called andHandleConnunblocks.Internals:
openChannelSet(per-connection, stashed onssh.Context) tracks accepted channels.trackingNewChannelwraps everygossh.NewChannelsoAccept()transparently registers; the per-channelreqsstream is also wrapped so close drives unregistration. No changes needed to existing user-suppliedChannelHandlers.notePeerActivity(unexported) bumpslastReceivedand resets the probe ticker on every inbound packet, matching OpenSSH'spacket.csemantics wherekeep_alive_timeouts = 0on any successfully-received packet.NewChannelUnwrapperexported interface gives any handler that needs the underlyinggossh.NewChannel(type-assert against a custom impl, etc.) a stable way to recover it.The per-session keep-alive ticker driver is removed — the connection-level loop supersedes it. Client-initiated
keepalive@openssh.com(viaKeepAliveRequestHandler) is unchanged.2.
ConnectionClosingCallbackInvoked synchronously the moment
HandleConnobserves the inbound channels stream close — before anydefers fire and beforesshConn.Wait(). UnlikeConnectionCompleteCallback, which is wrapped in a defer aroundWait()and so may never fire on a stuck transport, this callback gives downstream callers a fast-path hook for cleanup.The
ssh.Contextargument lets callers reach per-connection state stashed viactx.SetValuedirectly — nosync.Map[*gossh.ServerConn]indirection needed.Caveat documented on the type: channel handler goroutines spawned by
HandleConnmay still be live when the callback fires. For "all work done" semantics, useConnectionCompleteCallback.Behavior changes for existing consumers
client_alive_check, runs alwayskeepalive@openssh.com(server-initiated) wire formConnectionClosingCallback,NewChannelUnwrapperTest additions
TestConnectionClosingCallbackTestConnectionKeepAliveClosesStalledConnTestConnectionKeepAliveClosesStalledConnWithOpenSessionTestConnectionKeepAliveUsesChannelRequestWhenSessionOpenTestConnectionKeepAlivePrunesClosedChannelsVerification
go build ./...cleanCGO_ENABLED=1 go test -race -timeout 180s ./...— 34 tests pass, no flakesgolangci-lint runcleanRelease
feat:commit at the head triggers release-please to cutv1.2.0on merge tomain. devsy can then bumpgo.modtogithub.com/devsy-org/ssh@v1.2.0.Summary by CodeRabbit
New Features
Improvements
Tests