Skip to content

experimental/ssh: survive a dropped tunnel connection instead of ending the session - #6558

Open
anton-107 wants to merge 1 commit into
mainfrom
deco-28433-ws-stability-investigate-connections
Open

experimental/ssh: survive a dropped tunnel connection instead of ending the session#6558
anton-107 wants to merge 1 commit into
mainfrom
deco-28433-ws-stability-investigate-connections

Conversation

@anton-107

Copy link
Copy Markdown
Contributor

A single TCP reset anywhere between the client and the workspace ended an
ssh connect session, and both ends destroyed their state within milliseconds, so
nothing survived to reconnect to. It happens several times a session for some
users, badly enough that they avoid the terminal entirely.

The proxy already knew how to swap a websocket underneath a running sshd — that is
what the periodic auth handover does — but only for a connection that still works:
the close frame is written after the last data frame, and TCP ordering is the whole
safety argument. An abrupt reset provides no such barrier, so this adds the byte
accounting that replaces it.

Three changes, smallest first.

A failed handover dial no longer ends the session

The handover dials a replacement websocket every --handover-timeout (30m by
default). Any failure of that dial ended the session, even though nothing had been
swapped yet and the connection being replaced was still carrying traffic — so one
transient blip dropped a healthy tunnel on a 30-minute clock. The pre-swap failure
is now tagged so the client stays on its current connection and waits for the next
tick. Deferring the refresh is safe: the driver proxy authenticates a websocket at
upgrade time, so a live connection is not re-checked.

Retrying the dial in place would not be safe. A dial can fail after the server has
already accepted it and begun acceptHandover, and a second dial would then race
the first for the same connection.

A dropped session is no longer reported as a successful one

isSuccess is set before the proxy loop starts and category() short-circuited on
it, so every mid-session drop reached telemetry as a clean, successful session: the
drop rate was unmeasurable and a fix for it unverifiable. The sites that end a
session now carry sentinels (ErrConnectFailed, ErrWebsocketDropped,
ErrHandoverFailed) mapped onto three new categories.

isSuccess keeps its meaning — the tunnel was established — so is_success
separates a failed connection attempt from a session that connected and was later
cut short, and error_category gives the cause of either. category() no longer
discards a category just because the tunnel came up, but still refuses to invent
one: an established session that ended unattributed stays TYPE_UNSPECIFIED,
because the ssh client and the user's own remote command both exit non-zero there
and neither is a tunnel failure. Only a failed connection attempt falls back to
UNKNOWN, as before.

The three new SshTunnelErrorCategory values (WEBSOCKET_CONNECT_FAILED,
WEBSOCKET_DROPPED, HANDOVER_FAILED) get a companion PR in universe adding them to
proto/logs/frontend/databricks_cli/enum.proto. Ingestion ignores unknown enum
values, so the two can land in either order.

A dropped connection is reattached rather than mourned

Both ends now count the payload they have handed to the websocket and the payload
they have written to their destination, and keep a bounded tail of the former for
replay. The buffer, not the socket, is the source of truth: bytes are appended
before they are written, so a write that fails on a dying connection loses nothing.
On a drop the client redials with the offset it has delivered, the server answers
with its own, and each side replays exactly the difference. SSH verifies a MAC over
the byte stream, so this has to be exact — one lost or duplicated byte would
disconnect the session rather than repair it. Delivered counts are acknowledged
every 64 KB so neither replay buffer grows.

Server side, a dropped connection parks the session rather than reaping it: sshd
stays alive, the client slot stays held (which is also what keeps the shutdown timer
cancelled) and a grace timer tears it down exactly as before if nobody comes back.

The three timings are each derived from a constraint rather than picked: the client
gives up after 60s because ssh itself does at roughly 90s
(ServerAliveInterval 30 x OpenSSH's default ServerAliveCountMax 3); the server's
grace is 90s so it never reaps a client still trying; the 1 MB replay cap is far
above the unacknowledged window that acks every 64 KB can produce, so reaching it
means the peer stopped acknowledging and the connection is already beyond saving.

Three ordering hazards, each handled

  • The write mutex cannot be held across the wait for the peer, because that is the
    lock the other side needs to complete the reattach — it deadlocks the server.
    sendGate throttles the sending loop instead; correctness rests on the write
    mutex plus appending before writing.
  • A reattach must not report a delivered count that can still move. The client often
    notices a drop first, and the server may still be draining data buffered on the
    dying connection, so acceptReattach retires that connection and waits for the
    receiving loop to stop delivering before it answers. Otherwise the client replays
    from a stale offset and the server writes those bytes to sshd twice.
  • A handover tick landing during a client reattach would wait out its own timeout
    for a receiving loop that is busy reattaching, and end the session. The client
    holds the write mutex across its whole reattach so the two serialize.

Both ends must agree, and version skew is real

The generated ssh config pins --metadata into the ProxyCommand, which skips the
version-scoped server lookup, so an upgraded client can reach a server an older CLI
started. The client therefore probes a new /capabilities endpoint first; an older
server has no such route, 404s, and the session behaves exactly as it does today. A
resume-capable server also refuses a reattach for a session it no longer holds (410)
rather than starting a fresh sshd, since replaying into that would fail the SSH
stream with a corrupted MAC — a worse error than the drop.

The test server reports resume: false, because it drives sshd directly over the
websocket rather than running the CLI's own proxy server and has none of the session
bookkeeping a reattach needs.

Tests

Byte-exact delivery across single and repeated resets, driven through a TCP relay
that sends real RSTs (httptest's CloseClientConnections is no use: it does not
touch the hijacked connections a websocket upgrade leaves behind). The reattach
exchange is asserted frame by frame with a deliberately non-empty replay, so the
handshake ordering is pinned rather than incidentally satisfied. Plus the replay
arithmetic in isolation, the non-resumable path unchanged, the grace period
releasing an abandoned session, and 410/400 on malformed reattach requests.

TestADropIsNeverReportedAsACleanExit guards the errgroup race fixed in 75fcc00
from the drop side rather than the keepalive side. That fix keeps a cancellation
from masking the error that caused it; reverting it fails this test within a few of
its 25 attempts. It was found independently here while chasing an intermittent
failure of the new drop test, and landed upstream first, so only the test remains.

category() needed a real merge with the extension-category split in 3adaa48. Its
precedence was cyclic across the two changes: an interruption had to outrank the
category recorded at the failure site, isSuccess had to outrank an interruption,
and a session-end category has to outrank isSuccess. Broken by making the split
explicit instead — once the tunnel is up, the only reportable category is a
session-end one, so an established session no longer consults the connection-attempt
rules at all.

Still to validate against real compute: whether the driver proxy passes the reattach
query parameter through a websocket upgrade unchanged. ?id= already shows query
parameters survive, so this is a confirmation rather than a risk, but it is
load-bearing for the handshake.

This pull request and its description were written by Isaac.

…ng the session

A single TCP reset anywhere between the client and the workspace ended an
`ssh connect` session, and both ends destroyed their state within milliseconds, so
nothing survived to reconnect to. It happens several times a session for some
users, badly enough that they avoid the terminal entirely.

The proxy already knew how to swap a websocket underneath a running sshd - that is
what the periodic auth handover does - but only for a connection that still works:
the close frame is written after the last data frame, and TCP ordering is the whole
safety argument. An abrupt reset provides no such barrier, so this adds the byte
accounting that replaces it.

Three changes, smallest first.

## A failed handover dial no longer ends the session

The handover dials a replacement websocket every `--handover-timeout` (30m by
default). Any failure of that dial ended the session, even though nothing had been
swapped yet and the connection being replaced was still carrying traffic - so one
transient blip dropped a healthy tunnel on a 30-minute clock. The pre-swap failure
is now tagged so the client stays on its current connection and waits for the next
tick. Deferring the refresh is safe: the driver proxy authenticates a websocket at
upgrade time, so a live connection is not re-checked.

Retrying the dial in place would not be safe. A dial can fail after the server has
already accepted it and begun `acceptHandover`, and a second dial would then race
the first for the same connection.

## A dropped session is no longer reported as a successful one

`isSuccess` is set before the proxy loop starts and `category()` short-circuited on
it, so every mid-session drop reached telemetry as a clean, successful session: the
drop rate was unmeasurable and a fix for it unverifiable. The sites that end a
session now carry sentinels (`ErrConnectFailed`, `ErrWebsocketDropped`,
`ErrHandoverFailed`) mapped onto three new categories.

`isSuccess` keeps its meaning - the tunnel was established - so `is_success`
separates a failed connection attempt from a session that connected and was later
cut short, and `error_category` gives the cause of either. `category()` no longer
discards a category just because the tunnel came up, but still refuses to invent
one: an established session that ended unattributed stays `TYPE_UNSPECIFIED`,
because the ssh client and the user's own remote command both exit non-zero there
and neither is a tunnel failure. Only a failed connection attempt falls back to
`UNKNOWN`, as before.

## A dropped connection is reattached rather than mourned

Both ends now count the payload they have handed to the websocket and the payload
they have written to their destination, and keep a bounded tail of the former for
replay. The buffer, not the socket, is the source of truth: bytes are appended
before they are written, so a write that fails on a dying connection loses nothing.
On a drop the client redials with the offset it has delivered, the server answers
with its own, and each side replays exactly the difference. SSH verifies a MAC over
the byte stream, so this has to be exact - one lost or duplicated byte would
disconnect the session rather than repair it. Delivered counts are acknowledged
every 64 KB so neither replay buffer grows.

Server side, a dropped connection parks the session rather than reaping it: sshd
stays alive, the client slot stays held (which is also what keeps the shutdown timer
cancelled) and a grace timer tears it down exactly as before if nobody comes back.

The three timings are each derived from a constraint rather than picked: the client
gives up after 60s because ssh itself does at roughly 90s
(`ServerAliveInterval` 30 x OpenSSH's default `ServerAliveCountMax` 3); the server's
grace is 90s so it never reaps a client still trying; the 1 MB replay cap is far
above the unacknowledged window that acks every 64 KB can produce, so reaching it
means the peer stopped acknowledging and the connection is already beyond saving.

### Three ordering hazards, each handled

- The write mutex cannot be held across the wait for the peer, because that is the
  lock the other side needs to complete the reattach - it deadlocks the server.
  `sendGate` throttles the sending loop instead; correctness rests on the write
  mutex plus appending before writing.
- A reattach must not report a delivered count that can still move. The client often
  notices a drop first, and the server may still be draining data buffered on the
  dying connection, so `acceptReattach` retires that connection and waits for the
  receiving loop to stop delivering before it answers. Otherwise the client replays
  from a stale offset and the server writes those bytes to sshd twice.
- A handover tick landing during a client reattach would wait out its own timeout
  for a receiving loop that is busy reattaching, and end the session. The client
  holds the write mutex across its whole reattach so the two serialize.

### Both ends must agree, and version skew is real

The generated ssh config pins `--metadata` into the `ProxyCommand`, which skips the
version-scoped server lookup, so an upgraded client can reach a server an older CLI
started. The client therefore probes a new `/capabilities` endpoint first; an older
server has no such route, 404s, and the session behaves exactly as it does today. A
resume-capable server also refuses a reattach for a session it no longer holds (410)
rather than starting a fresh sshd, since replaying into that would fail the SSH
stream with a corrupted MAC - a worse error than the drop.

The test server reports `resume: false`, because it drives sshd directly over the
websocket rather than running the CLI's own proxy server and has none of the session
bookkeeping a reattach needs.

## Tests

Byte-exact delivery across single and repeated resets, driven through a TCP relay
that sends real RSTs (`httptest`'s `CloseClientConnections` is no use: it does not
touch the hijacked connections a websocket upgrade leaves behind). The reattach
exchange is asserted frame by frame with a deliberately non-empty replay, so the
handshake ordering is pinned rather than incidentally satisfied. Plus the replay
arithmetic in isolation, the non-resumable path unchanged, the grace period
releasing an abandoned session, and 410/400 on malformed reattach requests.

`TestADropIsNeverReportedAsACleanExit` guards the errgroup race fixed in 75fcc00
from the drop side rather than the keepalive side. That fix keeps a cancellation
from masking the error that caused it; reverting it fails this test within a few of
its 25 attempts. It was found independently here while chasing an intermittent
failure of the new drop test, and landed upstream first, so only the test remains.

`category()` needed a real merge with the extension-category split in 3adaa48. Its
precedence was cyclic across the two changes: an interruption had to outrank the
category recorded at the failure site, `isSuccess` had to outrank an interruption,
and a session-end category has to outrank `isSuccess`. Broken by making the split
explicit instead - once the tunnel is up, the only reportable category is a
session-end one, so an established session no longer consults the connection-attempt
rules at all.

Still to validate against real compute: whether the driver proxy passes the reattach
query parameter through a websocket upgrade unchanged. `?id=` already shows query
parameters survive, so this is a confirmation rather than a risk, but it is
load-bearing for the handshake.

Co-authored-by: Isaac <no-reply@databricks.com>
@github-actions

github-actions Bot commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

Approval status: pending

/libs/telemetry/ - needs approval

Files: libs/telemetry/protos/ssh_tunnel.go
Suggested: @renaudhartert-db
Also eligible: @simonfaltum, @Divyansh-db, @hectorcast-db, @parthban-db, @tanmay-db, @tejaskochar-db, @mihaimitrea-db, @chrisst, @rauchy

General files (require maintainer)

18 files changed
Based on git history:

  • @janniklasrose -- recent work in .nextchanges/cli/, libs/testserver/, experimental/ssh/internal/proxy/

Any maintainer (@andrewnester, @denik, @pietern, @shreyas-goenka, @simonfaltum, @renaudhartert-db, @janniklasrose, @lennartkats-db, @rugpanov, @rclarey) can approve all areas.
See OWNERS for ownership rules.

@eng-dev-ecosystem-bot

Copy link
Copy Markdown
Collaborator

Integration test report

Commit: 7cb2120

Run: 34153220207

Env 🔄​flaky 💚​RECOVERED ✅​pass 🙈​skip Time
💚​ aws linux 1 275 15 12:37
🔄​ aws windows 1 1 276 13 7:25
💚​ azure linux 1 274 15 10:03
💚​ azure windows 1 276 13 8:29
💚​ gcp linux 1 275 15 10:00
💚​ gcp windows 1 277 13 8:18
Test Name aws linux aws windows azure linux azure windows gcp linux gcp windows
💚​ TestAccept 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R
🔄​ TestSyncFullFileSync ✅​p 🔄​f ✅​p ✅​p ✅​p ✅​p
Top 12 slowest tests (at least 2 minutes):
duration env testname
4:29 gcp windows TestAccept
4:03 aws linux TestAccept
4:02 aws windows TestAccept
4:00 azure windows TestAccept
3:55 aws linux TestFilerWorkspaceFilesExtensionsReadDir
3:52 gcp linux TestAccept
3:50 azure linux TestAccept
3:45 azure windows TestImportDirWithOverwriteFlag
3:01 aws linux TestFilerWorkspaceNotebook/sqlNb.sql
2:18 gcp linux TestFilerWorkspaceFilesExtensionsReadDir
2:14 azure linux TestFilerWorkspaceFilesExtensionsDelete
2:01 aws linux TestFilerRecursiveDelete/workspace_files

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants