Skip to content

fix(kms): send the Finish response before shutting down - #875

Merged
kvinwang merged 1 commit into
masterfrom
codex/fix-kms-finish-response-order
Aug 4, 2026
Merged

fix(kms): send the Finish response before shutting down#875
kvinwang merged 1 commit into
masterfrom
codex/fix-kms-finish-response-order

Conversation

@kvinwang

@kvinwang kvinwang commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

Problem

Onboard.Finish called std::process::exit(0) from inside the RPC handler, so the
process died before Rocket could write the response. The onboarding client saw an EOF /
connection reset and could not tell whether finalization had succeeded.

Fix

Use Rocket's graceful shutdown, which is what the web UI's GET /finish route has always
done. Rocket finishes in-flight responses and then stops accepting connections, so the
client gets its response. run_onboard_service returns normally afterwards and main
continues into the regular KMS startup path — the same thing that already happens after
the web-UI flow, so no caller sees new behavior.

The shutdown handle only exists after ignition, so the launch is split into ignite() and
launch() with the handle stashed in OnboardState in between. It is written once, before
any request can be served, so a OnceLock is enough.

Scope change since the first revision

The earlier revision also rewrote JSON null response bodies to an empty body in
ra-rpc/src/rocket_helper.rs. That has been dropped from this PR: it was a
transport-layer workaround for a codec-layer problem, and it silently changed the wire
format of every unit-returning RPC in vmm and gateway, which does not belong in a KMS
fix. It is being fixed properly upstream in prpc instead
(Phala-Network/prpc#1), where codegen emits an empty body
for unit responses so the JSON codec matches the protobuf codec. dstack will pick it up
with a prpc-build bump in a separate PR.

This PR is no longer stacked: #874 is merged, so it targets master directly.

Verification

  • cargo clippy -p dstack-kms --all-features: clean.
  • cargo fmt: clean.

Copilot AI review requested due to automatic review settings July 31, 2026 02:52

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@kvinwang
kvinwang force-pushed the codex/fix-kms-finish-response-order branch from 3cf26d8 to a964b07 Compare August 3, 2026 14:53
@kvinwang
kvinwang force-pushed the codex/fix-kms-finish-response-order branch from a964b07 to 82efc4e Compare August 3, 2026 16:39
@kvinwang
kvinwang force-pushed the codex/fix-kms-finish-response-order branch from 82efc4e to 78a0dd2 Compare August 3, 2026 16:57
Base automatically changed from codex/fix-kms-private-key-permissions to master August 4, 2026 02:26
@kvinwang
kvinwang force-pushed the codex/fix-kms-finish-response-order branch from 78a0dd2 to d92c2d8 Compare August 4, 2026 02:26
`Onboard.Finish` called `std::process::exit(0)` from inside the RPC handler, so
the process died before Rocket could write the response. The onboarding client
saw an EOF and could not tell whether finalization had succeeded.

Use Rocket's graceful shutdown instead, which is what the web UI's `GET /finish`
route has always done: the server finishes in-flight responses and then stops
accepting connections. `run_onboard_service` then returns normally and `main`
continues into the regular KMS startup path, exactly as it does after the web-UI
flow.

The shutdown handle only exists after ignition, so the launch is split into
`ignite()` and `launch()` with the handle stashed in `OnboardState` in between.
It is written once, before any request can be served, so a `OnceLock` is enough.
@kvinwang
kvinwang force-pushed the codex/fix-kms-finish-response-order branch from d92c2d8 to 40b07b5 Compare August 4, 2026 02:47
@kvinwang kvinwang changed the title [STACKED on #874] fix(kms): return Finish response before shutdown fix(kms): send the Finish response before shutting down Aug 4, 2026
@kvinwang
kvinwang merged commit 93bbf75 into master Aug 4, 2026
15 checks passed
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