feat(gateway): upgrade to wavekv 2.0 delta-state sync with v1 dual-stack - #1031
Open
kvinwang wants to merge 9 commits into
Open
feat(gateway): upgrade to wavekv 2.0 delta-state sync with v1 dual-stack#1031kvinwang wants to merge 9 commits into
kvinwang wants to merge 9 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR upgrades dstack-gateway’s replicated KV sync layer to wavekv 2.0 (delta-state replication) while remaining compatible with wavekv 1.x peers during rolling upgrades, and adds schema-based admission control plus new sync observability fields exposed via the admin RPC.
Changes:
- Add dual-stack HTTP sync endpoints (
/wavekv/syncv1 +/wavekv/sync2v2) and an opportunistic push route (/wavekv/push) to reduce propagation latency. - Enforce per-store key-shape admission via a new schema policy integrated into wavekv node config.
- Extend admin/RPC status reporting with per-store digests and per-peer negotiated protocol / mismatch telemetry, and update wavekv dependency to the v2 branch.
Reviewed changes
Copilot reviewed 9 out of 10 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| dstack/gateway/src/web_routes/wavekv_sync.rs | Adds v2 sync and push HTTP endpoints; refactors gzip handling and introduces envelope decoding. |
| dstack/gateway/src/web_routes.rs | Mounts the new wavekv v2 sync + push routes alongside v1. |
| dstack/gateway/src/kv/sync_service.rs | Extends the sync network interface to use wavekv v2 envelopes and probing for v1/v2 negotiation. |
| dstack/gateway/src/kv/schema.rs | Introduces per-store key admission policy (schema) with tests. |
| dstack/gateway/src/kv/mod.rs | Wires admission policy into wavekv node configs; adds gateway-level wire-compat tests for v1/v2 sync. |
| dstack/gateway/src/kv/https_client.rs | Adds raw-bytes probe POST helper for v2 negotiation and opportunistic push transport. |
| dstack/gateway/src/admin_service.rs | Plumbs new wavekv v2 telemetry (digest, merged/rejected, per-peer protocol/mismatches) into admin RPC responses. |
| dstack/gateway/rpc/proto/gateway_rpc.proto | Extends sync status protos with digest + v2 peer telemetry; deprecates buffered_logs. |
| dstack/Cargo.toml | Switches wavekv dependency to the v2 git branch (with TODO to repoint to crates.io 2.0). |
| dstack/Cargo.lock | Locks wavekv to the v2 git revision and updates transitive deps accordingly. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+240
to
+264
| cert: Option<Certificate<'_>>, | ||
| store: &str, | ||
| data: Data<'_>, | ||
| ) -> Result<Status, Status> { | ||
| verify_gateway_peer(state, cert)?; | ||
|
|
||
| let Some(ref wavekv_sync) = state.wavekv_sync else { | ||
| return Err(Status::ServiceUnavailable); | ||
| }; | ||
|
|
||
| let env = read_envelope(data).await?; | ||
| if env.sender_id == 0 { | ||
| warn!("rejected push from invalid node_id 0"); | ||
| return Err(Status::BadRequest); | ||
| } | ||
|
|
||
| let Some(result) = wavekv_sync.handle_push(store, env) else { | ||
| return Err(Status::NotFound); | ||
| }; | ||
| result.map_err(|e| { | ||
| tracing::error!("{store} push failed: {e:#}"); | ||
| Status::InternalServerError | ||
| })?; | ||
| Ok(Status::Ok) | ||
| } |
Comment on lines
371
to
375
| /// Encode a KV value as MessagePack. | ||
| /// | ||
| /// Structs are encoded as maps keyed by field name rather than as positional | ||
| /// arrays. Field-name keys let a reader skip fields it does not know and fill | ||
| /// in `#[serde(default)]` fields it does not receive, so the value types below |
Comment on lines
+87
to
+92
| let bytes = data | ||
| .open(16.mebibytes()) | ||
| .into_bytes() | ||
| .await | ||
| .map_err(|_| Status::BadRequest)?; | ||
| let decompressed = gunzip(&bytes)?; |
Pick up the wavekv fix for the opportunistic push envelope, which was built without a `sender_uuid` and so failed `check_uuid` on every push — this gateway implements `query_uuid`, so the push channel never worked here. Writes still converged over the periodic round, but each one waited a full sync interval instead of the coalesce window and the receiver logged an error per push blaming node-id reuse. That fix also widens `link_status` to report every known peer rather than only those in the link cache. A peer whose rounds all fail was previously absent from `WaveKvStatus` entirely: a 5xx deliberately does not demote a peer to "v1", so nothing about it moved. Report the new `consecutive_failures` streak so that stall is visible. Document the one direction in which the store schema is not forward compatible: values may gain fields freely, but a new *key* is rejected by nodes that predate it, and a rejection parks ack adoption for the whole round (rule R1). The pair then re-exchanges the same batch indefinitely with no error. New keys therefore ship in two releases — widen the schema everywhere first, write the key second. Also silence a `manual_repeat_n` lint in the pp tests, unrelated but newly raised by the toolchain and enough to fail `clippy -D warnings`.
The HTTP layer was the one part of the sync path with no coverage. It was skipped on the grounds that constructing a `WaveKvSyncService` needs real TLS material; that was wrong. `rcgen` is already a dependency and already used by the cert_store tests, and `verify_gateway_peer` short-circuits under `insecure_skip_attestation`, so a self-signed CA plus a leaf written to a TempDir is enough to build a serving gateway. What this pins that nothing else did: - 503, not 404, when sync is disabled. 404 is the negotiation signal, so a sync-disabled node answering 404 would be cached as "v1" by every peer for a whole reprobe window — and sync is off, so nothing would correct it. - 404 for an unknown store, which is the same signal used deliberately. - An unstamped push is refused at the route and writes nothing. This is the server-side view of the envelope-identity bug; the sender-side view lives in the wavekv push test. - A well-formed push reaches the store, a v2 round trip returns a decodable envelope, and node id 0 is refused. Also stop reporting a 404 on the push route as a delivered push. `post_bytes_probe` maps 404/405 to `Ok(None)` so the v2 probe can read it as "not upgraded yet", but `push_to` discarded the `Option`. A mistyped push URL was therefore indistinguishable from success — the same shape of silent failure that let the unstamped-envelope bug survive, since pushes are best-effort and only debug-logged.
Takes the wavekv fix that verifies the responder's uuid on a v2 response. The field was already on the wire and populated by the responder; only the initiator never read it, so node-id-reuse detection ran in one direction.
The responder-side identity check shipped in the previous bump wedged any peer that regenerated its uuid — an ordinary CVM rebuild, since the uuid is derived from the data directory while the node id comes from config.
The sync wire is gzipped and the 16 MiB cap on a request body caps the *compressed* size, which bounds nothing on its own — gzip expands by three orders of magnitude on attacker-chosen input, so that cap admits a payload that expands into the gigabytes. Every gateway in a cluster shares one app_id, so mTLS proves only that the sender is some gateway of this deployment; it is the same trust level the key schema already treats as insufficient. All four decompression points are now bounded through one helper: both server routes and both client response paths. The client also read peer responses with `Body::collect`, which has no limit at all, so the memory was already spent before any decoding bound could apply; response bodies now go through `Limited` with the same 16 MiB the routes accept on a request. The decompressed ceiling is 128 MiB, far above any legitimate payload: a v2 delta is capped by `max_delta_bytes` at 4 MiB, and the v1 shim answers with the whole live state, which is bounded by the gateway's own key set rather than by anything a peer controls. Tested at the limit as well as past it — a fixture landing exactly on the ceiling must still decode, or the bound could tighten by a byte with only the bomb test still passing.
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.
Upgrades dstack-gateway onto wavekv 2.0 (delta-state replication, Phala-Network/wavekv#3) as a dual-stack node: it serves the native v2 protocol and keeps serving v1 peers, so a gateway cluster can be upgraded one CVM at a time.
Problem
The gateway's replicated state is an LWW CRDT, but wavekv 1.x replicates it with per-origin operation logs. Op-based replication needs exactly-once ordered delivery, and the machinery that buys it is where the gateway's real failure modes live:
local_ack/peer_acktrack log positions, not state. Two gateways whose WireGuard peer sets have drifted apart report identical healthy status, and the log that would repair them has been truncated. There is no metric an operator can alarm on.apply_pushed_entriesdiscards a whole batch when the first entry's seq is ahead oflocal_ack + 1, and leaves recovery to the pull path noticing later.gateway.tomlshipsinterval = "1m"), so an instance registered on node A is unroutable through node B for up to a minute.Fix
Dual-stack sync
HttpSyncNetworkgains a v2 leg posting to/wavekv/sync2/{store}. A peer still on wavekv 1.x has no such route and answers 404, whichpost_bytes_probesurfaces asOk(None)— distinct from a transport error. wavekv'sSyncManagerreads that as "this peer is v1", falls back to/wavekv/sync, caches the verdict per peer, and re-probes everyprotocol_reprobeso an upgraded peer is picked up without a restart.Serving the other direction needs nothing beyond mounting the route: a v2 gateway answers v1 peers through wavekv's compatibility shim, whose
is_snapshot = trueresponse makes an unmodified v1 client adopt coverage and merge in exactly delta-state order./wavekv/push/{store}carries opportunistic pushes. Per wavekv's rule R3 these merge data only and never move ack coverage, so loss, duplication and reordering are harmless and the periodic round stays the anti-entropy backstop. This is what cuts propagation latency from the sync interval to the coalescing window.Both new routes reuse
verify_gateway_peer(same-app_id mTLS) and the 16 MiB body cap, and decode throughSyncEnvelope::decode, which enforces the schema version and rejects trailing bytes — deliberately not the genericdecodeused for KV values.Admission control
kv/schema.rsconfines each store to the key shapes the gateway actually defines. wavekv enforces it insidemerge, which is the only place covering both sync directions — a check in the HTTP handler would see inbound requests but not entries arriving in a response. A rejection also parks that round's ack adoption (rule R1), so a peer sending inadmissible data keeps re-offering it rather than having it silently dropped.The two stores have disjoint schemas, so an ephemeral-store peer cannot plant
cert/...orinst/...keys.Observability
WaveKvStatusnow reports, per store, the state digest (hex SHA-256 over the replicated state) plus merged/rejected counters, and per peer the negotiatedprotocol("v1"/"v2"),heard_from, anddigest_mismatches.The digest is the operational point of this whole change: two converged replicas produce equal digests by construction, so comparing them across the cluster is both the promotion gate for the rollout and the standing divergence check afterwards.
buffered_logsis kept and marked deprecated — it is always 0 now — so existing clients keep decoding.Verification
cargo test -p dstack-gateway: 90 pass. The cross-version behaviour itself is covered exhaustively in Phala-Network/wavekv#3, whose suite runs the real, unmodified wavekv 1.0 crate from crates.io against v2 (mixed clusters, shim adoption, rollback, tombstones across versions, fault injection, clock skew). This PR adds the gateway-layer wire tests that suite cannot see:a_positionally_encoded_v1_request_is_still_acceptedSyncMessageencoded by a wavekv 1.x gateway still decodes herea_v1_peer_can_decode_our_sync_responsea_v2_envelope_survives_the_transport_framingthe_v1_shim_serves_a_complete_deltais_snapshot = truemerged_entries_outside_the_schema_are_refusedkv::schema(3 tests)cargo fmt --all --checkclean; clippy clean apart from the pre-existingmanual_repeat_ningateway/src/pp.rs:254.Rollout
Per wavekv RFC 0001 §8.4, upgrade one gateway CVM at a time. After each node, the promotion gate is cluster-wide digest equality via
WaveKvStatus, plusprotocolflipping to"v2"for upgraded pairs anddigest_mismatchesstaying at 0. Any anomaly: roll that node back alone — v2 writes the v1 snapshot container and a WAL that is a strict subset of the v1 op set, so a v1 binary reads the same data directory.Note that a pre-existing divergence in a live cluster will surface as a digest mismatch during the rollout. That is the tool working as intended — wavekv 1.x could not have told you — but operators should expect it rather than read it as an upgrade regression.
Follow-up
dstack/Cargo.tomlpoints wavekv at the PR branch. It must be repointed towavekv = "2.0"once Phala-Network/wavekv#3 is merged and released; the TODO is inline. This PR should not merge before that.