refactor #409: decouple tonic::Streaming from RaftEvent — channel bridge at server/core boundary#412
Conversation
…dge at server/core boundary Replace tonic::Streaming<T> in StreamSnapshot and InstallSnapshotChunk with mpsc channels. gRPC types now stay in d-engine-server; core only sees tokio primitives. - Delete StreamResponseSender, create_production_snapshot_stream - Remove unsafe impl Sync for RaftEvent - Fix IncompleteSnapshot bug: data stream exhaustion was returning Ok(()) - handle_ack: async fn → fn, clean up dead code in handle_retries - Add 6 pull-transfer unit tests (happy path, retry, timeout, error cases)
|
Warning Review limit reached
More reviews will be available in 34 minutes and 55 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan refill rate. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, the refill rate gradually slows as usage increases. The highest same-day bursts are limited more strictly. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (10)
📝 WalkthroughWalkthroughSnapshot streaming in ChangesSnapshot pipeline decoupled from tonic::Streaming
Sequence Diagram(s)sequenceDiagram
participant Follower as gRPC Follower Client
participant Service as grpc_raft_service
participant Core as RaftCore / LeaderState
participant Transfer as BackgroundSnapshotTransfer
participant SM as StateMachineHandler
rect rgba(100, 149, 237, 0.5)
note over Service,Core: stream_snapshot (pull path)
Follower->>Service: stream_snapshot(ack_stream)
Service->>Service: spawn task: tonic SnapshotAck -> mpsc ack_rx
Service->>Core: RaftEvent::StreamSnapshot(ack_rx, chunk_tx, startup_tx)
Core->>Transfer: run_pull_transfer(ack_rx, chunk_tx)
Transfer-->>Service: Arc<SnapshotChunk> via chunk_tx
Service-->>Follower: ReceiverStream<SnapshotChunk>
Follower->>Service: SnapshotAck
Service->>Transfer: SnapshotAck via ack_rx
end
rect rgba(60, 179, 113, 0.5)
note over Service,SM: install_snapshot (push path)
Follower->>Service: install_snapshot(chunk_stream)
Service->>Service: spawn task: tonic SnapshotChunk -> mpsc rx
Service->>Core: RaftEvent::InstallSnapshotChunk(rx, resp_tx)
Core->>SM: apply_snapshot_stream_from_leader(rx)
SM->>SM: process_snapshot_stream(mpsc::Receiver)
SM-->>Core: Result
Core-->>Service: SnapshotResponse via resp_tx
Service-->>Follower: SnapshotResponse
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
d-engine-core/src/network/mod.rs (1)
319-325:⚠️ Potential issue | 🟠 MajorSnapshot stream errors after RPC handshake are silently collapsed to EOF.
Result<mpsc::Receiver<SnapshotChunk>>only reports initial RPC setup failures. Once the receiver is returned, post-handshake gRPC stream errors are logged and the sender dropped (grpc_transport.rs:584), leaving the state machine to receiveOk(None)indistinguishably from clean EOF (default_state_machine_handler.rs:922). This prevents proper error recovery: timeout detection works, but actual network/protocol errors are lost, breaking retry logic.Add either a
Result<SnapshotChunk, _>wrapper in the receiver, an explicit error side-channel, or carry a core/domain error instead of silently droppingtonic::Status.🤖 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 `@d-engine-core/src/network/mod.rs` around lines 319 - 325, The function request_snapshot_from_leader returns Result<mpsc::Receiver<SnapshotChunk>> which only captures initial RPC setup failures, but post-handshake gRPC stream errors are silently logged and dropped, causing the state machine to receive Ok(None) indistinguishably from a clean EOF and breaking retry logic. Modify the return type to properly propagate post-handshake errors by either changing the receiver to carry Result<SnapshotChunk, Error> as its item type instead of just SnapshotChunk, or by adding an explicit error side-channel (such as a separate oneshot or broadcast channel) that communicates tonic::Status errors that occur after the initial handshake. Update the corresponding error handling sites (such as the code in grpc_transport.rs that currently drops the error) to communicate failures through the chosen error propagation mechanism.d-engine-core/src/raft_role/leader_state.rs (1)
1392-1413:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftPreserve an explicit error path for missing or unloadable snapshots.
When
snapshot_metadata()isNoneorload_snapshot_data(...).await?fails,chunk_txis dropped and the gRPC bridge emits a successful empty stream. That replaces the previousnot_found("Snapshot not found")behavior with silent EOF. Add an out-of-band result/error channel, or make the core chunk channel carry a core error type that the server layer maps totonic::Status.🤖 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 `@d-engine-core/src/raft_role/leader_state.rs` around lines 1392 - 1413, The code currently drops the chunk_tx when snapshot_metadata() returns None or when load_snapshot_data() fails, which causes the gRPC stream to complete silently with empty data instead of returning an error to the caller. To fix this, introduce an error channel (such as an out-of-band result channel) or modify the chunk transmission mechanism to carry error information, then explicitly send an error response in both cases: when snapshot_metadata() is None and when load_snapshot_data() fails. This ensures the gRPC server layer receives the error and can properly map it to a tonic::Status error response instead of emitting a successful empty stream.d-engine-core/src/network/background_snapshot_transfer.rs (1)
312-324:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMake retry exhaustion terminal for missing ACKs.
retry_countsonly increases onChecksumMismatch, so a silent follower is retried until the global timeout instead of respectingmax_retries. Count actual resend attempts inhandle_retriesand fail when the retry budget is exhausted.Proposed fix
- Some(ack) => Self::handle_ack( + Some(ack) => Self::handle_ack( ack, &mut pending_acks, &mut retry_counts, - &config, )?,fn handle_ack( ack: SnapshotAck, pending_acks: &mut HashSet<u32>, retry_counts: &mut HashMap<u32, u32>, - config: &SnapshotConfig, ) -> Result<()> {Ok(ChunkStatus::ChecksumMismatch) => { trace!(?retry_counts, "ChecksumMismatch"); - - let count = retry_counts.entry(seq).or_insert(0); - *count += 1; - - if *count > config.max_retries { - trace!(%seq, "Max retries exceeded, removing from pending_acks"); - pending_acks.remove(&seq); // Remove if the retry limit is exceeded - return Err(SnapshotError::TransferFailed.into()); - } - + retry_counts.entry(seq).or_insert(0); // Will be resent in next retry cycle }for &seq in pending_acks.iter() { - let count = retry_counts.entry(seq).or_insert(0); - if *count > config.max_retries { - trace!(%seq, "skipping retry: max retries exceeded"); - continue; - } - - if let Some(chunk) = chunk_cache.get(&seq) { - Self::send_chunk(chunk_tx, chunk.clone(), max_bandwidth_mbps).await?; + let retries = retry_counts.get(&seq).copied().unwrap_or(0); + if retries >= config.max_retries { + trace!(%seq, retries, "max retries exceeded"); + return Err(SnapshotError::TransferFailed.into()); + } + + if let Some(chunk) = chunk_cache.get(&seq).cloned() { + Self::send_chunk(chunk_tx, chunk, max_bandwidth_mbps).await?; + retry_counts.insert(seq, retries + 1); } else { return Err(SnapshotError::ChunkNotCached(seq).into()); }Also applies to: 345-353
🤖 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 `@d-engine-core/src/network/background_snapshot_transfer.rs` around lines 312 - 324, The retry counting mechanism currently only increments retry_counts when a ChecksumMismatch response is received, which means entries that receive no response (silent followers) will never hit the max_retries limit and will be retried indefinitely until the global timeout. Instead of counting retries based on responses in the ChecksumMismatch handler, track actual resend attempts by incrementing the retry count in the handle_retries function itself every time an entry is resent. Then check if the retry count exceeds config.max_retries at the point of resend and fail with SnapshotError::TransferFailed if the retry budget is exhausted. Apply this same logic pattern to the similar retry logic mentioned in the secondary location around lines 345-353.
🧹 Nitpick comments (1)
d-engine-core/src/state_machine_handler/default_state_machine_handler.rs (1)
327-336: 💤 Low valueUnnecessary clone of
ack_tx.The
ack_tx.clone()on line 335 is redundant sinceack_txis moved into the function and not used after the call toprocess_snapshot_stream. You can pass it directly.♻️ Suggested simplification
let (final_metadata, snapshot_path) = self - .process_snapshot_stream(stream_chunk_receiver, ack_tx.clone(), config) + .process_snapshot_stream(stream_chunk_receiver, ack_tx, config) .await?;🤖 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 `@d-engine-core/src/state_machine_handler/default_state_machine_handler.rs` around lines 327 - 336, In the receive_snapshot_stream_from_leader function, the call to process_snapshot_stream is passing ack_tx.clone() when ack_tx can be passed directly. Since ack_tx is not used after this function call and will be moved into process_snapshot_stream, remove the unnecessary clone and pass ack_tx directly without cloning to the process_snapshot_stream method call.
🤖 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.
Inline comments:
In `@d-engine-core/src/event.rs`:
- Around line 159-162: The StreamSnapshot variant in the Event enum lacks a
mechanism to communicate rejection or startup status to the caller. When
snapshot metadata is unavailable or a non-leader receives the request, the
channels are silently dropped without feedback. Add a rejection/startup-result
communication path to StreamSnapshot by either including a oneshot channel field
to send initial rejection status, or modify the chunk sender field to carry
Result<SnapshotChunk, Error> instead of just SnapshotChunk, allowing the gRPC
handler to properly communicate failures back to the caller instead of returning
OK for all cases.
In `@d-engine-core/src/network/background_snapshot_transfer.rs`:
- Around line 211-219: The `None => break` pattern match arm in the ACK handling
block allows early ACK-channel closure to be treated as successful completion
(returning Ok(())), even when chunks are still pending. Instead of simply
breaking, check whether the transfer has been completed normally using the
appropriate completion predicate or by verifying that pending_acks is empty and
all chunks have been acknowledged. If the transfer is not complete when the ACK
channel closes, return an error indicating the unexpected channel closure,
otherwise allow the function to continue and return success only when
appropriate.
In `@d-engine-core/src/raft_role/candidate_state.rs`:
- Around line 467-470: The StreamSnapshot event handler in the
RaftEvent::StreamSnapshot arm is silently returning Ok(()) and dropping the
_chunk_tx channel, which causes wrong-role requests to appear as
successful-but-empty streams instead of being explicitly rejected. Instead of
just returning Ok(()), use the _chunk_tx channel to send an explicit error or
rejection signal to properly signal the role mismatch to the caller, ensuring
routing errors are not masked and recovery is not delayed.
In `@d-engine-core/src/raft_role/follower_state.rs`:
- Around line 503-505: The warning message in the RaftEvent::StreamSnapshot
match arm of FollowerState contains an incorrect role name. The message
currently states "Candidate should not receive StreamSnapshot event." but this
code is in the FollowerState implementation, not Candidate, so it should be
updated to reference "Follower" instead of "Candidate" to accurately reflect the
actual role and avoid confusion during debugging.
In `@d-engine-core/src/raft_role/learner_state.rs`:
- Around line 449-452: In the LearnerState implementation, the
RaftEvent::StreamSnapshot match arm contains a warning message that incorrectly
references "Candidate" instead of "Learner". Update the warn! macro call to
change the message from "Candidate should not receive StreamSnapshot event." to
"Learner should not receive StreamSnapshot event." to accurately reflect which
role is handling this event.
---
Outside diff comments:
In `@d-engine-core/src/network/background_snapshot_transfer.rs`:
- Around line 312-324: The retry counting mechanism currently only increments
retry_counts when a ChecksumMismatch response is received, which means entries
that receive no response (silent followers) will never hit the max_retries limit
and will be retried indefinitely until the global timeout. Instead of counting
retries based on responses in the ChecksumMismatch handler, track actual resend
attempts by incrementing the retry count in the handle_retries function itself
every time an entry is resent. Then check if the retry count exceeds
config.max_retries at the point of resend and fail with
SnapshotError::TransferFailed if the retry budget is exhausted. Apply this same
logic pattern to the similar retry logic mentioned in the secondary location
around lines 345-353.
In `@d-engine-core/src/network/mod.rs`:
- Around line 319-325: The function request_snapshot_from_leader returns
Result<mpsc::Receiver<SnapshotChunk>> which only captures initial RPC setup
failures, but post-handshake gRPC stream errors are silently logged and dropped,
causing the state machine to receive Ok(None) indistinguishably from a clean EOF
and breaking retry logic. Modify the return type to properly propagate
post-handshake errors by either changing the receiver to carry
Result<SnapshotChunk, Error> as its item type instead of just SnapshotChunk, or
by adding an explicit error side-channel (such as a separate oneshot or
broadcast channel) that communicates tonic::Status errors that occur after the
initial handshake. Update the corresponding error handling sites (such as the
code in grpc_transport.rs that currently drops the error) to communicate
failures through the chosen error propagation mechanism.
In `@d-engine-core/src/raft_role/leader_state.rs`:
- Around line 1392-1413: The code currently drops the chunk_tx when
snapshot_metadata() returns None or when load_snapshot_data() fails, which
causes the gRPC stream to complete silently with empty data instead of returning
an error to the caller. To fix this, introduce an error channel (such as an
out-of-band result channel) or modify the chunk transmission mechanism to carry
error information, then explicitly send an error response in both cases: when
snapshot_metadata() is None and when load_snapshot_data() fails. This ensures
the gRPC server layer receives the error and can properly map it to a
tonic::Status error response instead of emitting a successful empty stream.
---
Nitpick comments:
In `@d-engine-core/src/state_machine_handler/default_state_machine_handler.rs`:
- Around line 327-336: In the receive_snapshot_stream_from_leader function, the
call to process_snapshot_stream is passing ack_tx.clone() when ack_tx can be
passed directly. Since ack_tx is not used after this function call and will be
moved into process_snapshot_stream, remove the unnecessary clone and pass ack_tx
directly without cloning to the process_snapshot_stream method call.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 990cc355-e635-43ed-8795-7e42737e7266
📒 Files selected for processing (20)
d-engine-core/src/event.rsd-engine-core/src/maybe_clone_oneshot.rsd-engine-core/src/network/background_snapshot_transfer.rsd-engine-core/src/network/background_snapshot_transfer_test.rsd-engine-core/src/network/mod.rsd-engine-core/src/raft_role/candidate_state.rsd-engine-core/src/raft_role/candidate_state_test.rsd-engine-core/src/raft_role/follower_state.rsd-engine-core/src/raft_role/follower_state_test.rsd-engine-core/src/raft_role/leader_state.rsd-engine-core/src/raft_role/leader_state_test/event_handling_test.rsd-engine-core/src/raft_role/learner_state.rsd-engine-core/src/raft_role/learner_state_test.rsd-engine-core/src/state_machine_handler/default_state_machine_handler.rsd-engine-core/src/state_machine_handler/default_state_machine_handler_test.rsd-engine-core/src/state_machine_handler/mod.rsd-engine-core/src/utils/stream.rsd-engine-core/src/utils/stream_test.rsd-engine-server/src/network/grpc/grpc_raft_service.rsd-engine-server/src/network/grpc/grpc_transport.rs
💤 Files with no reviewable changes (2)
- d-engine-core/src/utils/stream_test.rs
- d-engine-core/src/maybe_clone_oneshot.rs
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
d-engine-core/src/raft_role/follower_state.rs (1)
503-507: 💤 Low valueConsider consistent error handling with other role states.
The warning message is now correct ("Follower" not "Candidate"). However, unlike
CandidateStateandLearnerState, this arm silently discards thestartup_tx.send()result. For observability consistency across roles, consider logging the error if the gRPC receiver was already dropped.♻️ Suggested change for consistency
RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx, startup_tx) => { - let _ = startup_tx.send(Err(Status::failed_precondition("Not the leader"))); debug!("Follower::RaftEvent::StreamSnapshot"); warn!("Follower should not receive StreamSnapshot event."); + if let Err(e) = startup_tx.send(Err(Status::failed_precondition("Not the leader"))) { + error!( + ?e, + "StreamSnapshot startup_tx send failed: gRPC receiver already dropped" + ); + } return Ok(()); }🤖 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 `@d-engine-core/src/raft_role/follower_state.rs` around lines 503 - 507, The RaftEvent::StreamSnapshot handler in FollowerState silently discards the startup_tx.send() result using let _ =, unlike CandidateState and LearnerState which log when the send fails. Remove the let _ = prefix from the startup_tx.send() call and instead check the result, logging a warning or debug message if the send fails to indicate the gRPC receiver was already dropped. This ensures consistent error observability and debugging capability across all role states when handling this event.
🤖 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.
Nitpick comments:
In `@d-engine-core/src/raft_role/follower_state.rs`:
- Around line 503-507: The RaftEvent::StreamSnapshot handler in FollowerState
silently discards the startup_tx.send() result using let _ =, unlike
CandidateState and LearnerState which log when the send fails. Remove the let _
= prefix from the startup_tx.send() call and instead check the result, logging a
warning or debug message if the send fails to indicate the gRPC receiver was
already dropped. This ensures consistent error observability and debugging
capability across all role states when handling this event.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e4c8f446-4bda-47cd-9110-4908298f8bf9
📒 Files selected for processing (10)
d-engine-core/src/event.rsd-engine-core/src/network/background_snapshot_transfer.rsd-engine-core/src/network/background_snapshot_transfer_test.rsd-engine-core/src/raft.rsd-engine-core/src/raft_role/candidate_state.rsd-engine-core/src/raft_role/follower_state.rsd-engine-core/src/raft_role/leader_state.rsd-engine-core/src/raft_role/leader_state_test/snapshot_test.rsd-engine-core/src/raft_role/learner_state.rsd-engine-server/src/network/grpc/grpc_raft_service.rs
💤 Files with no reviewable changes (1)
- d-engine-core/src/raft.rs
🚧 Files skipped from review as they are similar to previous changes (5)
- d-engine-core/src/raft_role/candidate_state.rs
- d-engine-core/src/raft_role/learner_state.rs
- d-engine-core/src/raft_role/leader_state.rs
- d-engine-core/src/network/background_snapshot_transfer_test.rs
- d-engine-core/src/network/background_snapshot_transfer.rs
d87185c to
e39c823
Compare
… regression
- Add oneshot startup channel to StreamSnapshot event: gRPC handler awaits
Ok/Err before returning stream, replacing silent channel drop with proper
gRPC status codes (not_found / failed_precondition)
- Leader: sends Ok(()) when metadata exists, not_found when snapshot missing
- Follower/Candidate/Learner: sends failed_precondition instead of dropping
- Fix wrong log messages ("Candidate" in follower/learner handlers)
- Fix early ACK channel closure: None arm now returns Err(TransferFailed)
when pending_acks is non-empty, instead of silently returning Ok(())
- Remove two yield_now() calls from raft main loop — benchmarks show +69%
on single-client write, +12% on high-concurrency write; select! await
already provides sufficient cooperative scheduling
- Add tests: StreamSnapshot rejection (no metadata, non-leader), early ACK
closure regression guard
e39c823 to
68d8de3
Compare
What Does This PR Do?
Removes
tonic::Streaming<T>fromRaftEvent::StreamSnapshotandRaftEvent::InstallSnapshotChunk, replacing them with tokiompscchannels. gRPC transport types are now confined tod-engine-server; core Raft logic only sees tokio primitives.Type:
Why Is This Needed?
Closes #409.
tonic::Streaming<T>is an HTTP/2 connection handle — it carries network state, is!Sync, and has no place in the core Raft protocol layer. Its presence forcedunsafe impl Sync for RaftEvent {}and created an architectural violation where the core layer had knowledge of the gRPC transport.The fix introduces a channel bridge in
d-engine-server: server spawns a task convertingtonic::Streaming<T>→mpsc::Receiver<T>before sending anyRaftEvent. Core only sees channels.Checklist
Required:
make testpassesIf changing APIs:
Testing
How tested:
run_pull_transfer_testmodule with 6 tests — happy path, checksum-mismatch retry, max-retries exceeded, incomplete data stream, missing metadata, out-of-order chunkmpsc::channel+drop(tx)patternFor bug fixes:
test_pull_transfer_fails_on_incomplete_data_stream— fails without theNone => break→Err(IncompleteSnapshot)fixDoes This Follow d-engine's Principles?
Reviewer Notes
Key files to focus on:
d-engine-core/src/event.rs— newStreamSnapshot/InstallSnapshotChunksignaturesd-engine-server/src/network/grpc/grpc_raft_service.rs— bridge task patternd-engine-core/src/network/background_snapshot_transfer.rs— pull transfer cleanup + bug fixSide fixes included:
handle_ackasync fn→fn, dead code removed fromhandle_retries, noisy debug log removed.unsafe impl Sync for RaftEventis gone — no longer needed oncetonic::Streamingis out.Estimated review complexity:
Summary by CodeRabbit
Release Notes
Refactor
Bug Fixes
Tests