Move the SM worker's implementation into SM trait

openraft/src/core/raft_core.rs

@@ -776,7 +776,7 @@ where
             network,
             snapshot_network,
             self.log_store.get_log_reader().await,
-            self.sm_handle.new_snapshot_reader(),
+            self.sm_handle.new_snapshot_reader(self.tx_notification.clone()),
             self.tx_notification.clone(),
             tracing::span!(parent: &self.span, Level::DEBUG, "replication", id=display(&self.id), target=display(&target)),
         )

openraft/src/core/sm/handle.rs

@@ -3,7 +3,9 @@
 use crate::async_runtime::MpscUnboundedSender;
 use crate::async_runtime::MpscUnboundedWeakSender;
 use crate::async_runtime::SendError;
+use crate::core::notification::Notification;
 use crate::core::sm;
+use crate::storage::RaftStateMachineCommand;
 use crate::storage::Snapshot;
 use crate::type_config::alias::JoinHandleOf;
 use crate::type_config::alias::MpscUnboundedSenderOf;
@@ -15,7 +17,9 @@ use crate::RaftTypeConfig;
 pub(crate) struct Handle<C>
 where C: RaftTypeConfig
 {
-    pub(in crate::core::sm) cmd_tx: MpscUnboundedSenderOf<C, sm::Command<C>>,
+    pub(in crate::core::sm) cmd_tx: MpscUnboundedSenderOf<C, RaftStateMachineCommand<C>>,
+
+    pub(in crate::core::sm) tx_notification: MpscUnboundedSenderOf<C, Notification<C>>,
 
     #[allow(dead_code)]
     pub(in crate::core::sm) join_handle: JoinHandleOf<C, ()>,
@@ -24,15 +28,20 @@ where C: RaftTypeConfig
 impl<C> Handle<C>
 where C: RaftTypeConfig
 {
-    pub(crate) fn send(&mut self, cmd: sm::Command<C>) -> Result<(), SendError<sm::Command<C>>> {
+    pub(crate) fn send(&mut self, cmd: sm::Command<C>) -> Result<(), SendError<RaftStateMachineCommand<C>>> {
         tracing::debug!("sending command to state machine worker: {:?}", cmd);
+        let cmd = RaftStateMachineCommand::new(cmd, &self.tx_notification);
         self.cmd_tx.send(cmd)
     }
 
     /// Create a [`SnapshotReader`] to get the current snapshot from the state machine.
-    pub(crate) fn new_snapshot_reader(&self) -> SnapshotReader<C> {
+    pub(crate) fn new_snapshot_reader(
+        &self,
+        tx_notification: MpscUnboundedSenderOf<C, Notification<C>>,
+    ) -> SnapshotReader<C> {
         SnapshotReader {
             cmd_tx: self.cmd_tx.downgrade(),
+            tx_notification,
         }
     }
 }
@@ -46,7 +55,9 @@ where C: RaftTypeConfig
     /// It is weak because the [`Worker`] watches the close event of this channel for shutdown.
     ///
     /// [`Worker`]: sm::worker::Worker
-    cmd_tx: MpscUnboundedWeakSenderOf<C, sm::Command<C>>,
+    cmd_tx: MpscUnboundedWeakSenderOf<C, RaftStateMachineCommand<C>>,
+
+    tx_notification: MpscUnboundedSenderOf<C, Notification<C>>,
 }
 
 impl<C> SnapshotReader<C>
@@ -68,7 +79,7 @@ where C: RaftTypeConfig
         };
 
         // If fail to send command, cmd is dropped and tx will be dropped.
-        let _ = cmd_tx.send(cmd);
+        let _ = cmd_tx.send(RaftStateMachineCommand::new(cmd, &self.tx_notification));
 
         let got = match rx.await {
             Ok(x) => x,

openraft/src/core/sm/worker.rs

@@ -3,14 +3,11 @@ use std::collections::BTreeMap;
 use anyerror::AnyError;
 use tracing_futures::Instrument;
 
-use crate::async_runtime::MpscUnboundedReceiver;
 use crate::async_runtime::MpscUnboundedSender;
 use crate::async_runtime::OneshotSender;
-use crate::base::BoxAsyncOnceMut;
 use crate::core::notification::Notification;
 use crate::core::raft_msg::ResultSender;
 use crate::core::sm::handle::Handle;
-use crate::core::sm::Command;
 use crate::core::sm::CommandResult;
 use crate::core::sm::Response;
 use crate::core::ApplyResult;
@@ -23,6 +20,7 @@ use crate::raft::ClientWriteResponse;
 #[cfg(doc)]
 use crate::storage::RaftLogStorage;
 use crate::storage::RaftStateMachine;
+use crate::storage::RaftStateMachineCommand;
 use crate::storage::Snapshot;
 use crate::type_config::alias::JoinHandleOf;
 use crate::type_config::alias::LogIdOf;
@@ -48,7 +46,7 @@ where
     log_reader: LR,
 
     /// Read command from RaftCore to execute.
-    cmd_rx: MpscUnboundedReceiverOf<C, Command<C>>,
+    cmd_rx: MpscUnboundedReceiverOf<C, RaftStateMachineCommand<C>>,
 
     /// Send back the result of the command to RaftCore.
     resp_tx: MpscUnboundedSenderOf<C, Notification<C>>,
@@ -68,6 +66,7 @@ where
         span: tracing::Span,
     ) -> Handle<C> {
         let (cmd_tx, cmd_rx) = C::mpsc_unbounded();
+        let tx_notification = resp_tx.clone();
 
         let worker = Worker {
             state_machine,
@@ -78,7 +77,11 @@ where
 
         let join_handle = worker.do_spawn(span);
 
-        Handle { cmd_tx, join_handle }
+        Handle {
+            cmd_tx,
+            tx_notification,
+            join_handle,
+        }
     }
 
     fn do_spawn(mut self, span: tracing::Span) -> JoinHandleOf<C, ()> {
@@ -98,76 +101,9 @@ where
 
     #[tracing::instrument(level = "debug", skip_all)]
     async fn worker_loop(&mut self) -> Result<(), StorageError<C>> {
-        loop {
-            let cmd = self.cmd_rx.recv().await;
-            let cmd = match cmd {
-                None => {
-                    tracing::info!("{}: rx closed, state machine worker quit", func_name!());
-                    return Ok(());
-                }
-                Some(x) => x,
-            };
-
-            tracing::debug!("{}: received command: {:?}", func_name!(), cmd);
-
-            match cmd {
-                Command::BuildSnapshot => {
-                    tracing::info!("{}: build snapshot", func_name!());
-
-                    // It is a read operation and is spawned, and it responds in another task
-                    self.build_snapshot(self.resp_tx.clone()).await;
-                }
-                Command::GetSnapshot { tx } => {
-                    tracing::info!("{}: get snapshot", func_name!());
-
-                    self.get_snapshot(tx).await?;
-                    // GetSnapshot does not respond to RaftCore
-                }
-                Command::InstallFullSnapshot { io_id, snapshot } => {
-                    tracing::info!("{}: install complete snapshot", func_name!());
-
-                    let meta = snapshot.meta.clone();
-                    self.state_machine.install_snapshot(&meta, snapshot.snapshot).await?;
-
-                    tracing::info!("Done install complete snapshot, meta: {}", meta);
-
-                    let res = CommandResult::new(Ok(Response::InstallSnapshot((io_id, Some(meta)))));
-                    let _ = self.resp_tx.send(Notification::sm(res));
-                }
-                Command::BeginReceivingSnapshot { tx } => {
-                    tracing::info!("{}: BeginReceivingSnapshot", func_name!());
-
-                    let snapshot_data = self.state_machine.begin_receiving_snapshot().await?;
-
-                    let _ = tx.send(Ok(snapshot_data));
-                    // No response to RaftCore
-                }
-                Command::Apply {
-                    first,
-                    last,
-                    mut client_resp_channels,
-                } => {
-                    let resp = self.apply(first, last, &mut client_resp_channels).await?;
-                    let res = CommandResult::new(Ok(Response::Apply(resp)));
-                    let _ = self.resp_tx.send(Notification::sm(res));
-                }
-                Command::Func { func, input_sm_type } => {
-                    tracing::debug!("{}: run user defined Func", func_name!());
-
-                    let res: Result<Box<BoxAsyncOnceMut<'static, SM>>, _> = func.downcast();
-                    if let Ok(f) = res {
-                        f(&mut self.state_machine).await;
-                    } else {
-                        tracing::warn!(
-                            "User-defined SM function uses incorrect state machine type, expected: {}, got: {}",
-                            std::any::type_name::<SM>(),
-                            input_sm_type
-                        );
-                    };
-                }
-            };
-        }
+        self.state_machine.worker(&mut self.cmd_rx, &self.log_reader).await
     }
+
     #[tracing::instrument(level = "debug", skip_all)]
     async fn apply(
         &mut self,

openraft/src/storage/helper.rs

@@ -186,7 +186,7 @@ where
             chunk_size,
         );
 
-        let mut log_reader = self.log_store.get_log_reader().await;
+        let log_reader = self.log_store.get_log_reader().await;
 
         while start < end {
             let chunk_end = std::cmp::min(end, start + chunk_size);
@@ -290,7 +290,7 @@ where
         let step = 64;
 
         let mut res = vec![];
-        let mut log_reader = self.log_store.get_log_reader().await;
+        let log_reader = self.log_store.get_log_reader().await;
 
         while start < end {
             let step_start = std::cmp::max(start, end.saturating_sub(step));

openraft/src/storage/mod.rs

@@ -19,8 +19,12 @@ pub use self::log_state::LogState;
 pub use self::snapshot::Snapshot;
 pub use self::snapshot_meta::SnapshotMeta;
 pub use self::snapshot_signature::SnapshotSignature;
+pub use self::v2::ApplyResultSender;
+pub use self::v2::BuildSnapshotResultSender;
+pub use self::v2::InstallSnapshotResultSender;
 pub use self::v2::RaftLogReader;
 pub use self::v2::RaftLogStorage;
 pub use self::v2::RaftLogStorageExt;
 pub use self::v2::RaftSnapshotBuilder;
 pub use self::v2::RaftStateMachine;
+pub use self::v2::RaftStateMachineCommand;

openraft/src/storage/v2/mod.rs

@@ -13,4 +13,8 @@ pub use self::raft_log_reader::RaftLogReader;
 pub use self::raft_log_storage::RaftLogStorage;
 pub use self::raft_log_storage_ext::RaftLogStorageExt;
 pub use self::raft_snapshot_builder::RaftSnapshotBuilder;
+pub use self::raft_state_machine::ApplyResultSender;
+pub use self::raft_state_machine::BuildSnapshotResultSender;
+pub use self::raft_state_machine::InstallSnapshotResultSender;
 pub use self::raft_state_machine::RaftStateMachine;
+pub use self::raft_state_machine::RaftStateMachineCommand;

openraft/src/storage/v2/raft_log_reader.rs

@@ -48,7 +48,7 @@ where C: RaftTypeConfig
     /// - The read operation must be transactional. That is, it should not reflect any state changes
     ///   that occur after the read operation has commenced.
     async fn try_get_log_entries<RB: RangeBounds<u64> + Clone + Debug + OptionalSend>(
-        &mut self,
+        &self,
         range: RB,
     ) -> Result<Vec<C::Entry>, StorageError<C>>;
 
@@ -70,7 +70,7 @@ where C: RaftTypeConfig
     ///
     /// The default implementation just returns the full range of log entries.
     #[since(version = "0.10.0")]
-    async fn limited_get_log_entries(&mut self, start: u64, end: u64) -> Result<Vec<C::Entry>, StorageError<C>> {
+    async fn limited_get_log_entries(&self, start: u64, end: u64) -> Result<Vec<C::Entry>, StorageError<C>> {
         self.try_get_log_entries(start..end).await
     }