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

starnix_sync/
thread_affinity.rs
