From bc4b4e54661308ae47b88b39373f82c54dd5780c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E7=82=8E=E6=B3=BC?= Date: Sun, 28 Jul 2024 16:49:21 +0800 Subject: [PATCH] Refactor: Implement custom async Mutex using Oneshot channel This commit introduces a custom implementation of an asynchronous Mutex using the `AsyncRuntime::Oneshot` functions, tailored specifically for Openraft's limited use of asynchronous locks. This custom mutex replaces the previously used Tokio mutex. - **Refactor of `RaftInner::tx_shutdown`:** The `tx_shutdown` member of `RaftInner` has been changed to use a standard (synchronous) Mutex instead of an asynchronous one. This change is made because `tx_shutdown` does not span across `.await` points, making the asynchronous capabilities unnecessary. - **OneshotSender `Drop` Implementation:** It is now documented that the `OneshotSender` should implement the `Drop` trait to ensure that when a sender is dropped, the receiver is notified and yields an error. This behavior is crucial for maintaining robust error handling in asynchronous communication patterns. --- openraft/src/lib.rs | 1 + openraft/src/raft/mod.rs | 2 +- openraft/src/raft/raft_inner.rs | 4 +- openraft/src/replication/mod.rs | 6 +- openraft/src/sync/mod.rs | 5 + openraft/src/sync/mutex.rs | 168 ++++++++++++++++++ .../src/type_config/async_runtime/oneshot.rs | 2 + 7 files changed, 182 insertions(+), 6 deletions(-) create mode 100644 openraft/src/sync/mod.rs create mode 100644 openraft/src/sync/mutex.rs diff --git a/openraft/src/lib.rs b/openraft/src/lib.rs index ba80f0d2d..4f87b13d0 100644 --- a/openraft/src/lib.rs +++ b/openraft/src/lib.rs @@ -62,6 +62,7 @@ pub mod metrics; pub mod network; pub mod raft; pub mod storage; +pub mod sync; pub mod testing; pub mod type_config; diff --git a/openraft/src/raft/mod.rs b/openraft/src/raft/mod.rs index 7a3350135..bb904b8f9 100644 --- a/openraft/src/raft/mod.rs +++ b/openraft/src/raft/mod.rs @@ -37,7 +37,6 @@ pub use message::InstallSnapshotResponse; pub use message::SnapshotResponse; pub use message::VoteRequest; pub use message::VoteResponse; -use tokio::sync::Mutex; use tracing::trace_span; use tracing::Instrument; use tracing::Level; @@ -78,6 +77,7 @@ pub use crate::raft::runtime_config_handle::RuntimeConfigHandle; use crate::raft::trigger::Trigger; use crate::storage::RaftLogStorage; use crate::storage::RaftStateMachine; +use crate::sync::Mutex; use crate::type_config::alias::JoinErrorOf; use crate::type_config::alias::ResponderOf; use crate::type_config::alias::ResponderReceiverOf; diff --git a/openraft/src/raft/raft_inner.rs b/openraft/src/raft/raft_inner.rs index 18c139230..785702990 100644 --- a/openraft/src/raft/raft_inner.rs +++ b/openraft/src/raft/raft_inner.rs @@ -3,7 +3,6 @@ use std::fmt::Debug; use std::future::Future; use std::sync::Arc; -use tokio::sync::Mutex; use tracing::Level; use crate::async_runtime::watch::WatchReceiver; @@ -18,6 +17,7 @@ use crate::error::RaftError; use crate::metrics::RaftDataMetrics; use crate::metrics::RaftServerMetrics; use crate::raft::core_state::CoreState; +use crate::sync::Mutex; use crate::type_config::alias::AsyncRuntimeOf; use crate::type_config::alias::MpscUnboundedSenderOf; use crate::type_config::alias::OneshotReceiverOf; @@ -48,7 +48,7 @@ where C: RaftTypeConfig pub(in crate::raft) core_state: std::sync::Mutex>, /// The ongoing snapshot transmission. - pub(in crate::raft) snapshot: Mutex>>, + pub(in crate::raft) snapshot: Mutex>>, } impl RaftInner diff --git a/openraft/src/replication/mod.rs b/openraft/src/replication/mod.rs index 9fe857d0b..2d2ec0f46 100644 --- a/openraft/src/replication/mod.rs +++ b/openraft/src/replication/mod.rs @@ -19,7 +19,6 @@ use request::Replicate; use response::ReplicationResult; pub(crate) use response::Response; use tokio::select; -use tokio::sync::Mutex; use tracing_futures::Instrument; use crate::async_runtime::MpscUnboundedReceiver; @@ -50,6 +49,7 @@ use crate::replication::request_id::RequestId; use crate::storage::RaftLogReader; use crate::storage::RaftLogStorage; use crate::storage::Snapshot; +use crate::sync::Mutex; use crate::type_config::alias::InstantOf; use crate::type_config::alias::JoinHandleOf; use crate::type_config::alias::LogIdOf; @@ -114,7 +114,7 @@ where /// Another `RaftNetwork` specific for snapshot replication. /// /// Snapshot transmitting is a long running task, and is processed in a separate task. - snapshot_network: Arc>, + snapshot_network: Arc>, /// The current snapshot replication state. /// @@ -754,7 +754,7 @@ where async fn send_snapshot( request_id: RequestId, - network: Arc>, + network: Arc>, vote: Vote, snapshot: Snapshot, option: RPCOption, diff --git a/openraft/src/sync/mod.rs b/openraft/src/sync/mod.rs new file mode 100644 index 000000000..c54dd4de9 --- /dev/null +++ b/openraft/src/sync/mod.rs @@ -0,0 +1,5 @@ +pub(crate) mod mutex; + +pub(crate) use mutex::Mutex; +#[allow(unused_imports)] +pub(crate) use mutex::MutexGuard; diff --git a/openraft/src/sync/mutex.rs b/openraft/src/sync/mutex.rs new file mode 100644 index 000000000..fa6d830ed --- /dev/null +++ b/openraft/src/sync/mutex.rs @@ -0,0 +1,168 @@ +use std::cell::UnsafeCell; +use std::ops::Deref; +use std::ops::DerefMut; + +use crate::type_config::alias::OneshotReceiverOf; +use crate::type_config::alias::OneshotSenderOf; +use crate::type_config::TypeConfigExt; +use crate::RaftTypeConfig; + +/// A simple async mutex implementation that uses oneshot channels to notify the next waiting task. +/// +/// Openraft use async mutex in non-performance critical path, +/// so it's ok to use this simple implementation. +/// +/// Since oneshot channel is already required by AsyncRuntime implementation, +/// there is no need for the application to implement Mutex. +pub(crate) struct Mutex +where C: RaftTypeConfig +{ + /// The current lock holder. + /// + /// When the acquired `MutexGuard` is dropped, it will notify the next waiting task via this + /// oneshot channel. + lock_holder: std::sync::Mutex>>, + + /// The value protected by the mutex. + value: UnsafeCell, +} + +impl Mutex +where C: RaftTypeConfig +{ + pub(crate) fn new(value: T) -> Self { + Self { + lock_holder: std::sync::Mutex::new(None), + value: UnsafeCell::new(value), + } + } + + pub(crate) async fn lock(&self) -> MutexGuard<'_, C, T> { + // Every lock() call puts a oneshot receiver into the holder + // and takes out the existing one. + // If the existing one is Some(rx), + // it means the lock is already held by another task. + // In this case, the current task should wait for the lock to be released. + // + // Such approach forms a queue in which every task waits for the previous one. + + let (tx, rx) = C::oneshot(); + let current_rx = { + let mut l = self.lock_holder.lock().unwrap(); + l.replace(rx) + }; + + if let Some(rx) = current_rx { + let _ = rx.await; + } + + MutexGuard { guard: tx, lock: self } + } + + #[allow(dead_code)] + pub(crate) fn into_inner(self) -> T { + self.value.into_inner() + } +} + +/// The guard of the mutex. +pub(crate) struct MutexGuard<'a, C, T> +where C: RaftTypeConfig +{ + /// This is only used to trigger `Drop` to notify the next waiting task. + #[allow(dead_code)] + guard: OneshotSenderOf, + lock: &'a Mutex, +} + +impl<'a, C, T> Deref for MutexGuard<'a, C, T> +where C: RaftTypeConfig +{ + type Target = T; + + fn deref(&self) -> &Self::Target { + unsafe { &*self.lock.value.get() } + } +} + +impl<'a, C, T> DerefMut for MutexGuard<'a, C, T> +where C: RaftTypeConfig +{ + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { &mut *self.lock.value.get() } + } +} + +/// T must be `Send` to make Mutex `Send` +unsafe impl Send for Mutex where T: Send {} + +/// To allow multiple threads to access T through a `&Mutex`, T must be `Send`, +/// because the caller acquires the ownership through `Mutex::lock()`. +unsafe impl Sync for Mutex where T: Send {} + +/// MutexGuard needs to be Sync to across `.await` point. +unsafe impl Sync for MutexGuard<'_, C, T> where T: Send + Sync {} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::engine::testing::UTConfig; + + #[test] + fn bounds() { + fn check_send() {} + fn check_unpin() {} + // This has to take a value, since the async fn's return type is unnameable. + fn check_send_sync_val(_t: T) {} + fn check_send_sync() {} + + check_send::>(); + check_unpin::>(); + check_send_sync::>(); + + let mutex = Mutex::::new(1); + check_send_sync_val(mutex.lock()); + } + + #[test] + fn test_mutex() { + let mutex = Arc::new(Mutex::::new(0)); + + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(8) + .enable_all() + .build() + .expect("Failed building the Runtime"); + + let big_prime_num = 1_000_000_009; + let n = 100_000; + let n_task = 10; + let mut tasks = vec![]; + + for _i in 0..n_task { + let mutex = mutex.clone(); + let h = rt.spawn(async move { + for k in 0..n { + { + let mut guard = mutex.lock().await; + *guard = (*guard + k) % big_prime_num; + } + } + }); + + tasks.push(h); + } + + let got = rt.block_on(async { + for t in tasks { + let _ = t.await; + } + *mutex.lock().await + }); + + println!("got: {}", got); + assert_eq!(got, n_task * n * (n - 1) / 2 % big_prime_num); + } +} diff --git a/openraft/src/type_config/async_runtime/oneshot.rs b/openraft/src/type_config/async_runtime/oneshot.rs index 08048ebac..38af0c930 100644 --- a/openraft/src/type_config/async_runtime/oneshot.rs +++ b/openraft/src/type_config/async_runtime/oneshot.rs @@ -25,6 +25,8 @@ pub trait Oneshot { where T: OptionalSend; } +/// This `Sender` must implement `Drop` to notify the [`Oneshot::Receiver`] that the sending end has +/// been dropped, causing the receiver to return a [`Oneshot::ReceiverError`]. pub trait OneshotSender: OptionalSend + OptionalSync + Sized where T: OptionalSend {