fix(kms): send the Finish response before shutting down - #875
Merged
Conversation
kvinwang
force-pushed
the
codex/fix-kms-finish-response-order
branch
from
August 3, 2026 14:53
3cf26d8 to
a964b07
Compare
kvinwang
force-pushed
the
codex/fix-kms-finish-response-order
branch
from
August 3, 2026 16:39
a964b07 to
82efc4e
Compare
kvinwang
force-pushed
the
codex/fix-kms-finish-response-order
branch
from
August 3, 2026 16:57
82efc4e to
78a0dd2
Compare
Base automatically changed from
codex/fix-kms-private-key-permissions
to
master
August 4, 2026 02:26
kvinwang
force-pushed
the
codex/fix-kms-finish-response-order
branch
from
August 4, 2026 02:26
78a0dd2 to
d92c2d8
Compare
`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
force-pushed
the
codex/fix-kms-finish-response-order
branch
from
August 4, 2026 02:47
d92c2d8 to
40b07b5
Compare
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
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.
Problem
Onboard.Finishcalledstd::process::exit(0)from inside the RPC handler, so theprocess 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 /finishroute has alwaysdone. Rocket finishes in-flight responses and then stops accepting connections, so the
client gets its response.
run_onboard_servicereturns normally afterwards andmaincontinues 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()andlaunch()with the handle stashed inOnboardStatein between. It is written once, beforeany request can be served, so a
OnceLockis enough.Scope change since the first revision
The earlier revision also rewrote JSON
nullresponse bodies to an empty body inra-rpc/src/rocket_helper.rs. That has been dropped from this PR: it was atransport-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-buildbump in a separate PR.This PR is no longer stacked: #874 is merged, so it targets
masterdirectly.Verification
cargo clippy -p dstack-kms --all-features: clean.cargo fmt: clean.