fix: snapshot replication does not need to send a last 0 size chunk

drmingdrmer · drmingdrmer · commit eee8e534e0b0 · 2021-08-22T17:52:06.000+08:00
diff --git a/async-raft/Cargo.toml b/async-raft/Cargo.toml
@@ -29,6 +29,7 @@ tracing-futures = "0.2.4"
 [dev-dependencies]
 maplit = "1.0.2"
 memstore = { version="0.2.0", path="../memstore" }
+pretty_assertions = "0.7.2"
 tracing-subscriber = "0.2.10"
 
 [features]
diff --git a/async-raft/src/replication/mod.rs b/async-raft/src/replication/mod.rs
@@ -910,31 +910,36 @@ impl<'a, D: AppData, R: AppDataResponse, N: RaftNetwork<D>, S: RaftStorage<D, R>
 
     #[tracing::instrument(level = "trace", skip(self, snapshot))]
     async fn stream_snapshot(&mut self, mut snapshot: CurrentSnapshotData<S::Snapshot>) -> RaftResult<()> {
+        let end = snapshot.snapshot.seek(SeekFrom::End(0)).await?;
+
         let mut offset = 0;
+
         self.replication_core.next_index = snapshot.meta.last_log_id.index + 1;
         self.replication_core.matched = snapshot.meta.last_log_id;
         let mut buf = Vec::with_capacity(self.replication_core.config.snapshot_max_chunk_size as usize);
+
         loop {
             // Build the RPC.
             snapshot.snapshot.seek(SeekFrom::Start(offset)).await?;
-            let nread = snapshot.snapshot.read_buf(&mut buf).await?;
-            let done = nread == 0; // If bytes read == 0, then we're done.
+            let n_read = snapshot.snapshot.read_buf(&mut buf).await?;
+
+            let done = (offset + n_read as u64) == end; // If bytes read == 0, then we're done.
             let req = InstallSnapshotRequest {
                 term: self.replication_core.term,
                 leader_id: self.replication_core.id,
                 meta: snapshot.meta.clone(),
                 offset,
-                data: Vec::from(&buf[..nread]),
+                data: Vec::from(&buf[..n_read]),
                 done,
             };
             buf.clear();
 
             // Send the RPC over to the target.
             tracing::debug!(
                 snapshot_size = req.data.len(),
-                nread,
-                req.done,
                 req.offset,
+                end,
+                req.done,
                 "sending snapshot chunk"
             );
 
@@ -978,7 +983,7 @@ impl<'a, D: AppData, R: AppDataResponse, N: RaftNetwork<D>, S: RaftStorage<D, R>
             }
 
             // Everything is good, so update offset for sending the next chunk.
-            offset += nread as u64;
+            offset += n_read as u64;
 
             // Check raft channel to ensure we are staying up-to-date, then loop.
             if let Some(Some((event, span))) = self.replication_core.repl_rx.recv().now_or_never() {
diff --git a/async-raft/tests/fixtures/mod.rs b/async-raft/tests/fixtures/mod.rs
@@ -36,6 +36,10 @@ use maplit::btreeset;
 use memstore::ClientRequest as MemClientRequest;
 use memstore::ClientResponse as MemClientResponse;
 use memstore::MemStore;
+#[allow(unused_imports)]
+use pretty_assertions::assert_eq;
+#[allow(unused_imports)]
+use pretty_assertions::assert_ne;
 use tokio::sync::RwLock;
 use tracing_subscriber::prelude::*;
 
diff --git a/async-raft/tests/leader_metrics.rs b/async-raft/tests/leader_metrics.rs
@@ -12,6 +12,10 @@ use fixtures::RaftRouter;
 use futures::stream::StreamExt;
 use maplit::btreeset;
 use maplit::hashmap;
+#[allow(unused_imports)]
+use pretty_assertions::assert_eq;
+#[allow(unused_imports)]
+use pretty_assertions::assert_ne;
 
 #[macro_use]
 mod fixtures;
