Skip to content

refactor #409: decouple tonic::Streaming from RaftEvent — channel bridge at server/core boundary#412

Merged
JoshuaChi merged 3 commits into
mainfrom
fix/409-decouple-streaming
Jun 21, 2026
Merged

refactor #409: decouple tonic::Streaming from RaftEvent — channel bridge at server/core boundary#412
JoshuaChi merged 3 commits into
mainfrom
fix/409-decouple-streaming

Conversation

@JoshuaChi

@JoshuaChi JoshuaChi commented Jun 21, 2026

Copy link
Copy Markdown
Contributor

What Does This PR Do?

Removes tonic::Streaming<T> from RaftEvent::StreamSnapshot and RaftEvent::InstallSnapshotChunk, replacing them with tokio mpsc channels. gRPC transport types are now confined to d-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 forced unsafe 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 converting tonic::Streaming<T>mpsc::Receiver<T> before sending any RaftEvent. Core only sees channels.


Checklist

Required:

  • make test passes
  • Added tests for new code
  • Commits squashed to 1-2 logical units

If changing APIs:

  • Explained why complexity is justified

Testing

How tested:

  • Unit tests: Added run_pull_transfer_test module with 6 tests — happy path, checksum-mismatch retry, max-retries exceeded, incomplete data stream, missing metadata, out-of-order chunk
  • All existing role-state tests (follower, learner, leader, candidate) updated to use mpsc::channel + drop(tx) pattern

For bug fixes:

  • test_pull_transfer_fails_on_incomplete_data_stream — fails without the None => breakErr(IncompleteSnapshot) fix

Does This Follow d-engine's Principles?

  • Solves a real problem for most users (not just my edge case)
  • Keeps implementation simple
  • Doesn't bloat the API surface

Reviewer Notes

Key files to focus on:

  • d-engine-core/src/event.rs — new StreamSnapshot / InstallSnapshotChunk signatures
  • d-engine-server/src/network/grpc/grpc_raft_service.rs — bridge task pattern
  • d-engine-core/src/network/background_snapshot_transfer.rs — pull transfer cleanup + bug fix

Side fixes included: handle_ack async fnfn, dead code removed from handle_retries, noisy debug log removed.

unsafe impl Sync for RaftEvent is gone — no longer needed once tonic::Streaming is out.

Estimated review complexity:

  • Deep (> 300 lines, 20 files touched)

Summary by CodeRabbit

Release Notes

  • Refactor

    • Reworked snapshot streaming and transfer to use message channels instead of gRPC streaming types, simplifying the snapshot pipeline end-to-end.
    • Updated snapshot delivery interfaces across the consensus roles and state machine.
  • Bug Fixes

    • Improved robustness of snapshot transfers with clearer failure handling for incomplete, out-of-order, and missing-metadata scenarios.
    • Enhanced retry behavior for transient chunk integrity issues.
  • Tests

    • Expanded coverage for snapshot streaming and transfer edge cases (retries, early closure, and validation failures).

…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)
@coderabbitai

coderabbitai Bot commented Jun 21, 2026

Copy link
Copy Markdown

Review Change Stack

Warning

Review limit reached

@JoshuaChi, we couldn't start this review because you've reached your PR review rate limit.

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 @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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 configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 4eeeb334-1690-46ee-9c03-d51f699a1bc5

📥 Commits

Reviewing files that changed from the base of the PR and between d87185c and 68d8de3.

📒 Files selected for processing (10)
  • d-engine-core/src/event.rs
  • d-engine-core/src/network/background_snapshot_transfer.rs
  • d-engine-core/src/network/background_snapshot_transfer_test.rs
  • d-engine-core/src/raft.rs
  • d-engine-core/src/raft_role/candidate_state.rs
  • d-engine-core/src/raft_role/follower_state.rs
  • d-engine-core/src/raft_role/leader_state.rs
  • d-engine-core/src/raft_role/leader_state_test/snapshot_test.rs
  • d-engine-core/src/raft_role/learner_state.rs
  • d-engine-server/src/network/grpc/grpc_raft_service.rs
📝 Walkthrough

Walkthrough

Snapshot streaming in d-engine-core is decoupled from tonic::Streaming by replacing all gRPC streaming types in RaftEvent::InstallSnapshotChunk and RaftEvent::StreamSnapshot with tokio::sync::mpsc channels. The gRPC layer now bridges tonic streams into mpsc at the network boundary. unsafe impl Sync for RaftEvent and StreamResponseSender are removed. All dependent roles, transfer logic, and tests are updated accordingly.

Changes

Snapshot pipeline decoupled from tonic::Streaming

Layer / File(s) Summary
RaftEvent variants, trait interfaces, and oneshot cleanup
d-engine-core/src/event.rs, d-engine-core/src/maybe_clone_oneshot.rs, d-engine-core/src/network/mod.rs, d-engine-core/src/state_machine_handler/mod.rs
InstallSnapshotChunk now carries mpsc::Receiver<SnapshotChunk>; StreamSnapshot carries mpsc::Receiver<SnapshotAck> and mpsc::Sender<Arc<SnapshotChunk>>; unsafe impl Sync and StreamResponseSender removed; Transport and StateMachineHandler trait method return/parameter types updated to mpsc.
gRPC boundary bridging
d-engine-server/src/network/grpc/grpc_raft_service.rs, d-engine-server/src/network/grpc/grpc_transport.rs
stream_snapshot spawns a task forwarding tonic ACKs into mpsc and returns a ReceiverStream over chunks; install_snapshot spawns a task forwarding tonic chunks into mpsc; grpc_transport spawns a task to forward the tonic snapshot response stream into a bounded mpsc receiver before returning it.
Leader StreamSnapshot event handling
d-engine-core/src/raft_role/leader_state.rs
Leader StreamSnapshot arm rewritten to accept ack_rx/chunk_tx, fetch snapshot metadata, signal startup readiness, and spawn BackgroundSnapshotTransfer::run_pull_transfer; prior stream construction and not_found response path removed.
Background pull transfer refactored
d-engine-core/src/network/background_snapshot_transfer.rs
run_pull_transfer accepts mpsc::Receiver<SnapshotAck> and mpsc::Sender<Arc<SnapshotChunk>>; handle_ack converted from async to sync; send_chunk sends Arc<SnapshotChunk> directly; incomplete stream returns IncompleteSnapshot; handle_retries gains max-retry skipping internally.
StateMachineHandler snapshot receive path
d-engine-core/src/state_machine_handler/default_state_machine_handler.rs, d-engine-core/src/state_machine_handler/mod.rs
apply_snapshot_stream_from_leader and process_snapshot_stream accept mpsc::Receiver<SnapshotChunk>; chunk receive loop uses recv() with timeout replacing tonic stream.next(); tonic/tokio-stream imports removed.
Non-leader role StreamSnapshot and learner fetch path
d-engine-core/src/raft_role/candidate_state.rs, d-engine-core/src/raft_role/follower_state.rs, d-engine-core/src/raft_role/learner_state.rs
Candidate, follower, and learner StreamSnapshot arms changed from permission_denied reply to log-and-send-error via startup_tx; learner fetch_initial_snapshot passes mpsc receiver into apply_snapshot_stream_from_leader.
Remove create_production_snapshot_stream utility
d-engine-core/src/utils/stream.rs, d-engine-core/src/utils/stream_test.rs
create_production_snapshot_stream and its gRPC frame-encoding test removed; stream.rs imports trimmed to GrpcStreamDecoder only; stream_test.rs retains only varint unit tests.
New run_pull_transfer test suite
d-engine-core/src/network/background_snapshot_transfer_test.rs
Seven new async tests cover happy path, checksum mismatch retry, max retries exceeded, incomplete stream, missing metadata, early ACK closure, and out-of-order chunk delivery for the mpsc-based pull transfer.
Role state and StateMachineHandler tests updated to mpsc
d-engine-core/src/raft_role/candidate_state_test.rs, d-engine-core/src/raft_role/follower_state_test.rs, d-engine-core/src/raft_role/leader_state_test/event_handling_test.rs, d-engine-core/src/raft_role/learner_state_test.rs, d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs
All snapshot install tests replace create_test_snapshot_stream/boxed stream construction with mpsc channel feeding; create_test_snapshot_stream import removed throughout.
Leader StreamSnapshot startup behavior tests
d-engine-core/src/raft_role/leader_state_test/snapshot_test.rs
New tests verify startup rejection when snapshot metadata unavailable and startup success when metadata present; chunk channel receives first chunk with seq == 0.

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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

  • deventlab/d-engine#360: Changes follower/learner snapshot install success computation to depend on apply_snapshot_stream_from_leader, which this PR refactors to accept an mpsc receiver.
  • deventlab/d-engine#368: Modifies the snapshot chunk receiving path to use SnapshotConfig.receive_chunk_timeout_in_sec instead of a hardcoded timeout, directly related to the process_snapshot_stream refactoring in this PR.

Poem

🐇 Hoppity hop through the pipeline we go,
No more tonic streams cluttering the flow!
mpsc channels, so clean and so bright,
unsafe Sync banished — the types are now right.
From gRPC boundary to state machine's door,
This rabbit sends chunks without unsafe anymore! 🎉

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately and specifically describes the main change: decoupling tonic::Streaming from RaftEvent and introducing a channel bridge at the server/core boundary.
Linked Issues check ✅ Passed The PR fully addresses issue #409 by removing tonic::Streaming from RaftEvent variants, replacing it with tokio mpsc channels, and eliminating the unsafe Sync impl.
Out of Scope Changes check ✅ Passed All changes directly support the core objective of decoupling tonic from RaftEvent. Minor additional improvements (yield_now removal, IncompleteSnapshot fix) are closely related bug fixes identified during the refactor.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/409-decouple-streaming

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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 | 🟠 Major

Snapshot 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 receive Ok(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 dropping tonic::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 lift

Preserve an explicit error path for missing or unloadable snapshots.

When snapshot_metadata() is None or load_snapshot_data(...).await? fails, chunk_tx is dropped and the gRPC bridge emits a successful empty stream. That replaces the previous not_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 to tonic::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 win

Make retry exhaustion terminal for missing ACKs.

retry_counts only increases on ChecksumMismatch, so a silent follower is retried until the global timeout instead of respecting max_retries. Count actual resend attempts in handle_retries and 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 value

Unnecessary clone of ack_tx.

The ack_tx.clone() on line 335 is redundant since ack_tx is moved into the function and not used after the call to process_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

📥 Commits

Reviewing files that changed from the base of the PR and between 09bd35c and d0c7fae.

📒 Files selected for processing (20)
  • d-engine-core/src/event.rs
  • d-engine-core/src/maybe_clone_oneshot.rs
  • d-engine-core/src/network/background_snapshot_transfer.rs
  • d-engine-core/src/network/background_snapshot_transfer_test.rs
  • d-engine-core/src/network/mod.rs
  • d-engine-core/src/raft_role/candidate_state.rs
  • d-engine-core/src/raft_role/candidate_state_test.rs
  • d-engine-core/src/raft_role/follower_state.rs
  • d-engine-core/src/raft_role/follower_state_test.rs
  • d-engine-core/src/raft_role/leader_state.rs
  • d-engine-core/src/raft_role/leader_state_test/event_handling_test.rs
  • d-engine-core/src/raft_role/learner_state.rs
  • d-engine-core/src/raft_role/learner_state_test.rs
  • d-engine-core/src/state_machine_handler/default_state_machine_handler.rs
  • d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs
  • d-engine-core/src/state_machine_handler/mod.rs
  • d-engine-core/src/utils/stream.rs
  • d-engine-core/src/utils/stream_test.rs
  • d-engine-server/src/network/grpc/grpc_raft_service.rs
  • d-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

Comment thread d-engine-core/src/event.rs
Comment thread d-engine-core/src/network/background_snapshot_transfer.rs
Comment thread d-engine-core/src/raft_role/candidate_state.rs Outdated
Comment thread d-engine-core/src/raft_role/follower_state.rs Outdated
Comment thread d-engine-core/src/raft_role/learner_state.rs Outdated
@codecov

codecov Bot commented Jun 21, 2026

Copy link
Copy Markdown

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🧹 Nitpick comments (1)
d-engine-core/src/raft_role/follower_state.rs (1)

503-507: 💤 Low value

Consider consistent error handling with other role states.

The warning message is now correct ("Follower" not "Candidate"). However, unlike CandidateState and LearnerState, this arm silently discards the startup_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

📥 Commits

Reviewing files that changed from the base of the PR and between d0c7fae and d87185c.

📒 Files selected for processing (10)
  • d-engine-core/src/event.rs
  • d-engine-core/src/network/background_snapshot_transfer.rs
  • d-engine-core/src/network/background_snapshot_transfer_test.rs
  • d-engine-core/src/raft.rs
  • d-engine-core/src/raft_role/candidate_state.rs
  • d-engine-core/src/raft_role/follower_state.rs
  • d-engine-core/src/raft_role/leader_state.rs
  • d-engine-core/src/raft_role/leader_state_test/snapshot_test.rs
  • d-engine-core/src/raft_role/learner_state.rs
  • d-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

@JoshuaChi
JoshuaChi force-pushed the fix/409-decouple-streaming branch from d87185c to e39c823 Compare June 21, 2026 08:56
… 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
@JoshuaChi
JoshuaChi force-pushed the fix/409-decouple-streaming branch from e39c823 to 68d8de3 Compare June 21, 2026 09:15
@JoshuaChi
JoshuaChi merged commit 158c356 into main Jun 21, 2026
9 checks passed
@JoshuaChi
JoshuaChi deleted the fix/409-decouple-streaming branch June 21, 2026 09:44
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.

refactor: decouple tonic::Streaming from RaftEvent — introduce SnapshotChunkStream trait

1 participant