cordyceps/
loom.rs

1pub(crate) use self::inner::*;
2
3#[cfg(loom)]
4mod inner {
5    #![allow(unused_imports)]
6
7    pub(crate) mod atomic {
8        pub use core::sync::atomic::Ordering;
9        pub use loom::sync::atomic::*;
10    }
11
12    pub(crate) use loom::{cell, hint, model, sync, thread};
13
14    pub(crate) mod alloc {
15        #![allow(dead_code)]
16        use core::fmt;
17        use loom::alloc;
18        /// Track allocations, detecting leaks
19        ///
20        /// This is a version of `loom::alloc::Track` that adds a missing
21        /// `Default` impl.
22        pub struct Track<T>(alloc::Track<T>);
23
24        impl<T> Track<T> {
25            /// Track a value for leaks
26            #[inline(always)]
27            pub fn new(value: T) -> Track<T> {
28                Track(alloc::Track::new(value))
29            }
30
31            /// Get a reference to the value
32            #[inline(always)]
33            pub fn get_ref(&self) -> &T {
34                self.0.get_ref()
35            }
36
37            /// Get a mutable reference to the value
38            #[inline(always)]
39            pub fn get_mut(&mut self) -> &mut T {
40                self.0.get_mut()
41            }
42
43            /// Stop tracking the value for leaks
44            #[inline(always)]
45            pub fn into_inner(self) -> T {
46                self.0.into_inner()
47            }
48        }
49
50        impl<T: fmt::Debug> fmt::Debug for Track<T> {
51            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52                self.0.fmt(f)
53            }
54        }
55
56        impl<T: Default> Default for Track<T> {
57            fn default() -> Self {
58                Self::new(T::default())
59            }
60        }
61    }
62}
63
64#[cfg(not(loom))]
65mod inner {
66    #![allow(dead_code, unused_imports)]
67    pub(crate) mod sync {
68        pub use core::sync::*;
69
70        #[cfg(all(feature = "alloc", not(test)))]
71        pub use alloc::sync::*;
72
73        #[cfg(test)]
74        pub use std::sync::*;
75    }
76
77    pub(crate) use core::sync::atomic;
78
79    #[cfg(test)]
80    pub(crate) mod thread {
81        pub(crate) use std::thread::{yield_now, JoinHandle};
82        pub(crate) fn spawn<F, T>(f: F) -> JoinHandle<T>
83        where
84            F: FnOnce() -> T + Send + 'static,
85            T: Send + 'static,
86        {
87            let track = super::alloc::track::Registry::current();
88            std::thread::spawn(move || {
89                let _tracking = track.map(|track| track.set_default());
90                f()
91            })
92        }
93    }
94    pub(crate) mod hint {
95        #[inline(always)]
96        pub(crate) fn spin_loop() {
97            // MSRV: std::hint::spin_loop() stabilized in 1.49.0
98            #[allow(deprecated)]
99            super::atomic::spin_loop_hint()
100        }
101    }
102
103    pub(crate) mod cell {
104        #[derive(Debug)]
105        pub(crate) struct UnsafeCell<T>(core::cell::UnsafeCell<T>);
106
107        impl<T> UnsafeCell<T> {
108            pub const fn new(data: T) -> UnsafeCell<T> {
109                UnsafeCell(core::cell::UnsafeCell::new(data))
110            }
111
112            #[inline(always)]
113            pub fn with<F, R>(&self, f: F) -> R
114            where
115                F: FnOnce(*const T) -> R,
116            {
117                f(self.0.get())
118            }
119
120            #[inline(always)]
121            pub fn with_mut<F, R>(&self, f: F) -> R
122            where
123                F: FnOnce(*mut T) -> R,
124            {
125                f(self.0.get())
126            }
127
128            #[inline(always)]
129            pub(crate) fn get_mut(&self) -> MutPtr<T> {
130                MutPtr(self.0.get())
131            }
132        }
133
134        #[derive(Debug)]
135        pub(crate) struct MutPtr<T: ?Sized>(*mut T);
136
137        impl<T: ?Sized> MutPtr<T> {
138            // Clippy knows that it's Bad and Wrong to construct a mutable reference
139            // from an immutable one...but this function is intended to simulate a raw
140            // pointer, so we have to do that here.
141            #[allow(clippy::mut_from_ref)]
142            #[inline(always)]
143            pub(crate) unsafe fn deref(&self) -> &mut T {
144                &mut *self.0
145            }
146
147            #[inline(always)]
148            pub fn with<F, R>(&self, f: F) -> R
149            where
150                F: FnOnce(*mut T) -> R,
151            {
152                f(self.0)
153            }
154        }
155    }
156
157    #[cfg(test)]
158    pub(crate) mod model {
159        #[non_exhaustive]
160        #[derive(Default)]
161        pub(crate) struct Builder {
162            pub(crate) max_threads: usize,
163            pub(crate) max_branches: usize,
164            pub(crate) max_permutations: Option<usize>,
165            // pub(crate) max_duration: Option<Duration>,
166            pub(crate) preemption_bound: Option<usize>,
167            // pub(crate) checkpoint_file: Option<PathBuf>,
168            pub(crate) checkpoint_interval: usize,
169            pub(crate) location: bool,
170            pub(crate) log: bool,
171        }
172
173        impl Builder {
174            pub(crate) fn new() -> Self {
175                Self::default()
176            }
177
178            pub(crate) fn check(&self, f: impl FnOnce()) {
179                let registry = super::alloc::track::Registry::default();
180                let _tracking = registry.set_default();
181                f();
182                registry.check();
183            }
184        }
185    }
186
187    #[cfg(test)]
188    pub(crate) fn model(f: impl FnOnce()) {
189        let collector = tracing_subscriber::fmt()
190            .with_max_level(tracing::Level::TRACE)
191            .with_test_writer()
192            .without_time()
193            .with_thread_ids(true)
194            .with_thread_names(false)
195            .finish();
196        let _ = tracing::subscriber::set_global_default(collector);
197        model::Builder::new().check(f)
198    }
199
200    pub(crate) mod alloc {
201        #[cfg(test)]
202        use std::sync::Arc;
203
204        #[cfg(test)]
205        pub(in crate::loom) mod track {
206            use std::{
207                cell::RefCell,
208                sync::{
209                    atomic::{AtomicBool, Ordering},
210                    Arc, Mutex, Weak,
211                },
212            };
213
214            #[derive(Clone, Debug, Default)]
215            pub(crate) struct Registry(Arc<Mutex<RegistryInner>>);
216
217            #[derive(Debug, Default)]
218            struct RegistryInner {
219                tracks: Vec<Weak<TrackData>>,
220                next_id: usize,
221            }
222
223            #[derive(Debug)]
224            pub(super) struct TrackData {
225                was_leaked: AtomicBool,
226                type_name: &'static str,
227                location: &'static core::panic::Location<'static>,
228                id: usize,
229            }
230
231            thread_local! {
232                static REGISTRY: RefCell<Option<Registry>> = const { RefCell::new(None) };
233            }
234
235            impl Registry {
236                pub(in crate::loom) fn current() -> Option<Registry> {
237                    REGISTRY.with(|current| current.borrow().clone())
238                }
239
240                pub(in crate::loom) fn set_default(&self) -> impl Drop {
241                    struct Unset(Option<Registry>);
242                    impl Drop for Unset {
243                        fn drop(&mut self) {
244                            let _ =
245                                REGISTRY.try_with(|current| *current.borrow_mut() = self.0.take());
246                        }
247                    }
248
249                    REGISTRY.with(|current| {
250                        let mut current = current.borrow_mut();
251                        let unset = Unset(current.clone());
252                        *current = Some(self.clone());
253                        unset
254                    })
255                }
256
257                #[track_caller]
258                pub(super) fn start_tracking<T>() -> Option<Arc<TrackData>> {
259                    // we don't use `Option::map` here because it creates a
260                    // closure, which breaks `#[track_caller]`, since the caller
261                    // of `insert` becomes the closure, which cannot have a
262                    // `#[track_caller]` attribute on it.
263                    #[allow(clippy::manual_map)]
264                    match Self::current() {
265                        Some(registry) => Some(registry.insert::<T>()),
266                        _ => None,
267                    }
268                }
269
270                #[track_caller]
271                pub(super) fn insert<T>(&self) -> Arc<TrackData> {
272                    let mut inner = self.0.lock().unwrap();
273                    let id = inner.next_id;
274                    inner.next_id += 1;
275                    let location = core::panic::Location::caller();
276                    let type_name = std::any::type_name::<T>();
277                    let data = Arc::new(TrackData {
278                        type_name,
279                        location,
280                        id,
281                        was_leaked: AtomicBool::new(false),
282                    });
283                    let weak = Arc::downgrade(&data);
284                    test_trace!(
285                        target: "maitake::alloc",
286                        id,
287                        "type" = %type_name,
288                        %location,
289                        "started tracking allocation",
290                    );
291                    inner.tracks.push(weak);
292                    data
293                }
294
295                pub(in crate::loom) fn check(&self) {
296                    let leaked = self
297                        .0
298                        .lock()
299                        .unwrap()
300                        .tracks
301                        .iter()
302                        .filter_map(|weak| {
303                            let data = weak.upgrade()?;
304                            data.was_leaked.store(true, Ordering::SeqCst);
305                            Some(format!(
306                                " - id {}, {} allocated at {}",
307                                data.id, data.type_name, data.location
308                            ))
309                        })
310                        .collect::<Vec<_>>();
311                    if !leaked.is_empty() {
312                        let leaked = leaked.join("\n  ");
313                        panic!("the following allocations were leaked:\n  {leaked}");
314                    }
315                }
316            }
317
318            impl Drop for TrackData {
319                fn drop(&mut self) {
320                    if !self.was_leaked.load(Ordering::SeqCst) {
321                        test_trace!(
322                            target: "maitake::alloc",
323                            id = self.id,
324                            "type" = %self.type_name,
325                            location = %self.location,
326                            "dropped all references to a tracked allocation",
327                        );
328                    }
329                }
330            }
331        }
332
333        /// Track allocations, detecting leaks
334        #[derive(Debug, Default)]
335        pub struct Track<T> {
336            value: T,
337
338            #[cfg(test)]
339            track: Option<Arc<track::TrackData>>,
340        }
341
342        impl<T> Track<T> {
343            pub const fn new_const(value: T) -> Track<T> {
344                Track {
345                    value,
346
347                    #[cfg(test)]
348                    track: None,
349                }
350            }
351
352            /// Track a value for leaks
353            #[inline(always)]
354            #[track_caller]
355            pub fn new(value: T) -> Track<T> {
356                Track {
357                    value,
358
359                    #[cfg(test)]
360                    track: track::Registry::start_tracking::<T>(),
361                }
362            }
363
364            /// Get a reference to the value
365            #[inline(always)]
366            pub fn get_ref(&self) -> &T {
367                &self.value
368            }
369
370            /// Get a mutable reference to the value
371            #[inline(always)]
372            pub fn get_mut(&mut self) -> &mut T {
373                &mut self.value
374            }
375
376            /// Stop tracking the value for leaks
377            #[inline(always)]
378            pub fn into_inner(self) -> T {
379                self.value
380            }
381        }
382    }
383}