Merge pull request #757 from drmingdrmer/42-refact

Refactor: Add Command::SendResult to send back response

Author: drmingdrmer

openraft/src/core/install_snapshot.rs

@@ -5,11 +5,10 @@
 //! storage or forward messages to other raft nodes.
 
 mod building_state;
-mod install_snapshot;
 mod raft_core;
 mod replication_state;
 mod server_state;
-mod streaming_state;
+pub(crate) mod streaming_state;
 mod tick;
 
 pub(crate) mod snapshot_state;

openraft/src/core/mod.rs

@@ -15,6 +15,7 @@ use maplit::btreeset;
 use pin_utils::pin_mut;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncSeek;
+use tokio::io::AsyncWriteExt;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
 use tokio::sync::watch;
@@ -30,12 +31,14 @@ use crate::config::RuntimeConfig;
 use crate::config::SnapshotPolicy;
 use crate::core::building_state;
 use crate::core::snapshot_state;
+use crate::core::snapshot_state::SnapshotRequestId;
+use crate::core::streaming_state::Streaming;
 use crate::core::ServerState;
 use crate::core::SnapshotResult;
 use crate::display_ext::DisplaySlice;
 use crate::engine::Command;
 use crate::engine::Engine;
-use crate::engine::SendResult;
+use crate::engine::Respond;
 use crate::entry::FromAppData;
 use crate::entry::RaftEntry;
 use crate::entry::RaftPayload;
@@ -46,6 +49,7 @@ use crate::error::ForwardToLeader;
 use crate::error::InitializeError;
 use crate::error::QuorumNotEnough;
 use crate::error::RPCError;
+use crate::error::SnapshotMismatch;
 use crate::error::Timeout;
 use crate::log_id::LogIdOptionExt;
 use crate::log_id::RaftLogId;
@@ -62,8 +66,11 @@ use crate::raft::AppendEntriesTx;
 use crate::raft::ClientWriteResponse;
 use crate::raft::ClientWriteTx;
 use crate::raft::ExternalCommand;
+use crate::raft::InstallSnapshotRequest;
+use crate::raft::InstallSnapshotResponse;
+use crate::raft::InstallSnapshotTx;
 use crate::raft::RaftMsg;
-use crate::raft::RaftRespTx;
+use crate::raft::ResultSender;
 use crate::raft::VoteRequest;
 use crate::raft::VoteResponse;
 use crate::raft::VoteTx;
@@ -87,8 +94,10 @@ use crate::RaftNetworkFactory;
 use crate::RaftStorage;
 use crate::RaftTypeConfig;
 use crate::SnapshotId;
+use crate::SnapshotSegmentId;
 use crate::StorageError;
 use crate::StorageHelper;
+use crate::StorageIOError;
 use crate::Update;
 use crate::Vote;
 
@@ -218,7 +227,7 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftCore<C,
     #[tracing::instrument(level = "trace", skip(self, tx))]
     pub(super) async fn handle_check_is_leader_request(
         &mut self,
-        tx: RaftRespTx<(), CheckIsLeaderError<C::NodeId, C::Node>>,
+        tx: ResultSender<(), CheckIsLeaderError<C::NodeId, C::Node>>,
     ) -> Result<(), StorageError<C::NodeId>> {
         // Setup sentinel values to track when we've received majority confirmation of leadership.
 
@@ -370,7 +379,7 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftCore<C,
         &mut self,
         changes: ChangeMembers<C::NodeId, C::Node>,
         retain: bool,
-        tx: RaftRespTx<ClientWriteResponse<C>, ClientWriteError<C::NodeId, C::Node>>,
+        tx: ResultSender<ClientWriteResponse<C>, ClientWriteError<C::NodeId, C::Node>>,
     ) {
         let res = self.engine.state.membership_state.change_handler().apply(changes, retain);
         let new_membership = match res {
@@ -519,17 +528,111 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftCore<C,
     pub(crate) fn handle_initialize(
         &mut self,
         member_nodes: BTreeMap<C::NodeId, C::Node>,
-        tx: RaftRespTx<(), InitializeError<C::NodeId, C::Node>>,
+        tx: ResultSender<(), InitializeError<C::NodeId, C::Node>>,
     ) {
         let membership = Membership::from(member_nodes);
 
         let entry = C::Entry::new_membership(LogId::default(), membership);
         let res = self.engine.initialize(entry);
-        self.engine.output.push_command(Command::SendInitializeResult {
-            send: SendResult::new(res, tx),
+        self.engine.output.push_command(Command::Respond {
+            resp: Respond::new(res, tx),
         });
     }
 
+    /// Invoked by leader to send chunks of a snapshot to a follower.
+    ///
+    /// Leaders always send chunks in order. It is important to note that, according to the Raft
+    /// spec, a log may only have one snapshot at any time. As snapshot contents are application
+    /// specific, the Raft log will only store a pointer to the snapshot file along with the
+    /// index & term.
+    #[tracing::instrument(level = "debug", skip_all)]
+    pub(crate) async fn handle_install_snapshot_request(
+        &mut self,
+        req: InstallSnapshotRequest<C>,
+        tx: InstallSnapshotTx<C::NodeId>,
+    ) -> Result<(), StorageError<C::NodeId>> {
+        // TODO: move receiving to another thread.
+        tracing::debug!(req = display(req.summary()));
+
+        let snapshot_meta = req.meta.clone();
+        let done = req.done;
+        let offset = req.offset;
+
+        let req_id = SnapshotRequestId::new(*req.vote.leader_id(), snapshot_meta.snapshot_id.clone(), offset);
+
+        let res = self.engine.vote_handler().accept_vote(&req.vote, tx, |state, _rejected| {
+            Ok(InstallSnapshotResponse {
+                vote: *state.vote_ref(),
+            })
+        });
+
+        let tx = match res {
+            Ok(tx) => tx,
+            Err(_) => return Ok(()),
+        };
+
+        let curr_id = self.snapshot_state.streaming.as_ref().map(|s| &s.snapshot_id);
+
+        // Changed to another stream. re-init snapshot state.
+        if curr_id != Some(&req.meta.snapshot_id) {
+            if req.offset > 0 {
+                let mismatch = SnapshotMismatch {
+                    expect: SnapshotSegmentId {
+                        id: snapshot_meta.snapshot_id.clone(),
+                        offset: 0,
+                    },
+                    got: SnapshotSegmentId {
+                        id: snapshot_meta.snapshot_id.clone(),
+                        offset,
+                    },
+                };
+
+                self.engine.output.push_command(Command::Respond {
+                    resp: Respond::new(Err(mismatch.into()), tx),
+                });
+
+                return Ok(());
+            }
+
+            let snapshot_data = self.storage.begin_receiving_snapshot().await?;
+            self.snapshot_state.streaming = Some(Streaming::new(req.meta.snapshot_id.clone(), snapshot_data));
+        }
+
+        tracing::info!("Received snapshot request: {:?}", req_id);
+
+        let streaming = self.snapshot_state.streaming.as_mut().unwrap();
+
+        // Receive the data.
+        streaming.receive(req).await?;
+
+        if done {
+            let streaming = self.snapshot_state.streaming.take().unwrap();
+            let mut data = streaming.snapshot_data;
+
+            data.as_mut()
+                .shutdown()
+                .await
+                .map_err(|e| StorageIOError::write_snapshot(snapshot_meta.signature(), &e))?;
+
+            self.received_snapshot.insert(snapshot_meta.snapshot_id.clone(), data);
+        }
+
+        if done {
+            self.engine.following_handler().install_snapshot(snapshot_meta);
+        }
+
+        self.engine.output.push_command(Command::Respond {
+            resp: Respond::new(
+                Ok(InstallSnapshotResponse {
+                    vote: *self.engine.state.vote_ref(),
+                }),
+                tx,
+            ),
+        });
+
+        Ok(())
+    }
+
     fn handle_building_snapshot_result(
         &mut self,
         result: SnapshotResult<C::NodeId, C::Node>,
@@ -619,7 +722,7 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftCore<C,
 
     /// Reject a request due to the Raft node being in a state which prohibits the request.
     #[tracing::instrument(level = "trace", skip(self, tx))]
-    pub(crate) fn reject_with_forward_to_leader<T, E>(&self, tx: RaftRespTx<T, E>)
+    pub(crate) fn reject_with_forward_to_leader<T, E>(&self, tx: ResultSender<T, E>)
     where E: From<ForwardToLeader<C::NodeId, C::Node>> {
         let mut leader_id = self.current_leader();
         let leader_node = self.get_leader_node(leader_id);
@@ -926,8 +1029,8 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftCore<C,
         tracing::debug!(req = display(req.summary()), func = func_name!());
 
         let resp = self.engine.handle_vote_req(req);
-        self.engine.output.push_command(Command::SendVoteResult {
-            send: SendResult::new(Ok(resp), tx),
+        self.engine.output.push_command(Command::Respond {
+            resp: Respond::new(Ok(resp), tx),
         });
     }
 
@@ -1368,25 +1471,14 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftRuntime
                 debug_assert!(got.is_some(), "there has to be a buffered snapshot data");
             }
             Command::InstallSnapshot { snapshot_meta } => {
-                let snapshot_data = self.received_snapshot.remove(&snapshot_meta.snapshot_id);
+                // Safe unwrap: it is guaranteed that the snapshot data is buffered. Otherwise it is a bug.
+                let data = self.received_snapshot.remove(&snapshot_meta.snapshot_id).unwrap();
 
-                if let Some(data) = snapshot_data {
-                    self.storage.install_snapshot(&snapshot_meta, data).await?;
-                    tracing::debug!("Done install_snapshot, meta: {:?}", snapshot_meta);
-                } else {
-                    unreachable!("buffered snapshot not found: snapshot meta: {:?}", snapshot_meta)
-                }
-            }
-            Command::SendVoteResult { send } => {
-                send.send();
-            }
-            Command::SendAppendEntriesResult { send } => {
-                send.send();
-            }
-            Command::SendInstallSnapshotResult { send } => {
-                send.send();
+                tracing::info!("Start to install_snapshot, meta: {:?}", snapshot_meta);
+                self.storage.install_snapshot(&snapshot_meta, data).await?;
+                tracing::info!("Done install_snapshot, meta: {:?}", snapshot_meta);
             }
-            Command::SendInitializeResult { send } => {
+            Command::Respond { resp: send } => {
                 send.send();
             }
         }