starnix_sync/
thread_affinity.rs

1// Copyright 2024 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
5#[cfg(feature = "detect_lock_dep_cycles")]
6mod tracking {
7    use std::sync::atomic::{AtomicUsize, Ordering};
8
9    static NEXT_THREAD_ID: AtomicUsize = AtomicUsize::new(1);
10
11    thread_local! {
12        static THREAD_ID: usize = NEXT_THREAD_ID.fetch_add(1, Ordering::Relaxed);
13    }
14
15    fn current_thread_id() -> usize {
16        THREAD_ID.with(|id| *id)
17    }
18
19    pub struct ThreadAffinityState {
20        owner: AtomicUsize,
21    }
22
23    impl ThreadAffinityState {
24        pub const fn new() -> Self {
25            Self { owner: AtomicUsize::new(0) }
26        }
27
28        #[inline(always)]
29        pub fn attach(&self) -> ThreadAffinityToken<'_> {
30            let id = current_thread_id();
31            let previous = self.owner.swap(id, Ordering::Relaxed);
32            assert_eq!(previous, 0, "ThreadAffinity: Object is already attached to a thread!");
33            ThreadAffinityToken { state: self }
34        }
35
36        #[inline(always)]
37        pub fn assert_attached(&self) {
38            let id = current_thread_id();
39            assert_eq!(
40                self.owner.load(Ordering::Relaxed),
41                id,
42                "ThreadAffinity: Object is not attached to the current thread!"
43            );
44        }
45
46        #[inline(always)]
47        pub fn assert_not_attached(&self) {
48            let id = current_thread_id();
49            assert_ne!(
50                self.owner.load(Ordering::Relaxed),
51                id,
52                "ThreadAffinity: Object is already attached to the current thread!"
53            );
54        }
55    }
56
57    pub struct ThreadAffinityToken<'a> {
58        state: &'a ThreadAffinityState,
59    }
60
61    impl<'a> Drop for ThreadAffinityToken<'a> {
62        fn drop(&mut self) {
63            self.state.owner.store(0, Ordering::Relaxed);
64        }
65    }
66}
67
68#[cfg(not(feature = "detect_lock_dep_cycles"))]
69mod tracking {
70    pub struct ThreadAffinityState {}
71
72    impl ThreadAffinityState {
73        #[inline(always)]
74        pub const fn new() -> Self {
75            Self {}
76        }
77
78        #[inline(always)]
79        pub fn attach(&self) -> ThreadAffinityToken<'_> {
80            ThreadAffinityToken { _state: self }
81        }
82
83        #[inline(always)]
84        pub fn assert_attached(&self) {}
85
86        #[inline(always)]
87        pub fn assert_not_attached(&self) {}
88    }
89
90    pub struct ThreadAffinityToken<'a> {
91        _state: &'a ThreadAffinityState,
92    }
93}
94
95/// A synchronization primitive that tracks thread affinity.
96///
97/// It provides mechanisms to attach an object to a thread and assert that
98/// the object is or is not attached to the current thread. When the feature
99/// `detect_lock_dep_cycles` is disabled, this struct and its methods are zero-cost.
100pub struct ThreadAffinity {
101    state: tracking::ThreadAffinityState,
102}
103
104impl ThreadAffinity {
105    pub const fn new() -> Self {
106        Self { state: tracking::ThreadAffinityState::new() }
107    }
108
109    /// Attaches this object to the current thread.
110    ///
111    /// Returns a `ThreadAffinityGuard` that will detach the object when dropped.
112    #[inline(always)]
113    pub fn attach(&self) -> ThreadAffinityGuard<'_> {
114        ThreadAffinityGuard { _token: self.state.attach() }
115    }
116
117    /// Asserts that this object is attached to the current thread.
118    ///
119    /// Panics if it is not attached.
120    #[inline(always)]
121    pub fn assert_attached(&self) {
122        self.state.assert_attached();
123    }
124
125    /// Asserts that this object is NOT attached to the current thread.
126    ///
127    /// Panics if it is attached.
128    #[inline(always)]
129    pub fn assert_not_attached(&self) {
130        self.state.assert_not_attached();
131    }
132}
133
134impl Default for ThreadAffinity {
135    fn default() -> Self {
136        Self::new()
137    }
138}
139
140pub struct ThreadAffinityGuard<'a> {
141    _token: tracking::ThreadAffinityToken<'a>,
142}
starnix_sync/thread_affinity.rs

starnix_sync/
thread_affinity.rs
