Skip to main content

concurrent/
chainlock.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use core::marker::{PhantomData, PhantomPinned};
6use core::mem::offset_of;
7use core::ptr::NonNull;
8use core::sync::atomic::Ordering;
9
10use crate::chainlock_transaction::{
11    AtomicToken, ChainLockTransactionToken, ContentionState, Token, Transaction,
12};
13
14// # General Chain Lock Facility
15//
16// Chain locks are a type of spin lock designed for use cases where multiple CPUs need to acquire
17// various sets of locks in different orders, such as in dynamic object graphs where nodes have
18// lock-synchronized states and references to other nodes. Chain locks provide an arbitration
19// mechanism that is particularly useful in operations that require locking multiple nodes
20// simultaneously, such as propagating state changes or altering the graph's connectivity.
21// Operations initiated from different nodes in the graph may require locks on some of the same
22// nodes, but in different orders. If ordinary spin locks were used, such lock order variations
23// could lead to deadlocks.
24//
25// Similar to other spin locks, chain locks use atomic variables for lock state and memory
26// ordering, and require disabling interrupts to guarantee IRQ-safe critical sections. However,
27// chain locks differ from spin locks in that they don't use CPU IDs or other static values to
28// indicate the locked state and identify the owner. Instead, the locked value is derived from a
29// global generation counter that provides a priority invariant used in arbitration. An operation
30// that requires one or more chain locks is performed in a transaction that carries a unique
31// generation count. Within a transaction, a CPU can attempt to lock multiple chain locks using the
32// same generation count, giving rise to the name "chain lock."
33//
34// The generation count provides a fair arbitration mechanism when multiple CPUs attempt to acquire
35// the same locks in potentially different orders. During arbitration, the transaction with the
36// lowest generation count has the highest priority. If a CPU encounters contention and has a
37// transaction with a higher priority than the lock's current owner, it simply spins until the lock
38// is released. However, if the transaction has a lower priority, the CPU must back off by
39// releasing all held chain locks and restarting the locking sequence. This is because the CPU with
40// the higher priority transaction may be waiting to acquire one or more of these locks. Once a
41// transaction is complete its generation count is retired and never reused, ensuring that all
42// transactions eventually make progress, even under heavy contention.
43//
44
45/// Result of attempting to acquire a [`ChainLock`].
46#[derive(Clone, Copy, Debug, PartialEq, Eq)]
47pub enum Result {
48    /// Lock successfully acquired.
49    Ok,
50
51    /// Lock is held by a lower priority transaction. Returned only by internal try-acquire logic.
52    /// Other acquire methods spin internally and do not return this result.
53    Retry,
54
55    /// Lock is held by a higher priority transaction, requiring the current transaction to follow
56    /// the back off protocol.
57    Backoff,
58
59    /// Lock is already held by the current transaction and would otherwise deadlock, requiring the
60    /// current transaction to follow the cycle detected protocol. May be fatal in some
61    /// environments.
62    Cycle,
63}
64
65/// Specifies whether [`ChainLock::try_acquire_with_options`] is allowed when the active
66/// transaction has already been finalized.
67#[repr(u8)]
68#[derive(Clone, Copy, Debug, PartialEq, Eq)]
69pub enum AllowFinalized {
70    /// Disallow acquiring the lock if the active transaction is already finalized.
71    No = 0,
72    /// Allow acquiring the lock even if the active transaction is already finalized.
73    Yes = 1,
74}
75
76/// Specifies whether [`ChainLock::try_acquire_with_options`] should record contention backoff
77/// state if acquisition fails due to a higher-priority transaction holding the lock.
78#[repr(u8)]
79#[derive(Clone, Copy, Debug, PartialEq, Eq)]
80pub enum RecordBackoff {
81    /// Do not record backoff state on contention.
82    No = 0,
83    /// Record backoff state on contention with a higher-priority transaction.
84    Yes = 1,
85}
86
87#[derive(Clone, Copy)]
88struct AttemptResult {
89    lock_result: Result,
90    observed_token: Token,
91}
92
93/// Environment-independent chain lock interface and protocol.
94#[repr(C)]
95pub struct ChainLock<T: Transaction> {
96    state: AtomicToken,
97    contention_state: T::ContentionState,
98    // Wrapped in `PhantomData` so `#[repr(C)]` structs embedding a `ChainLock` (e.g. the kernel's
99    // `Thread`) stay FFI-safe when named in `extern` signatures (`improper_ctypes` rejects a bare
100    // `PhantomPinned`, but ignores `PhantomData` fields).
101    _pin: PhantomData<PhantomPinned>,
102}
103
104impl<T: Transaction> ChainLock<T> {
105    const UNLOCKED_TOKEN: Token = Token::UNLOCKED;
106    const MAX_TOKEN: Token = Token::MAX;
107
108    /// Creates a new unlocked `ChainLock`.
109    #[inline]
110    pub const fn new() -> Self {
111        Self {
112            state: AtomicToken::new(Self::UNLOCKED_TOKEN),
113            contention_state: T::ContentionState::INIT,
114            _pin: PhantomData,
115        }
116    }
117
118    /// Acquires the lock without backoff, spinning until acquired. This method may be used for the
119    /// first chain lock acquired in a transaction as an optimization, since backoff is unnecessary
120    /// when no other chain locks are held.
121    ///
122    /// Asserts if a cycle is detected or if any chain locks are currently held.
123    #[inline]
124    pub fn acquire_first_in_chain(&self, token: &ChainLockTransactionToken<T>) {
125        let transaction = T::active_ref(token);
126        transaction.assert_not_finalized();
127        transaction.assert_num_locks_held(0);
128
129        let tx_token = transaction.token();
130        loop {
131            let result = self.attempt_acquire_common(tx_token);
132            if result.lock_result == Result::Ok {
133                transaction.on_acquire(token);
134                return;
135            }
136
137            transaction.on_contention(token);
138
139            debug_assert_ne!(result.lock_result, Result::Cycle);
140            T::yield_cpu();
141        }
142    }
143
144    /// Acquires the lock, spinning until acquired or returning failure because the caller should
145    /// back off due to contention with a higher priority transaction. Returns `true` if the lock
146    /// was acquired or `false` if the current transaction should back off. Asserts if a cycle is
147    /// detected.
148    #[inline]
149    #[must_use]
150    pub fn acquire_or_backoff(&self, token: &ChainLockTransactionToken<T>) -> bool {
151        let transaction = T::active_ref(token);
152        transaction.assert_not_finalized();
153
154        let tx_token = transaction.token();
155        loop {
156            let result = self.attempt_acquire_common(tx_token);
157            if result.lock_result == Result::Ok {
158                transaction.on_acquire(token);
159                return true;
160            }
161
162            transaction.on_contention(token);
163
164            // If recording the backoff fails, the transaction owning the lock changed, which could
165            // be either higher or lower priority than the current transaction. In that case, retry
166            // the acquisition instead of returning false and unwinding the transaction.
167            if result.lock_result == Result::Backoff
168                && transaction.record_backoff_or_retry(
169                    result.observed_token,
170                    &self.state,
171                    &self.contention_state,
172                    token,
173                )
174            {
175                return false;
176            }
177
178            debug_assert_ne!(result.lock_result, Result::Cycle);
179            T::yield_cpu();
180        }
181    }
182
183    /// Acquires the lock, spinning until acquired or returning `Err` with [`Result::Backoff`] or
184    /// [`Result::Cycle`]. Returns `Ok(())` if the lock was acquired or `Err(result)` if acquiring
185    /// the lock failed.
186    #[inline]
187    pub fn acquire_or_result(
188        &self,
189        token: &ChainLockTransactionToken<T>,
190    ) -> core::result::Result<(), Result> {
191        let transaction = T::active_ref(token);
192        transaction.assert_not_finalized();
193
194        let tx_token = transaction.token();
195        loop {
196            let mut result = self.attempt_acquire_common(tx_token);
197            if result.lock_result == Result::Ok {
198                transaction.on_acquire(token);
199                return Ok(());
200            }
201
202            transaction.on_contention(token);
203
204            // If recording the backoff fails, the transaction owning the lock changed, which could
205            // be either higher or lower priority than the current transaction. In that case, retry
206            // the acquisition instead of failing and returning the result.
207            if result.lock_result == Result::Backoff
208                && !transaction.record_backoff_or_retry(
209                    result.observed_token,
210                    &self.state,
211                    &self.contention_state,
212                    token,
213                )
214            {
215                result.lock_result = Result::Retry;
216            }
217
218            if result.lock_result != Result::Retry {
219                return Err(result.lock_result);
220            }
221
222            T::yield_cpu();
223        }
224    }
225
226    /// Attempts to acquire the lock with default options (`AllowFinalized::No`,
227    /// `RecordBackoff::Yes`). Returns `true` if the lock was acquired or `false` if acquiring the
228    /// lock failed. Asserts if a cycle is detected.
229    #[inline]
230    #[must_use]
231    pub fn try_acquire(&self, token: &ChainLockTransactionToken<T>) -> bool {
232        self.try_acquire_with_options(AllowFinalized::No, RecordBackoff::Yes, token)
233    }
234
235    /// Attempts to acquire the lock with the specified `allow_finalized` and `record_backoff`
236    /// options. Returns `true` if the lock was acquired or `false` if acquiring the lock failed.
237    /// Asserts if a cycle is detected.
238    #[inline]
239    #[must_use]
240    pub fn try_acquire_with_options(
241        &self,
242        allow_finalized: AllowFinalized,
243        record_backoff: RecordBackoff,
244        token: &ChainLockTransactionToken<T>,
245    ) -> bool {
246        let transaction = T::active_ref(token);
247        if allow_finalized == AllowFinalized::No {
248            transaction.assert_not_finalized();
249        }
250
251        let tx_token = transaction.token();
252        let result = self.attempt_acquire_common(tx_token);
253        if result.lock_result == Result::Ok {
254            transaction.on_acquire(token);
255            return true;
256        }
257
258        transaction.on_contention(token);
259
260        if result.lock_result == Result::Backoff && record_backoff == RecordBackoff::Yes {
261            let _ = transaction.record_backoff_or_retry(
262                result.observed_token,
263                &self.state,
264                &self.contention_state,
265                token,
266            );
267        }
268
269        debug_assert_ne!(result.lock_result, Result::Cycle);
270        false
271    }
272
273    /// Releases the lock.
274    #[inline]
275    pub fn release(&self, token: &ChainLockTransactionToken<T>) {
276        let transaction = T::active_ref(token);
277        debug_assert_eq!(transaction.token(), self.state.load(Ordering::Relaxed));
278
279        transaction.on_release(token);
280
281        // The environment-specific contention backoff strategy may require a careful order of
282        // operations for notifying backed-off contenders of the release and the actual lock
283        // release operation to ensure that BOTH of the following properties hold:
284        //
285        // 1. Races between backing off and releasing the lock do not result in missed
286        //    notifications.
287        // 2. The contention state of the lock is not accessed after the lock is released to avoid
288        //    use-after-free hazards.
289        //
290        // The following protocol ensures these properties hold:
291        //
292        // 1. The lock token is set to the maximum token value to temporarily prevent further
293        //    updates to the contention state by new contenders, since the maximum token value is
294        //    the lowest possible priority for a chain lock transaction.
295        // 2. The lock contention state is sampled to determine which backed-off contenders to
296        //    notify after the lock is released.
297        // 3. The lock is released, allowing contenders to attempt to acquire the lock.
298        // 4. Contenders that backed off prior to the lock release are notified.
299        //
300        // The environment-specific implementation of `release_and_notify` is responsible for
301        // ensuring that steps 2 through 4 are implemented correctly, using the provided closure to
302        // implement step 3.
303        //
304        // TODO(johngro): Look into relaxing the memory ordering being used here.  I
305        // _think_ that we can perform this store to the lock's state using relaxed
306        // semantics, but I'm not going to take that chance without having first built
307        // a model and run it through CDSChecker.
308        //
309        self.state.store(Self::MAX_TOKEN, Ordering::Release);
310        transaction.release_and_notify(
311            &self.contention_state,
312            || {
313                self.state.store(Self::UNLOCKED_TOKEN, Ordering::Release);
314            },
315            token,
316        );
317    }
318
319    /// Asserts in debug builds that the lock is held by the current transaction.
320    #[inline]
321    pub fn assert_held(&self, token: &ChainLockTransactionToken<T>) {
322        if cfg!(debug_assertions) {
323            let transaction = T::active_ref(token);
324            let tx_token = transaction.token();
325            let lock_token = self.state.load(Ordering::Relaxed);
326            debug_assert_eq!(
327                lock_token,
328                tx_token,
329                "Failed to assert ChainLock is held (lock token {:#x}, our token {:#x}).",
330                lock_token.value(),
331                tx_token.value()
332            );
333        }
334    }
335
336    /// Asserts in debug builds that the lock is held by the current transaction and must be
337    /// released by the caller's scope.
338    #[inline]
339    pub fn assert_acquired(&self, token: &ChainLockTransactionToken<T>) {
340        self.assert_held(token);
341    }
342
343    /// Returns `true` if the lock is held by the current transaction.
344    #[inline]
345    pub fn is_held(&self, token: &ChainLockTransactionToken<T>) -> bool {
346        let transaction = T::active_ref(token);
347        let tx_token = transaction.token();
348        self.state.load(Ordering::Relaxed) == tx_token
349    }
350
351    /// Returns `true` if the lock is unlocked.
352    #[inline]
353    pub fn is_unlocked(&self) -> bool {
354        self.state.load(Ordering::Relaxed) == Self::UNLOCKED_TOKEN
355    }
356
357    /// Syncs the held token value of this lock with the current transaction token. Used for saving
358    /// and restoring state in special circumstances that use reserved tokens.
359    #[inline]
360    pub fn sync_token(&self, token: &ChainLockTransactionToken<T>) {
361        let transaction = T::active_ref(token);
362        let tx_token = transaction.token();
363        self.state.store(tx_token, Ordering::Release);
364    }
365
366    /// Returns the currently stored [`Token`] of this lock with `Relaxed` memory ordering for
367    /// testing.
368    #[inline]
369    pub fn state_for_testing(&self) -> Token {
370        self.state.load(Ordering::Relaxed)
371    }
372
373    #[inline]
374    fn attempt_acquire_common(&self, token: Token) -> AttemptResult {
375        match self.state.compare_exchange(
376            Self::UNLOCKED_TOKEN,
377            token,
378            Ordering::Acquire,
379            Ordering::Acquire,
380        ) {
381            Ok(expected) => AttemptResult { lock_result: Result::Ok, observed_token: expected },
382            Err(expected) => {
383                let lock_result = if expected > token {
384                    Result::Retry
385                } else if expected < token {
386                    Result::Backoff
387                } else {
388                    Result::Cycle
389                };
390                AttemptResult { lock_result, observed_token: expected }
391            }
392        }
393    }
394}
395
396impl<T: Transaction> Default for ChainLock<T> {
397    #[inline]
398    fn default() -> Self {
399        Self::new()
400    }
401}
402
403impl<T: Transaction> Drop for ChainLock<T> {
404    #[inline]
405    fn drop(&mut self) {
406        let lock_token = self.state.load(Ordering::Relaxed);
407        debug_assert_eq!(
408            lock_token,
409            Self::UNLOCKED_TOKEN,
410            "Destroying locked ChainLock: token={}",
411            lock_token.value()
412        );
413    }
414}
415
416/// Utility to simplify expressions where two chain locks need to be acquired with backoff.
417/// Returns `true` if both locks were acquired or `false` if the current transaction needs to
418/// back off.
419#[inline]
420#[must_use]
421pub fn acquire_both_or_backoff<T: Transaction>(
422    first: &ChainLock<T>,
423    second: &ChainLock<T>,
424    token: &ChainLockTransactionToken<T>,
425) -> bool {
426    if first.acquire_or_backoff(token) {
427        if second.acquire_or_backoff(token) {
428            return true;
429        }
430        first.release(token);
431    }
432    false
433}
434
435// Verify layout equivalence with C++ `concurrent::ChainLock<Transaction>` for both 32-bit
436// (`AtomicU32`, i.e., `ktl::atomic<cpu_mask_t>` in the kernel) and 64-bit (`AtomicU64`) contention
437// states.
438struct DummyTxU32;
439impl Transaction for DummyTxU32 {
440    type ContentionState = core::sync::atomic::AtomicU32;
441    fn active() -> Option<NonNull<Self>> {
442        None
443    }
444    fn active_ref(_: &ChainLockTransactionToken<Self>) -> &Self {
445        unimplemented!()
446    }
447    fn yield_cpu() {}
448    fn token(&self) -> Token {
449        Token::UNLOCKED
450    }
451    fn assert_not_finalized(&self) {}
452    fn assert_num_locks_held(&self, _: u32) {}
453    fn on_acquire(&self, _: &ChainLockTransactionToken<Self>) {}
454    fn on_release(&self, _: &ChainLockTransactionToken<Self>) {}
455    fn on_contention(&self, _: &ChainLockTransactionToken<Self>) {}
456    fn record_backoff_or_retry(
457        &self,
458        _: Token,
459        _: &AtomicToken,
460        _: &Self::ContentionState,
461        _: &ChainLockTransactionToken<Self>,
462    ) -> bool {
463        true
464    }
465    fn release_and_notify<F: FnOnce()>(
466        &self,
467        _: &Self::ContentionState,
468        do_release: F,
469        _: &ChainLockTransactionToken<Self>,
470    ) {
471        do_release();
472    }
473}
474
475struct DummyTxU64;
476impl Transaction for DummyTxU64 {
477    type ContentionState = core::sync::atomic::AtomicU64;
478    fn active() -> Option<NonNull<Self>> {
479        None
480    }
481    fn active_ref(_: &ChainLockTransactionToken<Self>) -> &Self {
482        unimplemented!()
483    }
484    fn yield_cpu() {}
485    fn token(&self) -> Token {
486        Token::UNLOCKED
487    }
488    fn assert_not_finalized(&self) {}
489    fn assert_num_locks_held(&self, _: u32) {}
490    fn on_acquire(&self, _: &ChainLockTransactionToken<Self>) {}
491    fn on_release(&self, _: &ChainLockTransactionToken<Self>) {}
492    fn on_contention(&self, _: &ChainLockTransactionToken<Self>) {}
493    fn record_backoff_or_retry(
494        &self,
495        _: Token,
496        _: &AtomicToken,
497        _: &Self::ContentionState,
498        _: &ChainLockTransactionToken<Self>,
499    ) -> bool {
500        true
501    }
502    fn release_and_notify<F: FnOnce()>(
503        &self,
504        _: &Self::ContentionState,
505        do_release: F,
506        _: &ChainLockTransactionToken<Self>,
507    ) {
508        do_release();
509    }
510}
511
512zr::static_assert!(size_of::<ChainLock<DummyTxU32>>() == 16);
513zr::static_assert!(align_of::<ChainLock<DummyTxU32>>() == 8);
514zr::static_assert!(offset_of!(ChainLock<DummyTxU32>, state) == 0);
515zr::static_assert!(offset_of!(ChainLock<DummyTxU32>, contention_state) == 8);
516
517zr::static_assert!(size_of::<ChainLock<DummyTxU64>>() == 16);
518zr::static_assert!(align_of::<ChainLock<DummyTxU64>>() == 8);
519zr::static_assert!(offset_of!(ChainLock<DummyTxU64>, state) == 0);
520zr::static_assert!(offset_of!(ChainLock<DummyTxU64>, contention_state) == 8);
521
522#[cfg(test)]
523mod tests {
524    use core::cell::Cell;
525    use core::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
526    use std::boxed::Box;
527    use std::sync::{Arc, Condvar, Mutex};
528    use std::thread;
529    use std::time::Duration;
530
531    use super::*;
532    use crate::chainlock_transaction::{
533        Action, ChainLockTransaction as BaseChainLockTransaction, DONE, Option as TxOption,
534        SelectAccounting, SelectTracing, StateOption, StateSaver, TransactionHooks,
535        TransactionResult,
536    };
537
538    const DEBUG_ENABLED: bool = true;
539    const TRACING_ENABLED: bool = true;
540
541    #[repr(C)]
542    #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
543    struct CallsiteInfo {
544        tag: &'static str,
545    }
546
547    const fn clt_tag(tag: &'static str) -> CallsiteInfo {
548        CallsiteInfo { tag }
549    }
550
551    #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
552    enum Options {
553        #[default]
554        None,
555        A,
556        B,
557    }
558
559    const OPTION_NONE: TxOption<Options> = TxOption::new(Options::None);
560    const OPTION_A: TxOption<Options> = TxOption::new(Options::A);
561    const OPTION_B: TxOption<Options> = TxOption::new(Options::B);
562
563    struct TestHooks;
564
565    type ChainLockTransaction = BaseChainLockTransaction<TestHooks>;
566    type ChainLockTransactionToken =
567        crate::chainlock_transaction::ChainLockTransactionToken<ChainLockTransaction>;
568    type ChainLock = super::ChainLock<ChainLockTransaction>;
569    type ChainLockGuard<'a> = crate::chainlock_guard::ChainLockGuard<'a, ChainLockTransaction>;
570
571    static RESERVED_TOKEN_ALLOCATOR: AtomicUsize = AtomicUsize::new(0);
572
573    std::thread_local! {
574        static LAST_CONTENTION_START_TICKS: Cell<i64> = const { Cell::new(0) };
575        static LAST_CALLSITE_INFO: Cell<CallsiteInfo> = const {
576            Cell::new(CallsiteInfo { tag: "" })
577        };
578        static ACTIVE_OPTION: Cell<Options> = const { Cell::new(Options::None) };
579        static ACTIVE_TRANSACTION: Cell<Option<NonNull<ChainLockTransaction>>> = const {
580            Cell::new(None)
581        };
582        static SIMULATED_TICKS: Cell<i64> = const { Cell::new(100) };
583    }
584
585    // SAFETY: `ACTIVE_TRANSACTION` is stored in thread-local storage (`std::thread_local!`), and
586    // `active()` returns the exact pointer most recently stored by `set_active()` on the current
587    // thread.
588    unsafe impl TransactionHooks for TestHooks {
589        type ContentionState = AtomicU64;
590        type CallsiteInfo = CallsiteInfo;
591        type Accounting = SelectAccounting<DEBUG_ENABLED>;
592        type Tracing = SelectTracing<TRACING_ENABLED, Self::CallsiteInfo>;
593        type ExtraState = ();
594
595        fn active() -> Option<NonNull<ChainLockTransaction>> {
596            ACTIVE_TRANSACTION.with(Cell::get)
597        }
598
599        unsafe fn set_active(transaction: Option<NonNull<ChainLockTransaction>>) {
600            ACTIVE_TRANSACTION.with(|cell| cell.set(transaction));
601        }
602
603        fn yield_cpu() {}
604
605        fn allocate_reserved_token() -> Token {
606            let reserved_index = RESERVED_TOKEN_ALLOCATOR.fetch_add(1, Ordering::Relaxed);
607            assert!((reserved_index as u64) < ChainLockTransaction::RESERVED_TOKEN_VALUES);
608            ChainLockTransaction::reserved_token(reserved_index)
609        }
610
611        fn get_current_ticks() -> i64 {
612            SIMULATED_TICKS.with(|cell| {
613                let ticks = cell.get();
614                cell.set(ticks + 1);
615                ticks
616            })
617        }
618
619        fn on_finalized(contention_start_ticks: i64, callsite_info: Self::CallsiteInfo) {
620            LAST_CONTENTION_START_TICKS.with(|cell| cell.set(contention_start_ticks));
621            LAST_CALLSITE_INFO.with(|cell| cell.set(callsite_info));
622        }
623
624        fn on_backoff(_transaction: &ChainLockTransaction) {}
625
626        fn record_backoff_or_retry(
627            _transaction: &ChainLockTransaction,
628            _observed_token: Token,
629            _lock_state: &AtomicToken,
630            _contention_state: &Self::ContentionState,
631        ) -> bool {
632            true
633        }
634
635        fn release_and_notify<F: FnOnce()>(
636            _transaction: &ChainLockTransaction,
637            _contention_state: &Self::ContentionState,
638            do_release: F,
639        ) {
640            do_release();
641        }
642    }
643
644    impl TestHooks {
645        fn active_option() -> Options {
646            ACTIVE_OPTION.with(Cell::get)
647        }
648
649        fn last_contention_start_ticks() -> i64 {
650            LAST_CONTENTION_START_TICKS.with(Cell::get)
651        }
652
653        fn last_callsite_info() -> CallsiteInfo {
654            LAST_CALLSITE_INFO.with(Cell::get)
655        }
656    }
657
658    struct TestStateSaver {
659        reset_on_drop: bool,
660    }
661
662    impl Drop for TestStateSaver {
663        fn drop(&mut self) {
664            if self.reset_on_drop {
665                ACTIVE_OPTION.with(|cell| cell.set(Options::None));
666            }
667        }
668    }
669
670    impl StateSaver<ChainLockTransaction> for TestStateSaver {
671        fn relax(&mut self, _transaction: &ChainLockTransaction) {
672            thread::yield_now();
673        }
674    }
675
676    impl StateOption<ChainLockTransaction> for Options {
677        type Saver = TestStateSaver;
678
679        fn save_state(&self) -> Self::Saver {
680            match self {
681                Options::None => TestStateSaver { reset_on_drop: false },
682                Options::A => {
683                    ACTIVE_OPTION.with(|cell| cell.set(Options::A));
684                    TestStateSaver { reset_on_drop: true }
685                }
686                Options::B => {
687                    ACTIVE_OPTION.with(|cell| cell.set(Options::B));
688                    TestStateSaver { reset_on_drop: true }
689                }
690            }
691        }
692    }
693
694    struct HelperShared<S> {
695        lock: Mutex<S>,
696        helper_condition: Condvar,
697        test_condition: Condvar,
698    }
699
700    struct Helper<S> {
701        shared: Arc<HelperShared<S>>,
702    }
703
704    impl<S: Copy + PartialEq> Helper<S> {
705        fn set_state_and_notify_test(&self, state: S) {
706            {
707                let mut guard = self.shared.lock.lock().unwrap();
708                *guard = state;
709            }
710            self.shared.test_condition.notify_one();
711        }
712
713        fn wait_for_test_state(&self, state: S) {
714            let mut guard = self.shared.lock.lock().unwrap();
715            while *guard != state {
716                guard = self.shared.helper_condition.wait(guard).unwrap();
717            }
718        }
719    }
720
721    struct HelperThread<S> {
722        shared: Arc<HelperShared<S>>,
723        helper_thread: Option<thread::JoinHandle<()>>,
724    }
725
726    impl<S: Copy + PartialEq + Send + 'static> HelperThread<S> {
727        fn new<F>(initial: S, callable: F) -> Self
728        where
729            F: FnOnce(Helper<S>) + Send + 'static,
730        {
731            let shared = Arc::new(HelperShared {
732                lock: Mutex::new(initial),
733                helper_condition: Condvar::new(),
734                test_condition: Condvar::new(),
735            });
736            let helper = Helper { shared: Arc::clone(&shared) };
737            let handle = thread::spawn(move || callable(helper));
738            Self { shared, helper_thread: Some(handle) }
739        }
740
741        fn set_state_and_notify_helper(&self, state: S) {
742            {
743                let mut guard = self.shared.lock.lock().unwrap();
744                *guard = state;
745            }
746            self.shared.helper_condition.notify_one();
747        }
748
749        fn wait_for_helper_state(&self, state: S) {
750            let mut guard = self.shared.lock.lock().unwrap();
751            while *guard != state {
752                guard = self.shared.test_condition.wait(guard).unwrap();
753            }
754        }
755    }
756
757    impl<S> Drop for HelperThread<S> {
758        fn drop(&mut self) {
759            if let Some(handle) = self.helper_thread.take() {
760                handle.join().unwrap();
761            }
762        }
763    }
764
765    #[test]
766    fn transaction_options() {
767        assert_eq!(Options::None, TestHooks::active_option());
768
769        ChainLockTransaction::until_done(
770            OPTION_NONE,
771            clt_tag("TransactionOptions"),
772            |_token| -> TransactionResult<()> {
773                assert_eq!(Options::None, TestHooks::active_option());
774                DONE.into()
775            },
776        );
777
778        assert_eq!(Options::None, TestHooks::active_option());
779
780        ChainLockTransaction::until_done(
781            OPTION_A,
782            clt_tag("TransactionOptions"),
783            |_token| -> TransactionResult<()> {
784                assert_eq!(Options::A, TestHooks::active_option());
785                DONE.into()
786            },
787        );
788
789        assert_eq!(Options::None, TestHooks::active_option());
790
791        ChainLockTransaction::until_done(
792            OPTION_B,
793            clt_tag("TransactionOptions"),
794            |_token| -> TransactionResult<()> {
795                assert_eq!(Options::B, TestHooks::active_option());
796                DONE.into()
797            },
798        );
799
800        assert_eq!(Options::None, TestHooks::active_option());
801    }
802
803    #[test]
804    fn transaction_return_value() {
805        let mut value = Some(Box::new(42));
806        let value_pointer: *const i32 = &**value.as_ref().unwrap();
807
808        let return_value: Box<i32> = ChainLockTransaction::until_done(
809            OPTION_NONE,
810            clt_tag("TransactionOptions"),
811            |_token| -> TransactionResult<Box<i32>> {
812                TransactionResult::Value(value.take().unwrap())
813            },
814        );
815
816        assert_eq!(&*return_value as *const i32, value_pointer);
817        assert!(value.is_none());
818    }
819
820    #[test]
821    fn uncontested_acquire() {
822        let lock = ChainLock::new();
823
824        ChainLockTransaction::until_done(
825            OPTION_NONE,
826            clt_tag("UncontestedAcquire"),
827            |token| -> TransactionResult<()> {
828                assert!(!lock.is_held(token));
829                assert!(lock.is_unlocked());
830
831                lock.acquire_first_in_chain(token);
832
833                assert!(lock.is_held(token));
834                assert!(!lock.is_unlocked());
835                assert_eq!(
836                    lock.state_for_testing(),
837                    ChainLockTransaction::active_ref(token).token()
838                );
839
840                lock.release(token);
841
842                assert!(!lock.is_held(token));
843                assert!(lock.is_unlocked());
844                DONE.into()
845            },
846        );
847    }
848
849    #[test]
850    fn cyclic_acquire() {
851        let lock = ChainLock::new();
852        let get_lock = || -> &ChainLock { &lock };
853
854        ChainLockTransaction::until_done(
855            OPTION_NONE,
856            clt_tag("CyclicAcquire"),
857            |token| -> TransactionResult<()> {
858                lock.acquire_first_in_chain(token);
859
860                assert_eq!(Err(Result::Cycle), get_lock().acquire_or_result(token));
861
862                lock.release(token);
863                DONE.into()
864            },
865        );
866    }
867
868    #[test]
869    fn contention_priority() {
870        let locks = Arc::new((ChainLock::new(), ChainLock::new()));
871
872        #[derive(Clone, Copy, Debug, PartialEq, Eq)]
873        enum State {
874            Initial,
875            HelperTransactionStarted,
876            TestTransactionStarted,
877        }
878
879        let helper_locks = Arc::clone(&locks);
880        let helper_thread = HelperThread::new(State::Initial, move |helper| {
881            let (lock_a, lock_b) = &*helper_locks;
882            ChainLockTransaction::until_done(
883                OPTION_NONE,
884                clt_tag("ContentionPriority helper"),
885                |token| -> TransactionResult<()> {
886                    let _guard = ChainLockGuard::new(lock_a, token);
887
888                    helper.set_state_and_notify_test(State::HelperTransactionStarted);
889                    helper.wait_for_test_state(State::TestTransactionStarted);
890
891                    // Acquiring lock_b should spin until the lower priority transaction backs off.
892                    let acquired = lock_b.acquire_or_backoff(token);
893                    if acquired {
894                        lock_b.release(token);
895                    } else {
896                        assert!(acquired);
897                    }
898
899                    DONE.into()
900                },
901            );
902        });
903
904        // Wait for the helper thread to start a transaction so it has higher priority.
905        helper_thread.wait_for_helper_state(State::HelperTransactionStarted);
906
907        let (lock_a, lock_b) = &*locks;
908        ChainLockTransaction::until_done(
909            OPTION_NONE,
910            clt_tag("ContentionPriority"),
911            |token| -> TransactionResult<()> {
912                let _guard = ChainLockGuard::new(lock_b, token);
913
914                helper_thread.set_state_and_notify_helper(State::TestTransactionStarted);
915
916                // The helper transaction should be higher priority than this transaction.
917                assert!(lock_b.state_for_testing() > lock_a.state_for_testing());
918
919                // Let the helper thread spin for a while waiting for lock_b to be released.
920                thread::sleep(Duration::from_millis(50));
921
922                if !lock_a.acquire_or_backoff(token) {
923                    return Action::Backoff.into();
924                }
925                lock_a.release(token);
926
927                DONE.into()
928            },
929        );
930
931        // Verify contention tracing recorded the contention on the test thread.
932        assert!(TestHooks::last_contention_start_ticks() > 0);
933        assert_eq!(TestHooks::last_callsite_info(), clt_tag("ContentionPriority"));
934    }
935
936    struct TypeA {
937        lock: ChainLock,
938        value: i32,
939    }
940
941    impl TypeA {
942        fn new(value: i32) -> Self {
943            Self { lock: ChainLock::new(), value }
944        }
945
946        fn get_lock(&self) -> &ChainLock {
947            &self.lock
948        }
949
950        fn compare(&self, other: &TypeB, token: &ChainLockTransactionToken) -> bool {
951            self.get_lock().assert_held(token);
952            other.get_lock().assert_held(token);
953            self.value == other.value
954        }
955    }
956
957    struct TypeB {
958        lock: ChainLock,
959        value: i32,
960    }
961
962    impl TypeB {
963        fn new(value: i32) -> Self {
964            Self { lock: ChainLock::new(), value }
965        }
966
967        fn get_lock(&self) -> &ChainLock {
968            &self.lock
969        }
970
971        fn compare(&self, other: &TypeA, token: &ChainLockTransactionToken) -> bool {
972            self.get_lock().assert_held(token);
973            other.get_lock().assert_held(token);
974            self.value == other.value
975        }
976    }
977
978    #[test]
979    fn cross_referencing_types() {
980        let value_a = TypeA::new(0);
981        let value_b = TypeB::new(0);
982
983        let do_transaction = |token: &ChainLockTransactionToken| -> TransactionResult<()> {
984            let _guard_a = ChainLockGuard::new(value_a.get_lock(), token);
985            let mut guard_b = ChainLockGuard::defer(value_b.get_lock(), token);
986            if !guard_b.acquire_or_backoff() {
987                return Action::Backoff.into();
988            }
989            assert!(value_a.compare(&value_b, token));
990            assert!(value_b.compare(&value_a, token));
991            DONE.into()
992        };
993
994        ChainLockTransaction::until_done(
995            OPTION_NONE,
996            clt_tag("CrossReferencingTypes"),
997            do_transaction,
998        );
999    }
1000
1001    #[test]
1002    fn guard_operations_and_acquire_both() {
1003        let lock_a = ChainLock::new();
1004        let lock_b = ChainLock::new();
1005
1006        ChainLockTransaction::until_done(OPTION_NONE, clt_tag("GuardOps"), |token| {
1007            assert!(acquire_both_or_backoff(&lock_a, &lock_b, token));
1008            let guard_a = ChainLockGuard::adopt(&lock_a, token);
1009            let mut guard_b = ChainLockGuard::adopt(&lock_b, token);
1010            assert_eq!(ChainLockTransaction::active_ref(token).locks_held(), 2);
1011
1012            // Transfer guard_a to a new guard via `take` and `from_taken`.
1013            let mut transferred_a = ChainLockGuard::from_taken(guard_a.take(), &lock_a);
1014            assert!(transferred_a.is_held());
1015
1016            // Test `unguard` + manual release on lock_a.
1017            transferred_a.unguard();
1018            assert!(!transferred_a.is_held());
1019            lock_a.release(token);
1020            assert_eq!(ChainLockTransaction::active_ref(token).locks_held(), 1);
1021
1022            // Explicitly release guard_b and re-acquire with `try_acquire_with`.
1023            guard_b.release();
1024            assert_eq!(ChainLockTransaction::active_ref(token).locks_held(), 0);
1025            assert!(guard_b.try_acquire_with(|l| l.try_acquire(token)));
1026            assert_eq!(ChainLockTransaction::active_ref(token).locks_held(), 1);
1027
1028            DONE
1029        });
1030    }
1031
1032    #[test]
1033    fn save_and_restore_state_and_restart() {
1034        let lock = ChainLock::new();
1035        let mut restarted = false;
1036
1037        ChainLockTransaction::until_done(
1038            OPTION_NONE,
1039            clt_tag("SaveRestoreAndRestart"),
1040            |token| -> TransactionResult<()> {
1041                if !restarted {
1042                    lock.acquire_first_in_chain(token);
1043                    ChainLockTransaction::finalize(token);
1044                    assert!(ChainLockTransaction::active_ref(token).is_finalized());
1045
1046                    let original_token = ChainLockTransaction::active_ref(token).token();
1047                    let saved = ChainLockTransaction::save_state_and_use_reserved_token(token);
1048                    assert_eq!(saved.token(), original_token);
1049                    assert!(ChainLockTransaction::active_ref(token).token().is_reserved());
1050
1051                    // Sync lock to the reserved token, then restore the original token and sync
1052                    // back.
1053                    lock.sync_token(token);
1054                    assert!(lock.state_for_testing().is_reserved());
1055
1056                    // SAFETY: `saved` was created from the currently active transaction on this
1057                    // stack.
1058                    unsafe {
1059                        ChainLockTransaction::restore_state(saved);
1060                    }
1061                    assert_eq!(ChainLockTransaction::active_ref(token).token(), original_token);
1062                    lock.sync_token(token);
1063                    assert_eq!(lock.state_for_testing(), original_token);
1064
1065                    lock.release(token);
1066                    restarted = true;
1067                    return Action::Restart.into();
1068                }
1069
1070                assert!(!ChainLockTransaction::active_ref(token).is_finalized());
1071                DONE.into()
1072            },
1073        );
1074        assert!(restarted);
1075    }
1076}