Skip to main content

ri/core/
lock.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18//! # Safe Lock Utilities
19//!
20//! This module provides safe abstractions for lock acquisition and management,
21//! eliminating the need for `.unwrap()` and `.expect()` calls when working with
22//! synchronization primitives. These utilities ensure robust error handling
23//! in concurrent scenarios without risking thread panics.
24//!
25//! ## Key Components
26//!
27//! - **RiLockError**: Specialized error type for lock-related failures
28//! - **RiLockResult**: Result type alias for lock operations
29//! - **RwLockExtensions**: Extension traits for standard `RwLock` types
30//! - **MutexExtensions**: Extension traits for standard `Mutex` types
31//!
32//! ## Design Principles
33//!
34//! 1. **Never Panic**: All lock operations return `Result` instead of panicking
35//! 2. **Poison Error Handling**: Properly handles poisoned locks without panicking
36//! 3. **Consistent API**: Uniform error handling across all lock types
37//! 4. **Zero-Cost Abstraction**: Extension traits add no overhead when unused
38//!
39//! ## Usage
40//!
41//! ```rust,ignore
42//! use std::sync::Arc;
43//! use ri::core::lock::{RwLockExtensions, RiLockResult};
44//!
45//! struct SharedState {
46//!     counter: u64,
47//! }
48//!
49//! impl SharedState {
50//!     fn increment(&mut self) {
51//!         self.counter += 1;
52//!     }
53//!
54//!     fn get_value(&self) -> u64 {
55//!         self.counter
56//!     }
57//! }
58//!
59//! fn example() -> RiLockResult<()> {
60//!     let state = Arc::new(RwLock::new(SharedState { counter: 0 }));
61//!
62//!     // Write lock with error handling
63//!     {
64//!         let mut guard = state.write_safe()?;
65//!         guard.increment();
66//!     }
67//!
68//!     // Read lock with error handling
69//!     {
70//!         let guard = state.read_safe()?;
71//!         assert_eq!(guard.get_value(), 1);
72//!     }
73//!
74//!     Ok(())
75//! }
76//! ```
77
78use std::sync::{RwLock, Mutex, PoisonError, RwLockReadGuard, RwLockWriteGuard, MutexGuard};
79use std::fmt;
80
81#[cfg(feature = "pyo3")]
82use pyo3::pyclass;
83
84/// Specialized error type for lock-related failures.
85///
86/// This error type provides detailed information about lock acquisition failures,
87/// including whether the lock was poisoned or simply contested. The error includes
88/// context about the lock's purpose for better debugging.
89#[derive(Debug, Clone, PartialEq, Eq, Hash)]
90#[cfg_attr(feature = "pyo3", pyclass)]
91pub struct RiLockError {
92    context: String,
93    is_poisoned: bool,
94}
95
96impl RiLockError {
97    /// Creates a new lock error with the given context.
98    ///
99    /// # Arguments
100    ///
101    ///     context: Human-readable description of what was being locked
102    ///
103    /// # Returns
104    ///
105    ///     A new `RiLockError` instance
106    pub fn new(context: &str) -> Self {
107        Self {
108            context: context.to_string(),
109            is_poisoned: false,
110        }
111    }
112
113    /// Creates a poisoned lock error.
114    ///
115    /// # Arguments
116    ///
117    ///     context: Human-readable description of what was being locked
118    ///
119    /// # Returns
120    ///
121    ///     A new `RiLockError` instance marked as poisoned
122    pub fn poisoned(context: &str) -> Self {
123        Self {
124            context: context.to_string(),
125            is_poisoned: true,
126        }
127    }
128
129    /// Gets the context message for this lock error.
130    pub fn get_context(&self) -> &str {
131        &self.context
132    }
133
134    /// Returns whether the lock was poisoned.
135    pub fn is_poisoned(&self) -> bool {
136        self.is_poisoned
137    }
138}
139
140impl fmt::Display for RiLockError {
141    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
142        if self.is_poisoned {
143            write!(f, "Lock poisoned during acquisition: {}", self.context)
144        } else {
145            write!(f, "Lock acquisition failed: {}", self.context)
146        }
147    }
148}
149
150impl std::error::Error for RiLockError {}
151
152#[cfg(feature = "pyo3")]
153#[pyo3::prelude::pymethods]
154impl RiLockError {
155    #[new]
156    fn new_py(context: String, is_poisoned: bool) -> Self {
157        Self {
158            context,
159            is_poisoned,
160        }
161    }
162
163    #[staticmethod]
164    fn create_from_context(context: String) -> Self {
165        Self {
166            context,
167            is_poisoned: false,
168        }
169    }
170
171    #[staticmethod]
172    fn create_poisoned(context: String) -> Self {
173        Self {
174            context,
175            is_poisoned: true,
176        }
177    }
178
179    #[getter]
180    fn is_poisoned_py(&self) -> bool {
181        self.is_poisoned
182    }
183
184    #[getter]
185    fn context(&self) -> String {
186        self.context.clone()
187    }
188
189    fn __str__(&self) -> String {
190        self.to_string()
191    }
192
193    fn __repr__(&self) -> String {
194        format!("RiLockError {{ context: {:?}, is_poisoned: {} }}", self.context, self.is_poisoned)
195    }
196}
197
198/// Result type alias for lock operations.
199///
200/// This type alias simplifies error handling for lock-related operations,
201/// providing a consistent return type for all lock acquisitions.
202pub type RiLockResult<T> = Result<T, RiLockError>;
203
204/// Extension trait providing safe read lock acquisition for `RwLock`.
205///
206/// This trait adds a `read_safe` method to `RwLock` that returns a `Result`
207/// instead of panicking when the lock is poisoned or cannot be acquired.
208pub trait RwLockExtensions<T: Send + Sync> {
209    /// Acquires a read lock safely, returning a Result instead of panicking.
210    ///
211    /// This method attempts to acquire a read lock on the `RwLock`. If the lock
212    /// is held by a writer or if a previous holder panicked (poisoned), this
213    /// method returns an error instead of panicking.
214    ///
215    /// ## Arguments
216    ///
217    /// - `context`: A description of what is being locked for error messages
218    ///
219    /// ## Returns
220    ///
221    /// - `Ok(RwLockReadGuard<T>)` if the lock was acquired successfully
222    /// - `Err(RiLockError)` if the lock could not be acquired
223    ///
224    /// ## Examples
225    ///
226    /// ```rust,ignore
227    /// use std::sync::RwLock;
228    /// use ri::core::lock::RwLockExtensions;
229    ///
230    /// let lock = RwLock::new(42);
231    ///
232    /// match lock.read_safe("counter") {
233    ///     Ok(guard) => println!("Value: {}", *guard),
234    ///     Err(e) => println!("Failed to acquire lock: {}", e),
235    /// }
236    /// ```
237    fn read_safe(&self, context: &str) -> RiLockResult<RwLockReadGuard<'_, T>>;
238    
239    /// Acquires a write lock safely, returning a Result instead of panicking.
240    ///
241    /// This method attempts to acquire a write lock on the `RwLock`. If the lock
242    /// is currently held by any readers or if a previous holder panicked (poisoned),
243    /// this method returns an error instead of panicking.
244    ///
245    /// ## Arguments
246    ///
247    /// - `context`: A description of what is being locked for error messages
248    ///
249    /// ## Returns
250    ///
251    /// - `Ok(RwLockWriteGuard<T>)` if the lock was acquired successfully
252    /// - `Err(RiLockError)` if the lock could not be acquired
253    ///
254    /// ## Examples
255    ///
256    /// ```rust,ignore
257    /// use std::sync::RwLock;
258    /// use ri::core::lock::RwLockExtensions;
259    ///
260    /// let lock = RwLock::new(42);
261    ///
262    /// match lock.write_safe("counter") {
263    ///     Ok(mut guard) => {
264    ///         *guard += 1;
265    ///         println!("New value: {}", *guard);
266    ///     }
267    ///     Err(e) => println!("Failed to acquire lock: {}", e),
268    /// }
269    /// ```
270    fn write_safe(&self, context: &str) -> RiLockResult<RwLockWriteGuard<'_, T>>;
271    
272    /// Attempts to acquire a read lock, returning immediately if unavailable.
273    ///
274    /// This method tries to acquire a read lock without blocking. If the lock
275    /// is held by a writer, it returns an error immediately.
276    ///
277    /// ## Arguments
278    ///
279    /// - `context`: A description of what is being locked for error messages
280    ///
281    /// ## Returns
282    ///
283    /// - `Ok(Some(RwLockReadGuard<T>))` if the lock was acquired
284    /// - `Ok(None)` if the lock is held by a writer
285    /// - `Err(RiLockError)` if the lock is poisoned
286    fn try_read_safe(&self, context: &str) -> RiLockResult<Option<RwLockReadGuard<'_, T>>>;
287    
288    /// Attempts to acquire a write lock, returning immediately if unavailable.
289    ///
290    /// This method tries to acquire a write lock without blocking. If the lock
291    /// is held by any readers, it returns an error immediately.
292    ///
293    /// ## Arguments
294    ///
295    /// - `context`: A description of what is being locked for error messages
296    ///
297    /// ## Returns
298    ///
299    /// - `Ok(Some(RwLockWriteGuard<T>))` if the lock was acquired
300    /// - `Ok(None)` if the lock is held by readers or a writer
301    /// - `Err(RiLockError)` if the lock is poisoned
302    fn try_write_safe(&self, context: &str) -> RiLockResult<Option<RwLockWriteGuard<'_, T>>>;
303}
304
305impl<T: Send + Sync> RwLockExtensions<T> for RwLock<T> {
306    fn read_safe(&self, context: &str) -> RiLockResult<RwLockReadGuard<'_, T>> {
307        RwLock::read(self).map_err(|_| {
308            RiLockError::poisoned(context)
309        })
310    }
311    
312    fn write_safe(&self, context: &str) -> RiLockResult<RwLockWriteGuard<'_, T>> {
313        RwLock::write(self).map_err(|_| {
314            RiLockError::poisoned(context)
315        })
316    }
317    
318    fn try_read_safe(&self, context: &str) -> RiLockResult<Option<RwLockReadGuard<'_, T>>> {
319        match RwLock::try_read(self) {
320            Ok(guard) => Ok(Some(guard)),
321            Err(std::sync::TryLockError::Poisoned(_)) => {
322                Err(RiLockError::poisoned(context))
323            }
324            Err(std::sync::TryLockError::WouldBlock) => Ok(None),
325        }
326    }
327    
328    fn try_write_safe(&self, context: &str) -> RiLockResult<Option<RwLockWriteGuard<'_, T>>> {
329        match RwLock::try_write(self) {
330            Ok(guard) => Ok(Some(guard)),
331            Err(std::sync::TryLockError::Poisoned(_)) => {
332                Err(RiLockError::poisoned(context))
333            }
334            Err(std::sync::TryLockError::WouldBlock) => Ok(None),
335        }
336    }
337}
338
339/// Extension trait providing safe lock acquisition for `Mutex`.
340///
341/// This trait adds a `lock_safe` method to `Mutex` that returns a `Result`
342/// instead of panicking when the lock is poisoned.
343pub trait MutexExtensions<T: Send> {
344    /// Acquires a lock safely, returning a Result instead of panicking.
345    ///
346    /// This method attempts to acquire the mutex lock. If a previous holder
347    /// panicked while holding the lock (poisoned), this method returns an
348    /// error instead of panicking.
349    ///
350    /// ## Arguments
351    ///
352    /// - `context`: A description of what is being locked for error messages
353    ///
354    /// ## Returns
355    ///
356    /// - `Ok(MutexGuard<T>)` if the lock was acquired successfully
357    /// - `Err(RiLockError)` if the lock could not be acquired
358    ///
359    /// ## Examples
360    ///
361    /// ```rust,ignore
362    /// use std::sync::Mutex;
363    /// use ri::core::lock::MutexExtensions;
364    ///
365    /// let mutex = Mutex::new(42);
366    ///
367    /// match mutex.lock_safe("important_data") {
368    ///     Ok(guard) => println!("Value: {}", *guard),
369    ///     Err(e) => println!("Failed to acquire lock: {}", e),
370    /// }
371    /// ```
372    fn lock_safe(&self, context: &str) -> RiLockResult<MutexGuard<'_, T>>;
373    
374    /// Attempts to acquire the lock, returning immediately if unavailable.
375    ///
376    /// This method tries to acquire the mutex lock without blocking. If the
377    /// lock is held by another thread, it returns an error immediately.
378    ///
379    /// ## Arguments
380    ///
381    /// - `context`: A description of what is being locked for error messages
382    ///
383    /// ## Returns
384    ///
385    /// - `Ok(Some(MutexGuard<T>))` if the lock was acquired
386    /// - `Ok(None)` if the lock is held by another thread
387    /// - `Err(RiLockError)` if the lock is poisoned
388    fn try_lock_safe(&self, context: &str) -> RiLockResult<Option<MutexGuard<'_, T>>>;
389}
390
391impl<T: Send> MutexExtensions<T> for Mutex<T> {
392    fn lock_safe(&self, context: &str) -> RiLockResult<MutexGuard<'_, T>> {
393        Mutex::lock(self).map_err(|_| {
394            RiLockError::poisoned(context)
395        })
396    }
397    
398    fn try_lock_safe(&self, context: &str) -> RiLockResult<Option<MutexGuard<'_, T>>> {
399        match Mutex::try_lock(self) {
400            Ok(guard) => Ok(Some(guard)),
401            Err(std::sync::TryLockError::Poisoned(_)) => {
402                Err(RiLockError::poisoned(context))
403            }
404            Err(std::sync::TryLockError::WouldBlock) => Ok(None),
405        }
406    }
407}
408
409/// Utility function to convert from `PoisonError` to `RiLockError`.
410///
411/// This conversion is useful when working with standard library types that
412/// return `PoisonError` and you want to convert to our custom lock error type.
413///
414/// ## Arguments
415///
416/// - `error`: The poison error to convert
417/// - `context`: Description of what was being locked
418///
419/// ## Returns
420///
421/// A `RiLockError` with the appropriate context and poisoned flag
422pub fn from_poison_error<T>(_error: PoisonError<T>, context: &str) -> RiLockError {
423    RiLockError::poisoned(context)
424}
425
426#[cfg(test)]
427mod tests {
428    use super::*;
429    use std::sync::atomic::{AtomicU64, Ordering};
430    use std::sync::Arc;
431    use std::thread;
432
433    #[test]
434    fn test_read_safe_success() {
435        let lock = RwLock::new(42);
436        let result = lock.read_safe("test counter");
437        assert!(result.is_ok());
438        assert_eq!(*result.unwrap(), 42);
439    }
440
441    #[test]
442    fn test_write_safe_success() {
443        let lock = RwLock::new(42);
444        let result = lock.write_safe("test counter");
445        assert!(result.is_ok());
446        assert_eq!(*result.unwrap(), 42);
447    }
448
449    #[test]
450    fn test_try_read_safe_available() {
451        let lock = RwLock::new(42);
452        let result = lock.try_read_safe("test counter");
453        assert!(result.is_ok());
454        assert!(result.unwrap().is_some());
455    }
456
457    #[test]
458    fn test_try_write_safe_unavailable() {
459        let lock = RwLock::new(42);
460        let _write_guard = lock.write_safe("test counter").unwrap();
461        
462        let result = lock.try_write_safe("test counter");
463        assert!(result.is_ok());
464        assert!(result.unwrap().is_none());
465    }
466
467    #[test]
468    fn test_mutex_lock_safe() {
469        let mutex = Mutex::new(42);
470        let result = mutex.lock_safe("test mutex");
471        assert!(result.is_ok());
472        assert_eq!(*result.unwrap(), 42);
473    }
474
475    #[test]
476    fn test_lock_error_display() {
477        let error = RiLockError::new("test context");
478        assert_eq!(error.to_string(), "Lock acquisition failed: test context");
479        
480        let poisoned = RiLockError::poisoned("poisoned lock");
481        assert_eq!(poisoned.to_string(), "Lock poisoned during acquisition: poisoned lock");
482        assert!(poisoned.is_poisoned());
483    }
484
485    #[test]
486    fn test_concurrent_reads() {
487        let lock = Arc::new(RwLock::new(0));
488        let num_threads = 10;
489        let iterations = 1000;
490        
491        let handles: Vec<_> = (0..num_threads)
492            .map(|_| {
493                let lock = Arc::clone(&lock);
494                thread::spawn(move || {
495                    for _ in 0..iterations {
496                        let _guard = lock.read_safe("concurrent read").unwrap();
497                    }
498                })
499            })
500            .collect();
501        
502        for handle in handles {
503            handle.join().unwrap();
504        }
505        
506        assert_eq!(*lock.read_safe("final read").unwrap(), 0);
507    }
508
509    #[test]
510    fn test_concurrent_writes() {
511        let lock = Arc::new(RwLock::new(AtomicU64::new(0)));
512        let num_threads = 4;
513        let iterations = 100;
514        
515        let handles: Vec<_> = (0..num_threads)
516            .map(|_i| {
517                let lock = Arc::clone(&lock);
518                thread::spawn(move || {
519                    for _ in 0..iterations {
520                        let guard = lock.write_safe("concurrent write").unwrap();
521                        guard.fetch_add(1, Ordering::SeqCst);
522                    }
523                })
524            })
525            .collect();
526        
527        for handle in handles {
528            handle.join().unwrap();
529        }
530        
531        let final_value = lock.read_safe("final value").unwrap().load(Ordering::SeqCst);
532        assert_eq!(final_value, (num_threads * iterations) as u64);
533    }
534}