feat(common): Detect when the cross-process lock has been poisoned.

Author: Hywan

crates/matrix-sdk-common/src/cross_process_lock.rs

@@ -42,13 +42,13 @@ use std::{
     future::Future,
     sync::{
         Arc,
-        atomic::{self, AtomicU32},
+        atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering},
     },
     time::Duration,
 };
 
 use tokio::sync::Mutex;
-use tracing::{debug, error, instrument, trace};
+use tracing::{debug, error, instrument, trace, warn};
 
 use crate::{
     SendOutsideWasm,
@@ -60,9 +60,13 @@ use crate::{
 /// holder.
 pub type CrossProcessLockGeneration = u64;
 
+type AtomicCrossProcessLockGeneration = AtomicU64;
+
 /// Describe the first lock generation value (see [`LockGeneration`]).
 pub const FIRST_CROSS_PROCESS_LOCK_GENERATION: CrossProcessLockGeneration = 1;
 
+pub const NO_CROSS_PROCESS_LOCK_GENERATION: CrossProcessLockGeneration = 0;
+
 /// Trait used to try to take a lock. Foundation of [`CrossProcessLock`].
 pub trait TryLock {
     #[cfg(not(target_family = "wasm"))]
@@ -81,7 +85,11 @@ pub trait TryLock {
     /// - If there was no previous lease, we will acquire the lock.
     /// - Otherwise, we don't get the lock.
     ///
-    /// Returns whether taking the lock succeeded.
+    /// Returns `Some(_)` to indicate the lock succeeded, `None` otherwise. The
+    /// cross-process lock generation must be compared to the generation before
+    /// the call to see if the lock has been poisoned: a different generation
+    /// means the lock has been poisoned, i.e. taken by a different holder in
+    /// the meantime.
     fn try_lock(
         &self,
         lease_duration_ms: u32,
@@ -111,7 +119,7 @@ pub struct CrossProcessLockGuard {
 
 impl Drop for CrossProcessLockGuard {
     fn drop(&mut self) {
-        self.num_holders.fetch_sub(1, atomic::Ordering::SeqCst);
+        self.num_holders.fetch_sub(1, Ordering::SeqCst);
     }
 }
 
@@ -154,6 +162,10 @@ where
 
     /// Backoff time, in milliseconds.
     backoff: Arc<Mutex<WaitingTime>>,
+
+    lock_generation: Arc<AtomicCrossProcessLockGeneration>,
+
+    is_poisoned: Arc<AtomicBool>,
 }
 
 /// Amount of time a lease of the lock should last, in milliseconds.
@@ -193,9 +205,32 @@ where
             num_holders: Arc::new(0.into()),
             locking_attempt: Arc::new(Mutex::new(())),
             renew_task: Default::default(),
+            lock_generation: Arc::new(AtomicCrossProcessLockGeneration::new(
+                FIRST_CROSS_PROCESS_LOCK_GENERATION,
+            )),
+            is_poisoned: Arc::new(AtomicBool::new(false)),
         }
     }
 
+    /// Determine whether the cross-process lock is poisoned.
+    ///
+    /// If another process has taken the lock, then this lock becomes poisoned.
+    /// You should not trust a `false` value for program correctness without
+    /// additional synchronisation.
+    pub fn is_poisoned(&self) -> bool {
+        self.is_poisoned.load(Ordering::SeqCst)
+    }
+
+    /// Clear the poisoned state from this cross-process lock.
+    ///
+    /// If the cross-process lock is poisoned, it will remain poisoned until
+    /// this method is called. This allows recovering from a poisoned
+    /// state and marking that it has recovered.
+    pub fn clear_poison(&self) {
+        self.is_poisoned.store(false, Ordering::SeqCst);
+        self.lock_generation.store(NO_CROSS_PROCESS_LOCK_GENERATION, Ordering::SeqCst);
+    }
+
     /// Try to lock once, returns whether the lock was obtained or not.
     #[instrument(skip(self), fields(?self.lock_key, ?self.lock_holder))]
     pub async fn try_lock_once(
@@ -207,25 +242,48 @@ where
 
         // If another thread obtained the lock, make sure to only superficially increase
         // the number of holders, and carry on.
-        if self.num_holders.load(atomic::Ordering::SeqCst) > 0 {
+        if self.num_holders.load(Ordering::SeqCst) > 0 {
             // Note: between the above load and the fetch_add below, another thread may
             // decrement `num_holders`. That's fine because that means the lock
             // was taken by at least one thread, and after this call it will be
             // taken by at least one thread.
             trace!("We already had the lock, incrementing holder count");
-            self.num_holders.fetch_add(1, atomic::Ordering::SeqCst);
+            self.num_holders.fetch_add(1, Ordering::SeqCst);
             let guard = CrossProcessLockGuard { num_holders: self.num_holders.clone() };
             return Ok(Some(guard));
         }
 
-        let acquired = self
+        if let Some(new_generation) = self
             .locker
             .try_lock(LEASE_DURATION_MS, &self.lock_key, &self.lock_holder)
             .await
             .map_err(|err| CrossProcessLockError::TryLockError(Box::new(err)))?
-            .is_some();
+        {
+            match self.lock_generation.swap(new_generation, Ordering::SeqCst) {
+                // If there was no lock generation, it means this is the first time the lock is
+                // acquired. It cannot be poisoned.
+                NO_CROSS_PROCESS_LOCK_GENERATION => {
+                    trace!(?new_generation, "Setting the lock generation for the first time");
+                }
 
-        if !acquired {
+                // This was NOT the same generation, the lock has been poisoned!
+                previous_generation if previous_generation != new_generation => {
+                    warn!(
+                        ?previous_generation,
+                        ?new_generation,
+                        "The lock has been acquired, but it's been poisoned!"
+                    );
+                    self.is_poisoned.store(true, Ordering::SeqCst);
+                }
+
+                // This was the same generation, no problem.
+                _ => {
+                    trace!("Same lock generation; no problem");
+                }
+            }
+
+            trace!("Lock acquired!");
+        } else {
             trace!("Couldn't acquire the lock immediately.");
             return Ok(None);
         }
@@ -269,7 +327,7 @@ where
                     let _guard = this.locking_attempt.lock().await;
 
                     // If there are no more users, we can quit.
-                    if this.num_holders.load(atomic::Ordering::SeqCst) == 0 {
+                    if this.num_holders.load(Ordering::SeqCst) == 0 {
                         trace!("exiting the lease extension loop");
 
                         // Cancel the lease with another 0ms lease.
@@ -284,18 +342,55 @@ where
 
                 sleep(Duration::from_millis(EXTEND_LEASE_EVERY_MS)).await;
 
-                let fut =
-                    this.locker.try_lock(LEASE_DURATION_MS, &this.lock_key, &this.lock_holder);
+                match this
+                    .locker
+                    .try_lock(LEASE_DURATION_MS, &this.lock_key, &this.lock_holder)
+                    .await
+                {
+                    Ok(Some(new_generation)) => {
+                        match this.lock_generation.swap(new_generation, Ordering::SeqCst) {
+                            // It's impossible to renew the lock if the lock wasn't acquired at
+                            // least once. This is unreachable.
+                            NO_CROSS_PROCESS_LOCK_GENERATION => unreachable!(
+                                "It's impossible to renew a lock lease that has not been acquired once"
+                            ),
+
+                            // This was NOT the same generation, the lock has been poisoned!
+                            previous_generation if previous_generation != new_generation => {
+                                warn!(
+                                    ?previous_generation,
+                                    ?new_generation,
+                                    "The lock lease has been renewed, but it's been poisoned!"
+                                );
+                                this.is_poisoned.store(true, Ordering::SeqCst);
+
+                                // Exit the loop.
+                                break;
+                            }
+
+                            // This was the same generation, no problem.
+                            _ => {}
+                        }
+                    }
 
-                if let Err(err) = fut.await {
-                    error!("error when extending lock lease: {err:#}");
-                    // Exit the loop.
-                    break;
+                    Ok(None) => {
+                        error!("Failed to renew the lock lease: the lock could not be acquired");
+
+                        // Exit the loop.
+                        break;
+                    }
+
+                    Err(err) => {
+                        error!("Error when extending the lock lease: {err:#}");
+
+                        // Exit the loop.
+                        break;
+                    }
                 }
             }
         }));
 
-        self.num_holders.fetch_add(1, atomic::Ordering::SeqCst);
+        self.num_holders.fetch_add(1, Ordering::SeqCst);
 
         let guard = CrossProcessLockGuard { num_holders: self.num_holders.clone() };
         Ok(Some(guard))