From d0c7fae7e799854ad6b91144dcf9c5f62ca7215c Mon Sep 17 00:00:00 2001 From: Joshua Chi <112539+JoshuaChi@users.noreply.github.com> Date: Sun, 21 Jun 2026 12:01:49 +0800 Subject: [PATCH 1/3] =?UTF-8?q?refactor=20#409:=20decouple=20tonic::Stream?= =?UTF-8?q?ing=20from=20RaftEvent=20=E2=80=94=20channel=20bridge=20at=20se?= =?UTF-8?q?rver/core=20boundary?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace tonic::Streaming 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) --- d-engine-core/src/event.rs | 14 +- d-engine-core/src/maybe_clone_oneshot.rs | 43 --- .../network/background_snapshot_transfer.rs | 74 ++--- .../background_snapshot_transfer_test.rs | 293 +++++++++++++++++- d-engine-core/src/network/mod.rs | 2 +- .../src/raft_role/candidate_state.rs | 15 +- .../src/raft_role/candidate_state_test.rs | 7 +- d-engine-core/src/raft_role/follower_state.rs | 15 +- .../src/raft_role/follower_state_test.rs | 14 +- d-engine-core/src/raft_role/leader_state.rs | 28 +- .../leader_state_test/event_handling_test.rs | 8 +- d-engine-core/src/raft_role/learner_state.rs | 71 ++--- .../src/raft_role/learner_state_test.rs | 21 +- .../default_state_machine_handler.rs | 33 +- .../default_state_machine_handler_test.rs | 82 +++-- .../src/state_machine_handler/mod.rs | 2 +- d-engine-core/src/utils/stream.rs | 67 +--- d-engine-core/src/utils/stream_test.rs | 44 --- .../src/network/grpc/grpc_raft_service.rs | 73 +++-- .../src/network/grpc/grpc_transport.rs | 30 +- 20 files changed, 508 insertions(+), 428 deletions(-) diff --git a/d-engine-core/src/event.rs b/d-engine-core/src/event.rs index 80e8291f..2bbffbae 100644 --- a/d-engine-core/src/event.rs +++ b/d-engine-core/src/event.rs @@ -18,13 +18,13 @@ use d_engine_proto::server::storage::SnapshotAck; use d_engine_proto::server::storage::SnapshotChunk; use d_engine_proto::server::storage::SnapshotMetadata; use d_engine_proto::server::storage::SnapshotResponse; +use tokio::sync::mpsc; use tonic::Status; use crate::ApplyResult; use crate::MaybeCloneOneshotSender; use crate::Result; use crate::ScanResult; -use crate::StreamResponseSender; #[derive(Debug, Clone, PartialEq)] pub struct NewCommitData { @@ -151,12 +151,15 @@ pub enum RaftEvent { // Response snapshot stream from Leader InstallSnapshotChunk( - Box>, + mpsc::Receiver, MaybeCloneOneshotSender>, ), // Request snapshot stream from Leader - StreamSnapshot(Box>, StreamResponseSender), + StreamSnapshot( + mpsc::Receiver, + mpsc::Sender>, + ), JoinCluster( JoinRequest, @@ -193,11 +196,6 @@ pub enum RaftEvent { }, } -// SAFETY: tonic::Streaming variants are only accessed on the thread that created them; -// no shared references to snapshot variants cross thread boundaries. -// Full decoupling tracked in #409. -unsafe impl Sync for RaftEvent {} - #[cfg(test)] #[cfg_attr(test, derive(Debug, Clone))] #[allow(unused)] diff --git a/d-engine-core/src/maybe_clone_oneshot.rs b/d-engine-core/src/maybe_clone_oneshot.rs index a719d1fb..d1427409 100644 --- a/d-engine-core/src/maybe_clone_oneshot.rs +++ b/d-engine-core/src/maybe_clone_oneshot.rs @@ -17,11 +17,9 @@ use std::pin::Pin; use std::task::Context; use std::task::Poll; -use d_engine_proto::server::storage::SnapshotChunk; #[cfg(any(test, feature = "__test_support"))] use tokio::sync::broadcast; use tokio::sync::oneshot; -use tonic::Status; pub trait RaftOneshot { type Sender: Send + Sync; @@ -211,44 +209,3 @@ impl RaftOneshot for MaybeCloneOneshot { ) } } - -#[derive(Debug)] -pub struct StreamResponseSender { - inner: oneshot::Sender, Status>>, - - #[cfg(any(test, feature = "__test_support"))] - test_inner: - Option, Status>>>, -} - -impl StreamResponseSender { - pub fn new() -> ( - Self, - oneshot::Receiver, Status>>, - ) { - let (inner_tx, inner_rx) = oneshot::channel(); - ( - Self { - inner: inner_tx, - #[cfg(any(test, feature = "__test_support"))] - test_inner: None, - }, - inner_rx, - ) - } - - pub fn send( - self, - value: std::result::Result, Status>, - ) -> Result<(), Box, Status>>> { - #[cfg(not(any(test, feature = "__test_support")))] - return self.inner.send(value).map_err(Box::new); - - #[cfg(any(test, feature = "__test_support"))] - if let Some(tx) = self.test_inner { - tx.send(value).map(|_| ()).map_err(|e| Box::new(e.0)) - } else { - self.inner.send(value).map_err(Box::new) - } - } -} diff --git a/d-engine-core/src/network/background_snapshot_transfer.rs b/d-engine-core/src/network/background_snapshot_transfer.rs index cd9c304d..7bad0b2a 100644 --- a/d-engine-core/src/network/background_snapshot_transfer.rs +++ b/d-engine-core/src/network/background_snapshot_transfer.rs @@ -19,7 +19,6 @@ use tokio::sync::mpsc::error::TrySendError; use tokio::time::Instant; use tokio::time::sleep; use tokio_stream::wrappers::ReceiverStream; -use tonic::Status; use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use tracing::debug; @@ -169,16 +168,15 @@ where // Unified pull transfer entry point pub(crate) async fn run_pull_transfer( - ack_stream: Box>, - chunk_tx: mpsc::Sender, Status>>, + ack_rx: mpsc::Receiver, + chunk_tx: mpsc::Sender>, mut data_stream: BoxStream<'static, Result>, config: SnapshotConfig, ) -> Result<()> { debug!("Starting pull snapshot transfer"); // Create processing pipeline - let transfer_fut = - Self::process_transfer(ack_stream, data_stream, chunk_tx, config.clone()); + let transfer_fut = Self::process_transfer(ack_rx, chunk_tx, data_stream, config.clone()); // Run with timeout tokio::select! { @@ -191,9 +189,9 @@ where // Dedicated pull logic async fn process_transfer( - mut ack_stream: Box>, + mut ack_rx: mpsc::Receiver, + chunk_tx: mpsc::Sender>, mut data_stream: Pin> + Send>>, - chunk_tx: mpsc::Sender, Status>>, config: SnapshotConfig, ) -> Result<()> { let mut chunk_cache = LruCache::new(NonZero::new(config.cache_size).unwrap()); @@ -207,19 +205,16 @@ where loop { tokio::select! { // Process incoming ACKs - ack = ack_stream.next() => { + ack = ack_rx.recv() => { trace!("receive new ack"); match ack { - Some(Ok(ack)) => Self::handle_ack( + Some(ack) => Self::handle_ack( ack, &mut pending_acks, &mut retry_counts, &config, - ).await?, - - Some(Err(e)) => return Err(NetworkError::TonicStatusError(Box::new(e)).into()), - + )?, None => break, // ACK stream closed } }, @@ -253,37 +248,24 @@ where next_seq += 1; } Some(Err(e)) => return Err(e), - None => break, // Data stream exhausted + None => return Err(SnapshotError::IncompleteSnapshot.into()),// Data stream exhausted } }, // Retry mechanism (separate timer) _ = sleep(Duration::from_millis(config.retry_interval_in_ms)), if !pending_acks.is_empty() => { - // Add filtering: only process blocks that need to be retried - let needs_retry: Vec = pending_acks.iter() - .filter(|&&seq| { - // Only process blocks whose retry times are not exceeded - retry_counts.get(&seq).is_none_or(|&c| c <= config.max_retries) - }) - .copied() - .collect(); - - if !needs_retry.is_empty() { - trace!(?retry_counts, ?needs_retry, ?pending_acks); - - Self::handle_retries( - &pending_acks, - &mut retry_counts, - &mut chunk_cache, - &chunk_tx, - &config, - ).await?; - } + trace!(?retry_counts, ?pending_acks, "retrying pending chunks"); + Self::handle_retries( + &pending_acks, + &mut retry_counts, + &mut chunk_cache, + &chunk_tx, + &config, + ).await?; } } // Completion check - debug!(?total_chunks, "------ total_chunks"); if let Some(total) = total_chunks && next_seq >= total @@ -299,7 +281,7 @@ where } // Handle ACK messages with proper error management - async fn handle_ack( + fn handle_ack( ack: SnapshotAck, pending_acks: &mut HashSet, retry_counts: &mut HashMap, @@ -355,25 +337,15 @@ where pending_acks: &HashSet, retry_counts: &mut HashMap, chunk_cache: &mut LruCache>, - chunk_tx: &mpsc::Sender, Status>>, + chunk_tx: &mpsc::Sender>, config: &SnapshotConfig, ) -> Result<()> { let max_bandwidth_mbps = config.max_bandwidth_mbps; - // Create a snapshot of pending_acks to avoid concurrent modification issues - let pending_snapshot: Vec = pending_acks.iter().copied().collect(); - - for seq in pending_snapshot { - // Double-check if still pending before sending - if !pending_acks.contains(&seq) { - trace!(%seq, "Skipping retry for already-acked chunk"); - continue; - } - - // Skip if max retries exceeded + for &seq in pending_acks.iter() { let count = retry_counts.entry(seq).or_insert(0); if *count > config.max_retries { - trace!(%seq, "Skipping retry for chunk with max retries exceeded"); + trace!(%seq, "skipping retry: max retries exceeded"); continue; } @@ -447,14 +419,14 @@ where // Send chunk with proper error handling async fn send_chunk( - chunk_tx: &mpsc::Sender, Status>>, + chunk_tx: &mpsc::Sender>, chunk: Arc, max_bandwidth_mbps: u32, ) -> Result<()> { Self::apply_rate_limit(&chunk, max_bandwidth_mbps).await; chunk_tx - .send(Ok(chunk)) + .send(chunk) .await .map_err(|_| SnapshotError::ReceiverDisconnected.into()) } diff --git a/d-engine-core/src/network/background_snapshot_transfer_test.rs b/d-engine-core/src/network/background_snapshot_transfer_test.rs index 19381131..78c8e367 100644 --- a/d-engine-core/src/network/background_snapshot_transfer_test.rs +++ b/d-engine-core/src/network/background_snapshot_transfer_test.rs @@ -411,4 +411,295 @@ mod run_push_transfer_test { } } -// ... rest of the file remains unchanged ... +#[cfg(test)] +mod run_pull_transfer_test { + use std::sync::Arc; + + use d_engine_proto::server::storage::SnapshotAck; + use d_engine_proto::server::storage::snapshot_ack::ChunkStatus; + use futures::stream; + use tokio::sync::mpsc; + + use super::*; + use crate::MockTypeConfig; + use crate::SnapshotConfig; + use crate::SnapshotError; + use crate::create_test_chunk; + + /// Config tuned for pull tests: no rate limiting, fast retry, generous timeout. + fn pull_config() -> SnapshotConfig { + SnapshotConfig { + max_bandwidth_mbps: 0, + transfer_timeout_in_sec: 10, + retry_interval_in_ms: 20, + max_retries: 3, + cache_size: 16, + ..Default::default() + } + } + + /// Build a valid SnapshotChunk with the given seq / total. + /// Seq 0 always carries metadata (required by the protocol). + fn make_chunk( + seq: u32, + total: u32, + ) -> SnapshotChunk { + let data = vec![seq as u8; 64]; + let mut chunk = create_test_chunk(seq, &data, 1, 1, total); + // Only the first chunk should carry metadata; clear it for subsequent chunks. + if seq != 0 { + chunk.metadata = None; + } + chunk + } + + fn make_ack( + seq: u32, + status: ChunkStatus, + ) -> SnapshotAck { + SnapshotAck { + seq, + status: status as i32, + next_requested: seq + 1, + } + } + + // Test: All chunks sent and ACKed in order — transfer completes successfully. + // + // The leader sends 2 chunks; the follower ACKs each as Accepted. + // Expected: run_pull_transfer returns Ok(()). + // Expected: chunks arrive at the receiver in seq order (0, 1). + #[tokio::test] + async fn test_pull_transfer_happy_path() { + let config = pull_config(); + let (ack_tx, ack_rx) = mpsc::channel::(32); + let (chunk_tx, mut chunk_rx) = mpsc::channel::>(32); + + let data_stream = stream::iter(vec![Ok(make_chunk(0, 2)), Ok(make_chunk(1, 2))]).boxed(); + + let transfer = tokio::spawn(async move { + BackgroundSnapshotTransfer::::run_pull_transfer( + ack_rx, + chunk_tx, + data_stream, + config, + ) + .await + }); + + // Simulate follower: receive each chunk and ACK it as Accepted. + for expected_seq in 0..2u32 { + let chunk = chunk_rx.recv().await.expect("should receive chunk"); + assert_eq!(chunk.seq, expected_seq, "chunks must arrive in seq order"); + ack_tx + .send(make_ack(expected_seq, ChunkStatus::Accepted)) + .await + .expect("ack channel should be open"); + } + + let result = transfer.await.expect("task should not panic"); + assert!(result.is_ok(), "expected Ok(()), got: {result:?}"); + } + + // Test: Checksum mismatch causes the same chunk to be retransmitted. + // Transfer succeeds once the follower ACKs the retried chunk as Accepted. + // + // Setup: 1-chunk stream. Follower sends ChecksumMismatch on first receipt, + // then Accepted on the retry. + // Expected: chunk_rx receives the same seq twice. + // Expected: run_pull_transfer returns Ok(()). + #[tokio::test] + async fn test_pull_transfer_retry_on_checksum_mismatch() { + let mut config = pull_config(); + config.max_retries = 1; + config.retry_interval_in_ms = 10; + + let (ack_tx, ack_rx) = mpsc::channel::(32); + let (chunk_tx, mut chunk_rx) = mpsc::channel::>(32); + + let expected_seq = 0; + let data_stream = stream::iter(vec![Ok(make_chunk(expected_seq, 1))]).boxed(); + + let transfer = tokio::spawn(async move { + BackgroundSnapshotTransfer::::run_pull_transfer( + ack_rx, + chunk_tx, + data_stream, + config, + ) + .await + }); + + // First delivery: follower rejects with ChecksumMismatch. + let first = chunk_rx.recv().await.expect("should receive initial chunk"); + assert_eq!(first.seq, 0); + ack_tx + .send(make_ack(0, ChunkStatus::ChecksumMismatch)) + .await + .expect("ack channel should be open"); + + // After retry_interval_in_ms the same chunk must be retransmitted. + // Follower ACKs as Accepted this time. + let retried = chunk_rx.recv().await.expect("should receive retried chunk"); + assert_eq!(retried.seq, 0, "retried chunk must carry the same seq"); + ack_tx + .send(make_ack(0, ChunkStatus::Accepted)) + .await + .expect("ack channel should be open"); + + let result = transfer.await.expect("task should not panic"); + assert!(result.is_ok(), "expected Ok(()), got: {result:?}"); + } + + // Test: Exceeding max_retries on a single chunk causes the transfer to fail. + // + // Setup: 1-chunk stream; follower always replies ChecksumMismatch. + // max_retries = 1, so the second mismatch pushes the count over the limit. + // Expected: run_pull_transfer returns Err(TransferFailed). + #[tokio::test] + async fn test_pull_transfer_fails_on_max_retries_exceeded() { + let mut config = pull_config(); + config.max_retries = 1; + config.retry_interval_in_ms = 10; + + let (ack_tx, ack_rx) = mpsc::channel::(32); + let (chunk_tx, mut chunk_rx) = mpsc::channel::>(32); + + let data_stream = stream::iter(vec![Ok(make_chunk(0, 1))]).boxed(); + + let transfer = tokio::spawn(async move { + BackgroundSnapshotTransfer::::run_pull_transfer( + ack_rx, + chunk_tx, + data_stream, + config, + ) + .await + }); + + // Keep sending ChecksumMismatch until the transfer gives up. + // We drain chunk_rx in a loop so the transfer is never blocked waiting to send. + tokio::spawn(async move { + while let Some(chunk) = chunk_rx.recv().await { + let _ = ack_tx.send(make_ack(chunk.seq, ChunkStatus::ChecksumMismatch)).await; + } + }); + + let result = transfer.await.expect("task should not panic"); + assert!( + matches!( + result, + Err(crate::Error::Consensus(crate::ConsensusError::Snapshot( + SnapshotError::TransferFailed + ))) + ), + "expected TransferFailed, got: {result:?}" + ); + } + + // Test: Data stream ending before all declared chunks are delivered returns an error. + // + // Setup: data_stream declares total_chunks=2 but only yields seq=0 then closes. + // The transfer sends seq=0, then tries to read seq=1 — stream returns None. + // The transfer fails immediately; no ACK is needed. + // Expected: run_pull_transfer returns Err(IncompleteSnapshot). + #[tokio::test] + async fn test_pull_transfer_fails_on_incomplete_data_stream() { + let config = pull_config(); + let (_ack_tx, ack_rx) = mpsc::channel::(32); + // Buffer large enough so the send of seq=0 does not block before the error fires. + let (chunk_tx, _chunk_rx) = mpsc::channel::>(32); + + // Declare total_chunks=2 but only provide seq=0. + let data_stream = stream::iter(vec![Ok(make_chunk(0, 2))]).boxed(); + + let result = BackgroundSnapshotTransfer::::run_pull_transfer( + ack_rx, + chunk_tx, + data_stream, + config, + ) + .await; + + assert!( + matches!( + result, + Err(crate::Error::Consensus(crate::ConsensusError::Snapshot( + SnapshotError::IncompleteSnapshot + ))) + ), + "expected IncompleteSnapshot, got: {result:?}" + ); + } + + // Test: First chunk missing metadata causes an immediate protocol error. + // + // Expected: run_pull_transfer returns Err(MissingMetadata) without reading + // any further chunks. + #[tokio::test] + async fn test_pull_transfer_fails_on_missing_metadata() { + let config = pull_config(); + let (_ack_tx, ack_rx) = mpsc::channel::(32); + let (chunk_tx, _chunk_rx) = mpsc::channel::>(32); + + let mut bad_chunk = make_chunk(0, 1); + bad_chunk.metadata = None; // violates protocol requirement + + let data_stream = stream::iter(vec![Ok(bad_chunk)]).boxed(); + + let result = BackgroundSnapshotTransfer::::run_pull_transfer( + ack_rx, + chunk_tx, + data_stream, + config, + ) + .await; + + assert!( + matches!( + result, + Err(crate::Error::Consensus(crate::ConsensusError::Snapshot( + SnapshotError::MissingMetadata + ))) + ), + "expected MissingMetadata, got: {result:?}" + ); + } + + // Test: Chunks arriving out of sequence cause an immediate protocol error. + // + // Setup: data_stream yields seq=0 then seq=2 (skipping seq=1). + // Expected: run_pull_transfer returns Err(OutOfOrderChunk). + #[tokio::test] + async fn test_pull_transfer_fails_on_out_of_order_chunk() { + let config = pull_config(); + let (_ack_tx, ack_rx) = mpsc::channel::(32); + let (chunk_tx, mut chunk_rx) = mpsc::channel::>(32); + + let data_stream = stream::iter(vec![Ok(make_chunk(0, 3)), Ok(make_chunk(2, 3))]).boxed(); + + let transfer = tokio::spawn(async move { + BackgroundSnapshotTransfer::::run_pull_transfer( + ack_rx, + chunk_tx, + data_stream, + config, + ) + .await + }); + + // Drain chunk_rx so the transfer is not blocked on sends. + drop(chunk_rx.recv().await); // consume seq=0 + + let result = transfer.await.expect("task should not panic"); + assert!( + matches!( + result, + Err(crate::Error::Consensus(crate::ConsensusError::Snapshot( + SnapshotError::OutOfOrderChunk + ))) + ), + "expected OutOfOrderChunk, got: {result:?}" + ); + } +} diff --git a/d-engine-core/src/network/mod.rs b/d-engine-core/src/network/mod.rs index 094544c6..7c723993 100644 --- a/d-engine-core/src/network/mod.rs +++ b/d-engine-core/src/network/mod.rs @@ -322,7 +322,7 @@ where ack_tx: tokio::sync::mpsc::Receiver, retry: &crate::InstallSnapshotBackoffPolicy, membership: std::sync::Arc>, - ) -> Result>>; + ) -> Result>; /// Opens a persistent bidirectional AppendEntries stream to the given peer. /// diff --git a/d-engine-core/src/raft_role/candidate_state.rs b/d-engine-core/src/raft_role/candidate_state.rs index 3b87ecd5..b290a0e5 100644 --- a/d-engine-core/src/raft_role/candidate_state.rs +++ b/d-engine-core/src/raft_role/candidate_state.rs @@ -464,18 +464,9 @@ impl RaftRoleState for CandidateState { return Ok(()); } - RaftEvent::StreamSnapshot(request, sender) => { - debug!(?request, "Candidate::RaftEvent::StreamSnapshot"); - sender - .send(Err(Status::permission_denied( - "Candidate should not receive StreamSnapshot event.", - ))) - .map_err(|e| { - let error_str = format!("{e:?}"); - error!("Failed to send: {}", error_str); - NetworkError::SingalSendFailed(error_str) - })?; - + RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx) => { + debug!("Candidate::RaftEvent::StreamSnapshot"); + warn!("Candidate should not receive StreamSnapshot event."); return Ok(()); } diff --git a/d-engine-core/src/raft_role/candidate_state_test.rs b/d-engine-core/src/raft_role/candidate_state_test.rs index 497f578a..45b3de91 100644 --- a/d-engine-core/src/raft_role/candidate_state_test.rs +++ b/d-engine-core/src/raft_role/candidate_state_test.rs @@ -34,7 +34,6 @@ use crate::RoleEvent; use crate::raft_role::candidate_state::CandidateState; use crate::raft_role::role_state::RaftRoleState; use crate::test_utils::create_test_chunk; -use crate::test_utils::create_test_snapshot_stream; use crate::test_utils::mock::MockTypeConfig; use crate::test_utils::mock::mock_election_core; use crate::test_utils::mock::mock_raft_context; @@ -915,8 +914,10 @@ async fn test_handle_install_snapshot_returns_permission_denied() { let mut state = CandidateState::::new(1, context.node_config.clone()); let (resp_tx, mut resp_rx) = MaybeCloneOneshot::new(); - let stream = create_test_snapshot_stream(vec![create_test_chunk(0, b"chunk0", 1, 1, 2)]); - let raft_event = RaftEvent::InstallSnapshotChunk(Box::new(stream), resp_tx); + let (tx, rx) = mpsc::channel(32); + tx.send(create_test_chunk(0, b"chunk0", 1, 1, 2)).await.unwrap(); + drop(tx); + let raft_event = RaftEvent::InstallSnapshotChunk(rx, resp_tx); let result = state.handle_raft_event(raft_event, &context, mpsc::unbounded_channel().0).await; diff --git a/d-engine-core/src/raft_role/follower_state.rs b/d-engine-core/src/raft_role/follower_state.rs index 36d29899..4533d8e1 100644 --- a/d-engine-core/src/raft_role/follower_state.rs +++ b/d-engine-core/src/raft_role/follower_state.rs @@ -500,18 +500,9 @@ impl RaftRoleState for FollowerState { return Ok(()); } - RaftEvent::StreamSnapshot(request, sender) => { - debug!(?request, "Follower::RaftEvent::StreamSnapshot"); - sender - .send(Err(Status::permission_denied( - "Follower should not receive StreamSnapshot event.", - ))) - .map_err(|e| { - let error_str = format!("{e:?}"); - error!("Failed to send: {}", error_str); - NetworkError::SingalSendFailed(error_str) - })?; - + RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx) => { + debug!("Follower::RaftEvent::StreamSnapshot"); + warn!("Candidate should not receive StreamSnapshot event."); return Ok(()); } diff --git a/d-engine-core/src/raft_role/follower_state_test.rs b/d-engine-core/src/raft_role/follower_state_test.rs index 95d274ce..8b851c67 100644 --- a/d-engine-core/src/raft_role/follower_state_test.rs +++ b/d-engine-core/src/raft_role/follower_state_test.rs @@ -3,7 +3,6 @@ use crate::client::ClientWriteRequest; use crate::client::ErrorCode; use crate::client::WriteOperation; use crate::config::ReadConsistencyPolicy; -use crate::test_utils::create_test_snapshot_stream; use d_engine_proto::common::LogId; use d_engine_proto::common::NodeRole; use d_engine_proto::server::cluster::ClusterConfChangeRequest; @@ -2549,13 +2548,14 @@ async fn test_follower_install_snapshot_reports_success_when_apply_succeeds() { let mut state = FollowerState::::new(1, context.node_config.clone(), None, None); - let stream = create_test_snapshot_stream(vec![SnapshotChunk::default()]); let (resp_tx, mut resp_rx) = MaybeCloneOneshot::new(); let (role_tx, _role_rx) = mpsc::unbounded_channel(); - + let (tx, rx) = mpsc::channel(32); + tx.send(SnapshotChunk::default()).await.unwrap(); + drop(tx); state .handle_raft_event( - RaftEvent::InstallSnapshotChunk(Box::new(stream), resp_tx), + RaftEvent::InstallSnapshotChunk(rx, resp_tx), &context, role_tx, ) @@ -2616,14 +2616,16 @@ async fn test_follower_install_snapshot_reports_failure_when_apply_fails_after_t let mut state = FollowerState::::new(1, context.node_config.clone(), None, None); - let stream = create_test_snapshot_stream(vec![SnapshotChunk::default()]); let (resp_tx, mut resp_rx) = MaybeCloneOneshot::new(); let (role_tx, _role_rx) = mpsc::unbounded_channel(); // Follower absorbs the error and continues (does not propagate) + let (tx, rx) = mpsc::channel(32); + tx.send(SnapshotChunk::default()).await.unwrap(); + drop(tx); let _ = state .handle_raft_event( - RaftEvent::InstallSnapshotChunk(Box::new(stream), resp_tx), + RaftEvent::InstallSnapshotChunk(rx, resp_tx), &context, role_tx, ) diff --git a/d-engine-core/src/raft_role/leader_state.rs b/d-engine-core/src/raft_role/leader_state.rs index 923c0e2c..fa35359c 100644 --- a/d-engine-core/src/raft_role/leader_state.rs +++ b/d-engine-core/src/raft_role/leader_state.rs @@ -45,7 +45,6 @@ use crate::alias::TROF; use crate::client::{ClientReadRequest, ClientResponse, ErrorCode}; use crate::event::ClientCmd; use crate::network::Transport; -use crate::stream::create_production_snapshot_stream; use async_trait::async_trait; use d_engine_proto::common::AddNode; use d_engine_proto::common::BatchPromote; @@ -65,7 +64,6 @@ use d_engine_proto::server::election::VotedFor; use d_engine_proto::server::replication::AppendEntriesRequest; use d_engine_proto::server::replication::AppendEntriesResponse; use d_engine_proto::server::replication::append_entries_response; -use d_engine_proto::server::storage::SnapshotChunk; use d_engine_proto::server::storage::SnapshotMetadata; use rand::distr::SampleString; use std::collections::BTreeMap; @@ -1388,24 +1386,11 @@ impl RaftRoleState for LeaderState { panic!("{}", msg); } } - RaftEvent::StreamSnapshot(request, sender) => { + RaftEvent::StreamSnapshot(ack_rx, chunk_tx) => { debug!("Leader::RaftEvent::StreamSnapshot"); // Get the latest snapshot metadata if let Some(metadata) = ctx.state_machine().snapshot_metadata() { - // Create response channel - let (response_tx, response_rx) = - mpsc::channel::, Status>>(32); - // Convert to properly encoded tonic stream - let size = 1024 * 1024 * 1024; // 1GB max message size - let response_stream = create_production_snapshot_stream(response_rx, size); - // Immediately respond with the stream - sender.send(Ok(response_stream)).map_err(|e| { - let error_str = format!("{e:?}"); - error!("Stream response failed: {}", error_str); - NetworkError::SingalSendFailed(error_str) - })?; - // Spawn background transfer task let state_machine_handler = ctx.state_machine_handler().clone(); let config = ctx.node_config.raft.snapshot.clone(); @@ -1415,8 +1400,8 @@ impl RaftRoleState for LeaderState { tokio::spawn(async move { if let Err(e) = BackgroundSnapshotTransfer::::run_pull_transfer( - request, - response_tx, + ack_rx, + chunk_tx, data_stream, config, ) @@ -1425,13 +1410,6 @@ impl RaftRoleState for LeaderState { error!("StreamSnapshot failed: {:?}", e); } }); - } else { - warn!("No snapshot available for streaming"); - sender.send(Err(Status::not_found("Snapshot not found"))).map_err(|e| { - let error_str = format!("{e:?}"); - error!("Stream response failed: {}", error_str); - NetworkError::SingalSendFailed(error_str) - })?; } } RaftEvent::PromoteReadyLearners => { diff --git a/d-engine-core/src/raft_role/leader_state_test/event_handling_test.rs b/d-engine-core/src/raft_role/leader_state_test/event_handling_test.rs index ae61215a..a6ad073a 100644 --- a/d-engine-core/src/raft_role/leader_state_test/event_handling_test.rs +++ b/d-engine-core/src/raft_role/leader_state_test/event_handling_test.rs @@ -22,9 +22,9 @@ use crate::event::{NewCommitData, RoleEvent}; use crate::maybe_clone_oneshot::RaftOneshot; use crate::raft_role::leader_state::LeaderState; use crate::role_state::RaftRoleState; +use crate::test_utils::create_test_chunk; use crate::test_utils::mock::mock_raft_context; use crate::test_utils::mock::{MockBuilder, MockTypeConfig}; -use crate::test_utils::{create_test_chunk, create_test_snapshot_stream}; use crate::utils::convert::safe_kv_bytes; use d_engine_proto::server::cluster::{ ClusterConfChangeRequest, ClusterMembership, MetadataRequest, @@ -829,8 +829,10 @@ async fn test_handle_install_snapshot_returns_permission_denied() { use crate::maybe_clone_oneshot::MaybeCloneOneshot; let (resp_tx, mut resp_rx) = >::new(); - let stream = create_test_snapshot_stream(vec![create_test_chunk(0, b"chunk0", 1, 1, 2)]); - let raft_event = RaftEvent::InstallSnapshotChunk(Box::new(stream), resp_tx); + let (tx, rx) = mpsc::channel(32); + tx.send(create_test_chunk(0, b"chunk0", 1, 1, 2)).await.unwrap(); + drop(tx); + let raft_event = RaftEvent::InstallSnapshotChunk(rx, resp_tx); let (role_tx, mut role_rx) = mpsc::unbounded_channel(); assert!(state.handle_raft_event(raft_event, &context, role_tx).await.is_err()); diff --git a/d-engine-core/src/raft_role/learner_state.rs b/d-engine-core/src/raft_role/learner_state.rs index b7bc4eb9..d1fd8933 100644 --- a/d-engine-core/src/raft_role/learner_state.rs +++ b/d-engine-core/src/raft_role/learner_state.rs @@ -1,30 +1,3 @@ -use std::fmt::Debug; -use std::marker::PhantomData; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::Duration; - -use async_trait::async_trait; -use d_engine_proto::common::LogId; -use d_engine_proto::common::NodeRole; -use d_engine_proto::common::NodeRole::Learner; -use d_engine_proto::server::cluster::ClusterConfUpdateResponse; -use d_engine_proto::server::cluster::JoinRequest; -use d_engine_proto::server::cluster::LeaderDiscoveryRequest; -use d_engine_proto::server::cluster::LeaderDiscoveryResponse; -use d_engine_proto::server::election::VoteResponse; -use d_engine_proto::server::election::VotedFor; -use d_engine_proto::server::storage::SnapshotAck; -use d_engine_proto::server::storage::SnapshotResponse; -use tokio::sync::mpsc::{self}; -use tokio::time::Instant; -use tonic::Status; -use tracing::debug; -use tracing::error; -use tracing::info; -use tracing::trace; -use tracing::warn; - use super::RaftRole; use super::SharedState; use super::StateSnapshot; @@ -48,6 +21,31 @@ use crate::StateTransitionError; use crate::Transport; use crate::TypeConfig; use crate::alias::MOF; +use async_trait::async_trait; +use d_engine_proto::common::LogId; +use d_engine_proto::common::NodeRole; +use d_engine_proto::common::NodeRole::Learner; +use d_engine_proto::server::cluster::ClusterConfUpdateResponse; +use d_engine_proto::server::cluster::JoinRequest; +use d_engine_proto::server::cluster::LeaderDiscoveryRequest; +use d_engine_proto::server::cluster::LeaderDiscoveryResponse; +use d_engine_proto::server::election::VoteResponse; +use d_engine_proto::server::election::VotedFor; +use d_engine_proto::server::storage::SnapshotAck; +use d_engine_proto::server::storage::SnapshotResponse; +use std::fmt::Debug; +use std::marker::PhantomData; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; +use tokio::sync::mpsc::{self}; +use tokio::time::Instant; +use tonic::Status; +use tracing::debug; +use tracing::error; +use tracing::info; +use tracing::trace; +use tracing::warn; /// Learner node's state in Raft cluster. /// @@ -447,20 +445,13 @@ impl RaftRoleState for LearnerState { return Ok(()); } - RaftEvent::StreamSnapshot(request, sender) => { - debug!(?request, "Learner::RaftEvent::StreamSnapshot"); - sender - .send(Err(Status::permission_denied( - "Learner should not receive StreamSnapshot event.", - ))) - .map_err(|e| { - let error_str = format!("{e:?}"); - error!("Failed to send: {}", error_str); - NetworkError::SingalSendFailed(error_str) - })?; + RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx) => { + debug!("Learner::RaftEvent::StreamSnapshot"); + warn!("Candidate should not receive StreamSnapshot event."); return Ok(()); } + RaftEvent::PromoteReadyLearners => { return Err(ConsensusError::RoleViolation { current_role: "Learner", @@ -597,7 +588,7 @@ impl RaftRoleState for LearnerState { let (ack_tx, ack_rx) = mpsc::channel(32); // Get snapshot stream from leader (with ACK feedback) - let snapshot_stream = ctx + let snapshot_chunk_receiver = ctx .transport() .request_snapshot_from_leader( leader_id, @@ -613,7 +604,7 @@ impl RaftRoleState for LearnerState { .state_machine_handler .apply_snapshot_stream_from_leader( self.current_term(), - snapshot_stream, + snapshot_chunk_receiver, ack_tx, &ctx.node_config.raft.snapshot, ) diff --git a/d-engine-core/src/raft_role/learner_state_test.rs b/d-engine-core/src/raft_role/learner_state_test.rs index 8531c336..dd1247c6 100644 --- a/d-engine-core/src/raft_role/learner_state_test.rs +++ b/d-engine-core/src/raft_role/learner_state_test.rs @@ -16,7 +16,6 @@ use crate::client::WriteOperation; use crate::config::ReadConsistencyPolicy; use crate::raft_role::learner_state::LearnerState; use crate::raft_role::role_state::RaftRoleState; -use crate::test_utils::create_test_snapshot_stream; use crate::test_utils::mock::MockTypeConfig; use crate::test_utils::mock::mock_raft_context; use crate::test_utils::mock::mock_raft_context_with_temp; @@ -1790,14 +1789,15 @@ async fn test_learner_install_snapshot_reports_success_after_all_chunks_applied( let mut state = LearnerState::::new(1, context.node_config.clone()); state.update_current_term(2); - let chunks: Vec = vec![SnapshotChunk::default()]; - let stream = create_test_snapshot_stream(chunks); let (resp_tx, mut resp_rx) = MaybeCloneOneshot::new(); let (role_tx, _role_rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::channel(32); + tx.send(SnapshotChunk::default()).await.unwrap(); + drop(tx); state .handle_raft_event( - RaftEvent::InstallSnapshotChunk(Box::new(stream), resp_tx), + RaftEvent::InstallSnapshotChunk(rx, resp_tx), &context, role_tx, ) @@ -1863,15 +1863,16 @@ async fn test_learner_install_snapshot_does_not_report_success_on_mid_chunk_fail let mut state = LearnerState::::new(1, context.node_config.clone()); state.update_current_term(2); - let chunks: Vec = vec![SnapshotChunk::default()]; - let stream = create_test_snapshot_stream(chunks); let (resp_tx, mut resp_rx) = MaybeCloneOneshot::new(); let (role_tx, _role_rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::channel(32); + tx.send(SnapshotChunk::default()).await.unwrap(); + drop(tx); // handle_raft_event returns Err (apply failed) — that is expected let _ = state .handle_raft_event( - RaftEvent::InstallSnapshotChunk(Box::new(stream), resp_tx), + RaftEvent::InstallSnapshotChunk(rx, resp_tx), &context, role_tx, ) @@ -1933,14 +1934,16 @@ async fn test_learner_install_snapshot_reports_failure_when_apply_fails_after_tr let mut state = LearnerState::::new(1, context.node_config.clone()); state.update_current_term(2); - let stream = create_test_snapshot_stream(vec![SnapshotChunk::default()]); let (resp_tx, mut resp_rx) = MaybeCloneOneshot::new(); let (role_tx, _role_rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::channel(32); + tx.send(SnapshotChunk::default()).await.unwrap(); + drop(tx); // Learner propagates apply errors — ignore the return value here let _ = state .handle_raft_event( - RaftEvent::InstallSnapshotChunk(Box::new(stream), resp_tx), + RaftEvent::InstallSnapshotChunk(rx, resp_tx), &context, role_tx, ) diff --git a/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs b/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs index da36a56c..e69192a3 100644 --- a/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs +++ b/d-engine-core/src/state_machine_handler/default_state_machine_handler.rs @@ -34,9 +34,7 @@ use tokio::sync::RwLock; use tokio::sync::mpsc; use tokio::time::Instant; use tokio::time::timeout; -use tokio_stream::StreamExt; use tokio_tar::Archive; -use tonic::Streaming; use tracing::debug; use tracing::error; use tracing::info; @@ -326,15 +324,16 @@ where async fn apply_snapshot_stream_from_leader( &self, current_term: u64, - stream: Box>, + stream_chunk_receiver: mpsc::Receiver, ack_tx: mpsc::Sender, // ACK sender config: &SnapshotConfig, ) -> Result<()> { let _timer = ScopedTimer::new("receive_snapshot_stream_from_leader"); info!(?ack_tx, %current_term, "receive_snapshot_stream_from_leader"); - let (final_metadata, snapshot_path) = - self.process_snapshot_stream(stream, ack_tx.clone(), config).await?; + let (final_metadata, snapshot_path) = self + .process_snapshot_stream(stream_chunk_receiver, ack_tx.clone(), config) + .await?; debug!( ?final_metadata, @@ -898,10 +897,10 @@ where } /// Helper function to process snapshot stream - #[instrument(skip(self, stream))] + #[instrument(skip(self, stream_chunk_receiver))] async fn process_snapshot_stream( &self, - mut stream: Box>, + mut stream_chunk_receiver: mpsc::Receiver, ack_tx: mpsc::Sender, config: &SnapshotConfig, ) -> Result<(SnapshotMetadata, PathBuf)> { @@ -914,28 +913,12 @@ where let mut count = 0; loop { - let chunk = match timeout(chunk_timeout, stream.next()).await { - Ok(Some(Ok(chunk))) => { + let chunk = match timeout(chunk_timeout, stream_chunk_receiver.recv()).await { + Ok(Some(chunk)) => { debug!("receive new chunk."); last_received = Instant::now(); chunk } - Ok(Some(Err(e))) => { - // Send stream error ACK - ack_tx - .send(SnapshotAck { - seq: 0, // Best effort - status: ChunkStatus::Failed.into(), - next_requested: 0, - }) - .await - .map_err(|e| { - SnapshotError::OperationFailed(format!("Failed to send ACK: {e}")) - })?; - return Err( - SnapshotError::OperationFailed(format!("Stream error: {e:?}")).into(), - ); - } Ok(None) => { debug!("no more chunks available..."); diff --git a/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs b/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs index 87862008..93c879d4 100644 --- a/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs +++ b/d-engine-core/src/state_machine_handler/default_state_machine_handler_test.rs @@ -34,7 +34,6 @@ use crate::SnapshotError; use crate::StorageError; use crate::test_utils::create_test_chunk; use crate::test_utils::create_test_compressed_snapshot; -use crate::test_utils::create_test_snapshot_stream; use crate::test_utils::snapshot_config; // Case 1: normal update @@ -530,7 +529,6 @@ async fn test_apply_snapshot_stream_from_leader_case2() { // 3. Fake install snapshot request stream let total_chunks = 1; // Create compressed chunk data - let mut chunks: Vec = vec![]; let metadata = SnapshotMetadata { last_included: Some(LogId { index: 2, term: 1 }), checksum: Bytes::from(vec![2; 32]), @@ -563,19 +561,17 @@ async fn test_apply_snapshot_stream_from_leader_case2() { data: Bytes::from(compressed_data.clone()), chunk_checksum: Bytes::from(crc32fast::hash(&compressed_data).to_be_bytes().to_vec()), }; - chunks.push(chunk); - let streaming_request = create_test_snapshot_stream(chunks); let (ack_tx, mut ack_rx) = mpsc::channel::(1); + let (tx, rx) = mpsc::channel(32); + tx.send(chunk).await.unwrap(); + drop(tx); + // Spawn the handler in a separate task to prevent deadlock let handler_task = tokio::spawn({ let config = snapshot_config(temp_path.to_path_buf()); - async move { - handler - .apply_snapshot_stream_from_leader(1, Box::new(streaming_request), ack_tx, &config) - .await - } + async move { handler.apply_snapshot_stream_from_leader(1, rx, ack_tx, &config).await } }); // Verify intermediate response @@ -604,15 +600,13 @@ async fn test_apply_snapshot_stream_from_leader_case3() { // Create ACK channel let (ack_tx, mut ack_rx) = mpsc::channel::(1); - let stream = create_test_snapshot_stream(vec![bad_chunk]); + let (tx, rx) = mpsc::channel(32); + tx.send(bad_chunk).await.unwrap(); + drop(tx); let handler_task = tokio::spawn({ let config = snapshot_config(temp_path.to_path_buf()); - async move { - handler - .apply_snapshot_stream_from_leader(TEST_TERM, Box::new(stream), ack_tx, &config) - .await - } + async move { handler.apply_snapshot_stream_from_leader(TEST_TERM, rx, ack_tx, &config).await } }); let ack = ack_rx.recv().await.unwrap(); @@ -641,15 +635,15 @@ async fn test_apply_snapshot_stream_from_leader_case4() { ]; let (ack_tx, mut ack_rx) = mpsc::channel::(1); - let stream = create_test_snapshot_stream(chunks); + let (tx, rx) = mpsc::channel(32); + for chunk in chunks { + tx.send(chunk).await.unwrap(); + } + drop(tx); let handler_task = tokio::spawn({ let config = snapshot_config(temp_path.to_path_buf()); - async move { - handler - .apply_snapshot_stream_from_leader(TEST_TERM, Box::new(stream), ack_tx, &config) - .await - } + async move { handler.apply_snapshot_stream_from_leader(TEST_TERM, rx, ack_tx, &config).await } }); let ack = ack_rx.recv().await.unwrap(); @@ -683,16 +677,17 @@ async fn test_apply_snapshot_stream_from_leader_case5() { TEST_LEADER_ID, 2, )]; - let stream = create_test_snapshot_stream(chunks); + + let (tx, rx) = mpsc::channel(32); + for chunk in chunks { + tx.send(chunk).await.unwrap(); + } + drop(tx); let (ack_tx, mut ack_rx) = mpsc::channel::(1); let handler_task = tokio::spawn({ let config = snapshot_config(temp_path.to_path_buf()); - async move { - handler - .apply_snapshot_stream_from_leader(TEST_TERM, Box::new(stream), ack_tx, &config) - .await - } + async move { handler.apply_snapshot_stream_from_leader(TEST_TERM, rx, ack_tx, &config).await } }); let ack = ack_rx.recv().await.unwrap(); @@ -719,15 +714,13 @@ async fn test_apply_snapshot_stream_from_leader_case6() { invalid_chunk.metadata = None; let (ack_tx, mut ack_rx) = mpsc::channel::(1); - let stream = create_test_snapshot_stream(vec![invalid_chunk]); + let (tx, rx) = mpsc::channel(32); + tx.send(invalid_chunk).await.unwrap(); + drop(tx); let handler_task = tokio::spawn({ let config = snapshot_config(temp_path.to_path_buf()); - async move { - handler - .apply_snapshot_stream_from_leader(TEST_TERM, Box::new(stream), ack_tx, &config) - .await - } + async move { handler.apply_snapshot_stream_from_leader(TEST_TERM, rx, ack_tx, &config).await } }); let ack = ack_rx.recv().await.unwrap(); @@ -790,17 +783,17 @@ async fn test_apply_snapshot_stream_from_leader_case7() { chunks.push(chunk); } - let streaming_request = create_test_snapshot_stream(chunks); let (ack_tx, mut ack_rx) = mpsc::channel::(1); + let (tx, rx) = mpsc::channel(32); + for chunk in chunks { + tx.send(chunk).await.unwrap(); + } + drop(tx); // Spawn the handler in a separate task let handler_task = tokio::spawn({ let config = snapshot_config(temp_path.to_path_buf()); - async move { - handler - .apply_snapshot_stream_from_leader(1, Box::new(streaming_request), ack_tx, &config) - .await - } + async move { handler.apply_snapshot_stream_from_leader(1, rx, ack_tx, &config).await } }); // Verify intermediate ACKs @@ -1702,16 +1695,13 @@ async fn test_apply_snapshot_stream_from_leader_decompresses_before_apply() { chunk_checksum: Bytes::from(chunk_checksum), }; - let streaming_request = create_test_snapshot_stream(vec![chunk]); let (ack_tx, mut ack_rx) = mpsc::channel::(1); - + let (tx, rx) = mpsc::channel(32); + tx.send(chunk).await.unwrap(); + drop(tx); let handler_task = tokio::spawn({ let config = snapshot_config(temp_path.to_path_buf()); - async move { - handler - .apply_snapshot_stream_from_leader(1, Box::new(streaming_request), ack_tx, &config) - .await - } + async move { handler.apply_snapshot_stream_from_leader(1, rx, ack_tx, &config).await } }); // Verify ACK diff --git a/d-engine-core/src/state_machine_handler/mod.rs b/d-engine-core/src/state_machine_handler/mod.rs index bbf6ac66..eff03f13 100644 --- a/d-engine-core/src/state_machine_handler/mod.rs +++ b/d-engine-core/src/state_machine_handler/mod.rs @@ -112,7 +112,7 @@ where async fn apply_snapshot_stream_from_leader( &self, current_term: u64, - stream: Box>, + stream_chunk_receiver: tokio::sync::mpsc::Receiver, ack_tx: tokio::sync::mpsc::Sender, config: &crate::SnapshotConfig, ) -> Result<()>; diff --git a/d-engine-core/src/utils/stream.rs b/d-engine-core/src/utils/stream.rs index 2a8bce53..aad9db19 100644 --- a/d-engine-core/src/utils/stream.rs +++ b/d-engine-core/src/utils/stream.rs @@ -1,19 +1,7 @@ -use std::marker::PhantomData; -use std::sync::Arc; - use bytes::BufMut; -use bytes::BytesMut; -use futures::TryStreamExt; -use http_body::Frame; -use http_body_util::BodyExt; -use http_body_util::StreamBody; -use prost::Message; -use tokio::sync::mpsc; -use tokio_stream::StreamExt; -use tokio_stream::wrappers::ReceiverStream; +use std::marker::PhantomData; use tonic::Code; use tonic::Status; -use tonic::Streaming; // Adjust path as needed /// Generic gRPC stream decoder for any protobuf message /// @@ -79,56 +67,3 @@ where tonic::codec::BufferSettings::new(4 * 1024 * 1024, 4 * 1024 * 1025) } } - -/// Converts a receiver channel into a properly encoded `tonic::Streaming` -/// -/// This handles: -/// 1. Proper gRPC frame encoding -/// 2. Error conversion -/// 3. Backpressure through bounded channel -/// 4. Efficient memory usage -pub(crate) fn create_production_snapshot_stream( - rx: mpsc::Receiver, Status>>, - max_message_size: usize, -) -> Streaming -where - T: Message + Default + 'static, -{ - // Create byte stream with proper gRPC framing - let byte_stream = ReceiverStream::new(rx).map(|res| { - match res { - Ok(arc_chunk) => { - let chunk: &T = &arc_chunk; - - // Encode the T to bytes - let mut buf = Vec::new(); - chunk.encode(&mut buf).map_err(|e| { - Status::new(Code::Internal, format!("Snapshot encoding failed: {e}")) - })?; - - // Create gRPC frame with header - let mut frame = BytesMut::with_capacity(5 + buf.len()); - frame.put_u8(0); // No compression - frame.put_u32(buf.len() as u32); // Message length - frame.extend_from_slice(&buf); - - Ok(frame.freeze()) - } - Err(e) => Err(e), - } - }); - - // Create stream body with proper boxing - let body = StreamBody::new(byte_stream.map_ok(Frame::data).map_err(|e: Status| e)); - - // Create streaming with appropriate codec - Streaming::new_request( - GrpcStreamDecoder:: { - _marker: PhantomData, - }, - body.boxed_unsync(), - None, - Some(max_message_size), - // Some(1024 * 1024 * 1024), // 1GB max message size - ) -} diff --git a/d-engine-core/src/utils/stream_test.rs b/d-engine-core/src/utils/stream_test.rs index 19250977..e16a02e2 100644 --- a/d-engine-core/src/utils/stream_test.rs +++ b/d-engine-core/src/utils/stream_test.rs @@ -1,50 +1,6 @@ -use std::sync::Arc; - -use prost::Message; -use tokio::sync::mpsc; -use tokio_stream::StreamExt; - -use crate::stream::create_production_snapshot_stream; use crate::stream::encode_varint; use crate::stream::encoded_len_varint; -// Simple test message -#[derive(Clone, PartialEq, Message)] -pub struct TestMessage { - #[prost(string, tag = "1")] - pub data: String, -} - -#[tokio::test] -async fn test_create_production_snapshot_stream() { - let (tx, rx) = mpsc::channel(10); - let max_message_size = 1024 * 1024; // 1MB - - // Create the stream - let mut stream = create_production_snapshot_stream::(rx, max_message_size); - - // Send a test message - let test_msg = TestMessage { - data: "test".to_string(), - }; - tx.send(Ok(Arc::new(test_msg))).await.unwrap(); - - // Close the channel - drop(tx); - - // Receive from the stream - let received = stream.next().await; - assert!(received.is_some()); - - let result = received.unwrap(); - assert!(result.is_ok()); - let msg = result.unwrap(); - assert_eq!(msg.data, "test"); - - // Stream should be done now - assert!(stream.next().await.is_none()); -} - #[test] fn test_encoded_len_varint() { // Test small values diff --git a/d-engine-server/src/network/grpc/grpc_raft_service.rs b/d-engine-server/src/network/grpc/grpc_raft_service.rs index b34cad59..17039ace 100644 --- a/d-engine-server/src/network/grpc/grpc_raft_service.rs +++ b/d-engine-server/src/network/grpc/grpc_raft_service.rs @@ -2,15 +2,12 @@ //! and client requests. Implements core Raft protocol logic for leader election, //! log replication, and cluster configuration management. -use std::future::Future; -use std::pin::Pin; -use std::time::Duration; - +use crate::Node; +use crate::proto_convert; use d_engine_core::MaybeCloneOneshot; use d_engine_core::MaybeCloneOneshotReceiver; use d_engine_core::RaftEvent; use d_engine_core::RaftOneshot; -use d_engine_core::StreamResponseSender; use d_engine_core::TypeConfig; #[cfg(feature = "watch")] use d_engine_core::WatchError; @@ -45,9 +42,14 @@ use d_engine_proto::server::storage::SnapshotResponse; use d_engine_proto::server::storage::snapshot_service_server::SnapshotService; use futures::Stream; use futures::StreamExt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::Duration; use tokio::select; use tokio::sync::mpsc; use tokio::time::timeout; +use tokio_stream::wrappers::ReceiverStream; use tokio_util::sync::CancellationToken; use tonic::Request; use tonic::Response; @@ -58,9 +60,6 @@ use tracing::error; use tracing::info; use tracing::warn; -use crate::Node; -use crate::proto_convert; - #[tonic::async_trait] impl RaftElectionService for Node where @@ -223,7 +222,7 @@ impl SnapshotService for Node where T: TypeConfig, { - type StreamSnapshotStream = tonic::Streaming; + type StreamSnapshotStream = Pin> + Send>>; async fn stream_snapshot( &self, @@ -234,24 +233,34 @@ where return Err(Status::unavailable("Service is not ready")); } - let (resp_tx, resp_rx) = StreamResponseSender::new(); + // Bridge incoming ACK stream: tonic → mpsc + let mut ack_stream = request.into_inner(); + let (ack_tx, ack_rx) = mpsc::channel::(32); + tokio::spawn(async move { + while let Some(ack) = ack_stream.next().await { + match ack { + Ok(a) => { + if ack_tx.send(a).await.is_err() { + break; + } + } + Err(_) => break, + } + } + }); + + // Server creates chunk channel, passes tx to core, keeps rx for gRPC response + let (chunk_tx, chunk_rx) = mpsc::channel::>(32); self.event_tx - .send(RaftEvent::StreamSnapshot( - Box::new(request.into_inner()), - resp_tx, - )) + .send(RaftEvent::StreamSnapshot(ack_rx, chunk_tx)) .await .map_err(|_| Status::internal("Event channel closed"))?; - let timeout_duration = Duration::from_millis(self.node_config.raft.snapshot_rpc_timeout_ms); - - handle_rpc_timeout( - async { resp_rx.await.map_err(|_| Status::internal("Response channel closed")) }, - timeout_duration, - "stream_snapshot", - ) - .await + // Return chunk stream directly — no waiting needed + let response_stream = + ReceiverStream::new(chunk_rx).map(|arc_chunk| Ok((*arc_chunk).clone())); + Ok(Response::new(Box::pin(response_stream))) } async fn install_snapshot( @@ -265,11 +274,23 @@ where let (resp_tx, resp_rx) = MaybeCloneOneshot::new(); + let mut tonic_stream = request.into_inner(); + let (tx, rx) = mpsc::channel(32); + tokio::spawn(async move { + while let Some(chunk) = tonic_stream.next().await { + match chunk { + Ok(c) => { + if tx.send(c).await.is_err() { + break; + } + } + Err(_) => break, + } + } + }); + self.event_tx - .send(RaftEvent::InstallSnapshotChunk( - Box::new(request.into_inner()), - resp_tx, - )) + .send(RaftEvent::InstallSnapshotChunk(rx, resp_tx)) .await .map_err(|_| Status::internal("Event channel closed"))?; diff --git a/d-engine-server/src/network/grpc/grpc_transport.rs b/d-engine-server/src/network/grpc/grpc_transport.rs index 493e935f..162a9489 100644 --- a/d-engine-server/src/network/grpc/grpc_transport.rs +++ b/d-engine-server/src/network/grpc/grpc_transport.rs @@ -1,10 +1,6 @@ //! Centerialized all RPC client operations will make unit test eaiser. //! We also want to refactor all the APIs based its similar parttern. -use std::collections::HashSet; -use std::marker::PhantomData; -use std::sync::Arc; - use async_trait::async_trait; use d_engine_core::AppendResult; use d_engine_core::BackgroundSnapshotTransfer; @@ -44,6 +40,9 @@ use dashmap::DashMap; use futures::FutureExt; use futures::StreamExt; use futures::stream::FuturesUnordered; +use std::collections::HashSet; +use std::marker::PhantomData; +use std::sync::Arc; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::task; @@ -548,7 +547,7 @@ where ack_rx: mpsc::Receiver, _retry: &InstallSnapshotBackoffPolicy, membership: Arc>, - ) -> Result>> { + ) -> Result> { debug!("Fetching snapshot from leader {}", leader_id); // Get bulk connection channel @@ -569,7 +568,26 @@ where .await .map_err(|e| NetworkError::TonicStatusError(Box::new(e)))?; - Ok(Box::new(response.into_inner())) + let mut tonic_stream = response.into_inner(); + let (tx, rx) = mpsc::channel(32); + + tokio::spawn(async move { + while let Some(chunk) = tonic_stream.next().await { + match chunk { + Ok(c) => { + if tx.send(c).await.is_err() { + break; //receive stops + } + } + Err(e) => { + //network error, drop tx, rx will receive None + error!("{:?}", e); + break; + } + } + } + }); + Ok(rx) } } From 008c8a4af708b20d32bc3ee86e16b29acbbd4476 Mon Sep 17 00:00:00 2001 From: Joshua Chi <112539+JoshuaChi@users.noreply.github.com> Date: Sun, 21 Jun 2026 16:50:35 +0800 Subject: [PATCH 2/3] fix #409: StreamSnapshot explicit rejection + remove yield_now() perf regression MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- d-engine-core/src/event.rs | 18 +-- .../network/background_snapshot_transfer.rs | 10 +- .../background_snapshot_transfer_test.rs | 55 +++++++ d-engine-core/src/raft.rs | 14 -- .../src/raft_role/candidate_state.rs | 9 +- d-engine-core/src/raft_role/follower_state.rs | 5 +- d-engine-core/src/raft_role/leader_state.rs | 18 ++- .../leader_state_test/snapshot_test.rs | 150 ++++++++++++++++++ d-engine-core/src/raft_role/learner_state.rs | 11 +- .../src/network/grpc/grpc_raft_service.rs | 13 +- 10 files changed, 272 insertions(+), 31 deletions(-) diff --git a/d-engine-core/src/event.rs b/d-engine-core/src/event.rs index 2bbffbae..31d12e33 100644 --- a/d-engine-core/src/event.rs +++ b/d-engine-core/src/event.rs @@ -1,6 +1,6 @@ -use std::path::PathBuf; - -use crate::client::{ClientReadRequest, ClientResponse, ClientWriteRequest}; +use crate::client::ClientReadRequest; +use crate::client::ClientResponse; +use crate::client::ClientWriteRequest; use d_engine_proto::common::LogId; use d_engine_proto::server::cluster::ClusterConfChangeRequest; use d_engine_proto::server::cluster::ClusterConfUpdateResponse; @@ -18,7 +18,6 @@ use d_engine_proto::server::storage::SnapshotAck; use d_engine_proto::server::storage::SnapshotChunk; use d_engine_proto::server::storage::SnapshotMetadata; use d_engine_proto::server::storage::SnapshotResponse; -use tokio::sync::mpsc; use tonic::Status; use crate::ApplyResult; @@ -151,14 +150,15 @@ pub enum RaftEvent { // Response snapshot stream from Leader InstallSnapshotChunk( - mpsc::Receiver, + tokio::sync::mpsc::Receiver, MaybeCloneOneshotSender>, ), // Request snapshot stream from Leader StreamSnapshot( - mpsc::Receiver, - mpsc::Sender>, + tokio::sync::mpsc::Receiver, + tokio::sync::mpsc::Sender>, + tokio::sync::oneshot::Sender>, // startup confirmation ), JoinCluster( @@ -175,7 +175,7 @@ pub enum RaftEvent { LogPurgeCompleted(LogId), - SnapshotCreated(Result<(SnapshotMetadata, PathBuf)>), + SnapshotCreated(Result<(SnapshotMetadata, std::path::PathBuf)>), // Lightweight promotion trigger PromoteReadyLearners, @@ -248,7 +248,7 @@ pub(crate) fn raft_event_to_test_event(event: &RaftEvent) -> TestEvent { RaftEvent::ClusterConfUpdate(req, _) => TestEvent::ClusterConfUpdate(req.clone()), RaftEvent::AppendEntries(req, _) => TestEvent::AppendEntries(req.clone()), RaftEvent::InstallSnapshotChunk(_, _) => TestEvent::InstallSnapshotChunk, - RaftEvent::StreamSnapshot(_, _) => TestEvent::StreamSnapshot, + RaftEvent::StreamSnapshot(_, _, _) => TestEvent::StreamSnapshot, RaftEvent::JoinCluster(req, _) => TestEvent::JoinCluster(req.clone()), RaftEvent::DiscoverLeader(req, _) => TestEvent::DiscoverLeader(req.clone()), RaftEvent::CreateSnapshotEvent => TestEvent::CreateSnapshotEvent, diff --git a/d-engine-core/src/network/background_snapshot_transfer.rs b/d-engine-core/src/network/background_snapshot_transfer.rs index 7bad0b2a..d70aac23 100644 --- a/d-engine-core/src/network/background_snapshot_transfer.rs +++ b/d-engine-core/src/network/background_snapshot_transfer.rs @@ -215,7 +215,15 @@ where &mut retry_counts, &config, )?, - None => break, // ACK stream closed + None => { + if let Some(total) = total_chunks + && next_seq >= total + && pending_acks.is_empty() + { + break; + } + return Err(SnapshotError::TransferFailed.into()); + } } }, diff --git a/d-engine-core/src/network/background_snapshot_transfer_test.rs b/d-engine-core/src/network/background_snapshot_transfer_test.rs index 78c8e367..b7dc8ae9 100644 --- a/d-engine-core/src/network/background_snapshot_transfer_test.rs +++ b/d-engine-core/src/network/background_snapshot_transfer_test.rs @@ -666,6 +666,61 @@ mod run_pull_transfer_test { ); } + // Test: ACK channel closes before all chunks are acknowledged — transfer must fail. + // + // This test guards the fix in the `None` arm of `ack_rx.recv()`. + // Before the fix, `None => break` exited the loop and returned Ok(()) even + // with unacknowledged chunks still in pending_acks. + // + // Setup: 2-chunk stream. Follower ACKs seq=0 (Accepted), receives seq=1, + // then drops the ACK sender without ACKing seq=1 (simulates disconnect). + // Expected: run_pull_transfer returns Err(TransferFailed). + // Regression: revert `None` arm to just `break` — this test must then FAIL + // because the transfer would incorrectly return Ok(()). + #[tokio::test] + async fn test_pull_transfer_fails_on_early_ack_channel_closure() { + let config = pull_config(); + let (ack_tx, ack_rx) = mpsc::channel::(32); + let (chunk_tx, mut chunk_rx) = mpsc::channel::>(32); + + let data_stream = stream::iter(vec![Ok(make_chunk(0, 2)), Ok(make_chunk(1, 2))]).boxed(); + + let transfer = tokio::spawn(async move { + BackgroundSnapshotTransfer::::run_pull_transfer( + ack_rx, + chunk_tx, + data_stream, + config, + ) + .await + }); + + // Receive seq=0 and ACK it — pending_acks becomes {1} after seq=1 is sent. + let chunk0 = chunk_rx.recv().await.expect("should receive chunk 0"); + assert_eq!(chunk0.seq, 0); + ack_tx + .send(make_ack(0, ChunkStatus::Accepted)) + .await + .expect("ack channel should be open"); + + // Receive seq=1 but do NOT ACK it — then close the ACK channel. + // At this point pending_acks = {1}, so early closure must be an error. + let chunk1 = chunk_rx.recv().await.expect("should receive chunk 1"); + assert_eq!(chunk1.seq, 1); + drop(ack_tx); + + let result = transfer.await.expect("task should not panic"); + assert!( + matches!( + result, + Err(crate::Error::Consensus(crate::ConsensusError::Snapshot( + SnapshotError::TransferFailed + ))) + ), + "expected TransferFailed when ack channel closes with pending chunks, got: {result:?}" + ); + } + // Test: Chunks arriving out of sequence cause an immediate protocol error. // // Setup: data_stream yields seq=0 then seq=2 (skipping seq=1). diff --git a/d-engine-core/src/raft.rs b/d-engine-core/src/raft.rs index 6af9c249..ad80b95e 100644 --- a/d-engine-core/src/raft.rs +++ b/d-engine-core/src/raft.rs @@ -296,8 +296,6 @@ where debug!(%self.node_id, ?role_event, "receive role event"); self.buffered_role_event.push_back(role_event); self.drain_role_events().await?; - - } // P3: Client commands — push first, drain rest after select @@ -312,23 +310,11 @@ where trace!(%self.node_id, ?raft_event, "receive raft event"); self.buffered_raft_event.push_back(raft_event); self.drain_raft_events().await?; - - - - } } - // After any arm fires: drain all channels in order. - // role_rx first: processes ACKs/commits and sends responses to clients, - // naturally yielding at .await points so client tasks can enqueue their - // next writes into cmd_rx before drain_client_cmds runs. - tokio::task::yield_now().await; self.process_role_events().await?; - // Yield after role events so woken client tasks can enqueue their next - // writes into cmd_rx before drain_client_cmds runs. - tokio::task::yield_now().await; self.process_client_cmds().await?; self.process_raft_events().await?; } diff --git a/d-engine-core/src/raft_role/candidate_state.rs b/d-engine-core/src/raft_role/candidate_state.rs index b290a0e5..5e44a209 100644 --- a/d-engine-core/src/raft_role/candidate_state.rs +++ b/d-engine-core/src/raft_role/candidate_state.rs @@ -464,9 +464,16 @@ impl RaftRoleState for CandidateState { return Ok(()); } - RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx) => { + RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx, startup_tx) => { debug!("Candidate::RaftEvent::StreamSnapshot"); warn!("Candidate 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(()); } diff --git a/d-engine-core/src/raft_role/follower_state.rs b/d-engine-core/src/raft_role/follower_state.rs index 4533d8e1..f00e3a12 100644 --- a/d-engine-core/src/raft_role/follower_state.rs +++ b/d-engine-core/src/raft_role/follower_state.rs @@ -500,9 +500,10 @@ impl RaftRoleState for FollowerState { return Ok(()); } - RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx) => { + RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx, startup_tx) => { + let _ = startup_tx.send(Err(Status::failed_precondition("Not the leader"))); debug!("Follower::RaftEvent::StreamSnapshot"); - warn!("Candidate should not receive StreamSnapshot event."); + warn!("Follower should not receive StreamSnapshot event."); return Ok(()); } diff --git a/d-engine-core/src/raft_role/leader_state.rs b/d-engine-core/src/raft_role/leader_state.rs index fa35359c..7de086da 100644 --- a/d-engine-core/src/raft_role/leader_state.rs +++ b/d-engine-core/src/raft_role/leader_state.rs @@ -1386,11 +1386,19 @@ impl RaftRoleState for LeaderState { panic!("{}", msg); } } - RaftEvent::StreamSnapshot(ack_rx, chunk_tx) => { + RaftEvent::StreamSnapshot(ack_rx, chunk_tx, startup_tx) => { debug!("Leader::RaftEvent::StreamSnapshot"); // Get the latest snapshot metadata if let Some(metadata) = ctx.state_machine().snapshot_metadata() { + // confirm: transfer starting + if let Err(e) = startup_tx.send(Ok(())) { + error!( + ?e, + "StreamSnapshot startup_tx send failed: gRPC receiver already dropped" + ); + } + // Spawn background transfer task let state_machine_handler = ctx.state_machine_handler().clone(); let config = ctx.node_config.raft.snapshot.clone(); @@ -1410,6 +1418,14 @@ impl RaftRoleState for LeaderState { error!("StreamSnapshot failed: {:?}", e); } }); + } else if let Err(e) = + startup_tx.send(Err(Status::not_found("No snapshot available"))) + { + error!( + ?e, + "StreamSnapshot startup_tx send failed: gRPC receiver already dropped" + ); + // chunk_tx dropped } } RaftEvent::PromoteReadyLearners => { diff --git a/d-engine-core/src/raft_role/leader_state_test/snapshot_test.rs b/d-engine-core/src/raft_role/leader_state_test/snapshot_test.rs index 18523378..58031841 100644 --- a/d-engine-core/src/raft_role/leader_state_test/snapshot_test.rs +++ b/d-engine-core/src/raft_role/leader_state_test/snapshot_test.rs @@ -220,3 +220,153 @@ async fn test_handle_log_purge_completed() { .unwrap(); assert_eq!(state.last_purged_index, Some(LogId { term: 1, index: 50 })); } + +// ============================================================================ +// StreamSnapshot Startup Rejection Tests +// ============================================================================ + +/// Leader rejects StreamSnapshot when no snapshot metadata is available. +/// +/// # Given +/// - Leader with no snapshot created yet (snapshot_metadata() → None) +/// +/// # When +/// - StreamSnapshot event is received +/// +/// # Then +/// - startup_rx receives Err(Status::not_found) +/// - chunk_rx is closed immediately (no chunks sent) +#[tokio::test] +async fn test_stream_snapshot_rejects_when_no_metadata() { + let (_graceful_tx, graceful_rx) = watch::channel(()); + // Default mock: snapshot_metadata() returns None + let context = MockBuilder::new(graceful_rx).build_context(); + let mut state = LeaderState::::new(1, context.node_config()); + + let (ack_tx, ack_rx) = mpsc::channel::(4); + let (chunk_tx, mut chunk_rx) = + mpsc::channel::>(4); + let (startup_tx, startup_rx) = tokio::sync::oneshot::channel::>(); + + let (role_tx, _role_rx) = mpsc::unbounded_channel(); + state + .handle_raft_event( + RaftEvent::StreamSnapshot(ack_rx, chunk_tx, startup_tx), + &context, + role_tx, + ) + .await + .expect("handler must not return Err"); + + // startup_rx must carry a rejection, not Ok + let result = startup_rx.await.expect("startup_tx must be sent"); + assert!(result.is_err(), "expected Err from startup channel, got Ok"); + assert_eq!(result.unwrap_err().code(), tonic::Code::NotFound); + + // chunk channel must be closed — no chunks were sent + drop(ack_tx); // silence unused warning + assert!( + chunk_rx.recv().await.is_none(), + "expected chunk_rx to be closed" + ); +} + +/// Leader confirms startup and begins transfer when snapshot metadata exists. +/// +/// # Given +/// - Leader with valid snapshot metadata +/// - State machine handler returns a single-chunk data stream +/// +/// # When +/// - StreamSnapshot event is received +/// +/// # Then +/// - startup_rx receives Ok(()) +/// - chunk_rx eventually delivers the chunk (background task runs) +#[tokio::test] +async fn test_stream_snapshot_confirms_startup_when_metadata_exists() { + use bytes::Bytes; + use futures::StreamExt; + use futures::stream; + + let (_graceful_tx, graceful_rx) = watch::channel(()); + + // Build state machine from scratch so snapshot_metadata returns Some. + // (Can't override expectations added by mock_state_machine() — mockall + // uses the first matching expectation, so the None default would win.) + let mut sm = crate::MockStateMachine::new(); + sm.expect_start().returning(|| Ok(())); + sm.expect_stop().returning(|| Ok(())); + sm.expect_is_running().returning(|| true); + sm.expect_get().returning(|_| Ok(None)); + sm.expect_entry_term().returning(|_| None); + sm.expect_apply_chunk().returning(|_| Ok(vec![])); + sm.expect_len().returning(|| 0); + sm.expect_update_last_applied().returning(|_| ()); + sm.expect_last_applied().return_const(LogId::default()); + sm.expect_persist_last_applied().returning(|_| Ok(())); + sm.expect_update_last_snapshot_metadata().returning(|_| Ok(())); + sm.expect_snapshot_metadata().returning(|| { + Some(SnapshotMetadata { + last_included: Some(LogId { term: 1, index: 10 }), + ..Default::default() + }) + }); + sm.expect_persist_last_snapshot_metadata().returning(|_| Ok(())); + sm.expect_apply_snapshot_from_file().returning(|_, _| Ok(())); + sm.expect_generate_snapshot_data() + .returning(|_, _| Ok(bytes::Bytes::copy_from_slice(&[0u8; 32]))); + sm.expect_save_hard_state().returning(|| Ok(())); + sm.expect_flush().returning(|| Ok(())); + + // State machine handler returns a minimal 1-chunk stream + let mut smh = crate::test_utils::mock::mock_state_machine_handler(); + smh.expect_load_snapshot_data().returning(|_| { + let chunk = d_engine_proto::server::storage::SnapshotChunk { + seq: 0, + total_chunks: 1, + data: Bytes::from_static(b"hello"), + metadata: Some(SnapshotMetadata::default()), + ..Default::default() + }; + Ok(stream::iter(vec![Ok(chunk)]).boxed()) + }); + + let context = MockBuilder::new(graceful_rx) + .with_state_machine(sm) + .with_state_machine_handler(smh) + .build_context(); + + let mut state = LeaderState::::new(1, context.node_config()); + + let (_ack_tx, ack_rx) = mpsc::channel::(4); + let (chunk_tx, mut chunk_rx) = + mpsc::channel::>(4); + let (startup_tx, startup_rx) = tokio::sync::oneshot::channel::>(); + + let (role_tx, _role_rx) = mpsc::unbounded_channel(); + state + .handle_raft_event( + RaftEvent::StreamSnapshot(ack_rx, chunk_tx, startup_tx), + &context, + role_tx, + ) + .await + .expect("handler must not return Err"); + + // startup_rx must carry Ok — transfer has started + let result = startup_rx.await.expect("startup_tx must be sent"); + assert!( + result.is_ok(), + "expected Ok from startup channel, got {:?}", + result + ); + + // Background task delivers the chunk; give it a moment to run + let chunk = tokio::time::timeout(std::time::Duration::from_millis(200), chunk_rx.recv()) + .await + .expect("timed out waiting for chunk") + .expect("chunk_rx closed before first chunk"); + + assert_eq!(chunk.seq, 0); +} diff --git a/d-engine-core/src/raft_role/learner_state.rs b/d-engine-core/src/raft_role/learner_state.rs index d1fd8933..f5078258 100644 --- a/d-engine-core/src/raft_role/learner_state.rs +++ b/d-engine-core/src/raft_role/learner_state.rs @@ -446,9 +446,16 @@ impl RaftRoleState for LearnerState { return Ok(()); } - RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx) => { + RaftEvent::StreamSnapshot(_ack_rx, _chunk_tx, startup_tx) => { debug!("Learner::RaftEvent::StreamSnapshot"); - warn!("Candidate should not receive StreamSnapshot event."); + warn!("Learner 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(()); } diff --git a/d-engine-server/src/network/grpc/grpc_raft_service.rs b/d-engine-server/src/network/grpc/grpc_raft_service.rs index 17039ace..75236e7d 100644 --- a/d-engine-server/src/network/grpc/grpc_raft_service.rs +++ b/d-engine-server/src/network/grpc/grpc_raft_service.rs @@ -48,6 +48,7 @@ use std::sync::Arc; use std::time::Duration; use tokio::select; use tokio::sync::mpsc; +use tokio::sync::oneshot; use tokio::time::timeout; use tokio_stream::wrappers::ReceiverStream; use tokio_util::sync::CancellationToken; @@ -251,12 +252,22 @@ where // Server creates chunk channel, passes tx to core, keeps rx for gRPC response let (chunk_tx, chunk_rx) = mpsc::channel::>(32); + let (startup_tx, startup_rx) = oneshot::channel::>(); self.event_tx - .send(RaftEvent::StreamSnapshot(ack_rx, chunk_tx)) + .send(RaftEvent::StreamSnapshot(ack_rx, chunk_tx, startup_tx)) .await .map_err(|_| Status::internal("Event channel closed"))?; + // Wait for core to confirm it can serve the snapshot + match startup_rx.await { + Ok(Ok(())) => { + debug!("stream request been processed"); + } + Ok(Err(status)) => return Err(status), + Err(_) => return Err(Status::internal("Core dropped startup sender")), + } + // Return chunk stream directly — no waiting needed let response_stream = ReceiverStream::new(chunk_rx).map(|arc_chunk| Ok((*arc_chunk).clone())); From 68d8de37c94cfbd48719e13ac7e38a7a313ac3f9 Mon Sep 17 00:00:00 2001 From: Joshua Chi <112539+JoshuaChi@users.noreply.github.com> Date: Sun, 21 Jun 2026 17:15:38 +0800 Subject: [PATCH 3/3] fix #409: log StreamSnapshot startup_tx send failure on non-leader roles --- d-engine-core/src/raft_role/follower_state.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/d-engine-core/src/raft_role/follower_state.rs b/d-engine-core/src/raft_role/follower_state.rs index f00e3a12..7a108568 100644 --- a/d-engine-core/src/raft_role/follower_state.rs +++ b/d-engine-core/src/raft_role/follower_state.rs @@ -501,9 +501,15 @@ impl RaftRoleState for FollowerState { } 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(()); }