maitake/
loom.rs

1#[allow(unused_imports)]
2pub(crate) use self::inner::*;
3
4#[cfg(loom)]
5mod inner {
6    #![allow(dead_code)]
7    pub(crate) use loom::thread_local;
8
9    #[cfg(feature = "alloc")]
10    pub(crate) mod alloc {
11        use super::sync::Arc;
12        use core::{
13            future::Future,
14            pin::Pin,
15            task::{Context, Poll},
16        };
17        pub(crate) use loom::alloc::*;
18
19        #[derive(Debug)]
20        #[pin_project::pin_project]
21        pub(crate) struct TrackFuture<F> {
22            #[pin]
23            inner: F,
24            track: Arc<()>,
25        }
26
27        impl<F: Future> Future for TrackFuture<F> {
28            type Output = TrackFuture<F::Output>;
29            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
30                let this = self.project();
31                this.inner.poll(cx).map(|inner| TrackFuture {
32                    inner,
33                    track: this.track.clone(),
34                })
35            }
36        }
37
38        impl<F> TrackFuture<F> {
39            /// Wrap a `Future` in a `TrackFuture` that participates in Loom's
40            /// leak checking.
41            #[track_caller]
42            pub(crate) fn new(inner: F) -> Self {
43                Self {
44                    inner,
45                    track: Arc::new(()),
46                }
47            }
48
49            /// Stop tracking this future, and return the inner value.
50            pub(crate) fn into_inner(self) -> F {
51                self.inner
52            }
53        }
54
55        #[track_caller]
56        pub(crate) fn track_future<F: Future>(inner: F) -> TrackFuture<F> {
57            TrackFuture::new(inner)
58        }
59
60        // PartialEq impl so that `assert_eq!(..., Ok(...))` works
61        impl<F: PartialEq> PartialEq for TrackFuture<F> {
62            fn eq(&self, other: &Self) -> bool {
63                self.inner == other.inner
64            }
65        }
66    }
67
68    pub(crate) use loom::{cell, future, model, thread};
69
70    pub(crate) mod sync {
71        pub(crate) use loom::sync::*;
72
73        pub(crate) mod spin {
74            pub(crate) use loom::sync::MutexGuard;
75
76            /// Mock version of mycelium's spinlock, but using
77            /// `loom::sync::Mutex`. The API is slightly different, since the
78            /// mycelium mutex does not support poisoning.
79            #[derive(Debug)]
80            pub(crate) struct Mutex<T>(loom::sync::Mutex<T>);
81
82            impl<T> Mutex<T> {
83                #[track_caller]
84                pub(crate) fn new(t: T) -> Self {
85                    Self(loom::sync::Mutex::new(t))
86                }
87
88                #[track_caller]
89                pub fn try_lock(&self) -> Option<MutexGuard<'_, T>> {
90                    self.0.try_lock().ok()
91                }
92
93                #[track_caller]
94                pub fn lock(&self) -> MutexGuard<'_, T> {
95                    self.0.lock().expect("loom mutex will never poison")
96                }
97            }
98        }
99    }
100}
101
102#[cfg(not(loom))]
103mod inner {
104    #![allow(dead_code, unused_imports)]
105
106    #[cfg(test)]
107    pub(crate) use std::thread_local;
108
109    pub(crate) mod sync {
110        #[cfg(feature = "alloc")]
111        pub use alloc::sync::*;
112        pub use core::sync::*;
113
114        pub use mycelium_util::sync::spin;
115    }
116
117    pub(crate) mod atomic {
118        pub use portable_atomic::*;
119    }
120
121    pub(crate) use portable_atomic::hint;
122
123    #[cfg(test)]
124    pub(crate) mod thread {
125
126        pub(crate) use std::thread::{yield_now, JoinHandle};
127
128        pub(crate) fn spawn<F, T>(f: F) -> JoinHandle<T>
129        where
130            F: FnOnce() -> T + Send + 'static,
131            T: Send + 'static,
132        {
133            use super::atomic::{AtomicUsize, Ordering::Relaxed};
134            thread_local! {
135                static CHILDREN: AtomicUsize = const { AtomicUsize::new(1) };
136            }
137
138            let track = super::alloc::track::Registry::current();
139            let subscriber = tracing_02::Dispatch::default();
140            let span = tracing_02::Span::current();
141            let num = CHILDREN.with(|children| children.fetch_add(1, Relaxed));
142            std::thread::spawn(move || {
143                let _tracing = tracing_02::dispatch::set_default(&subscriber);
144                let _span =
145                    tracing_02::info_span!(parent: span.id(), "thread", message = num).entered();
146
147                tracing_02::info!(num, "spawned child thread");
148                let _tracking = track.map(|track| track.set_default());
149                let res = f();
150                tracing_02::info!(num, "child thread completed");
151
152                res
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 _trace = crate::util::test::trace_init();
180                let _span = tracing_02::info_span!(
181                    "test",
182                    message = std::thread::current().name().unwrap_or("<unnamed>")
183                )
184                .entered();
185                let registry = super::alloc::track::Registry::default();
186                let _tracking = registry.set_default();
187
188                tracing_02::info!("started test...");
189                f();
190                tracing_02::info!("test completed successfully!");
191
192                registry.check();
193            }
194        }
195    }
196
197    #[cfg(test)]
198    pub(crate) fn model(f: impl FnOnce()) {
199        model::Builder::new().check(f)
200    }
201
202    pub(crate) mod cell {
203        #[derive(Debug)]
204        pub(crate) struct UnsafeCell<T: ?Sized>(core::cell::UnsafeCell<T>);
205
206        impl<T> UnsafeCell<T> {
207            pub const fn new(data: T) -> UnsafeCell<T> {
208                UnsafeCell(core::cell::UnsafeCell::new(data))
209            }
210        }
211
212        impl<T: ?Sized> UnsafeCell<T> {
213            #[inline(always)]
214            pub fn with<F, R>(&self, f: F) -> R
215            where
216                F: FnOnce(*const T) -> R,
217            {
218                f(self.0.get())
219            }
220
221            #[inline(always)]
222            pub fn with_mut<F, R>(&self, f: F) -> R
223            where
224                F: FnOnce(*mut T) -> R,
225            {
226                f(self.0.get())
227            }
228
229            #[inline(always)]
230            pub(crate) fn get(&self) -> ConstPtr<T> {
231                ConstPtr(self.0.get())
232            }
233
234            #[inline(always)]
235            pub(crate) fn get_mut(&self) -> MutPtr<T> {
236                MutPtr(self.0.get())
237            }
238        }
239
240        #[derive(Debug)]
241        pub(crate) struct ConstPtr<T: ?Sized>(*const T);
242
243        impl<T: ?Sized> ConstPtr<T> {
244            #[inline(always)]
245            pub(crate) unsafe fn deref(&self) -> &T {
246                &*self.0
247            }
248
249            #[inline(always)]
250            pub fn with<F, R>(&self, f: F) -> R
251            where
252                F: FnOnce(*const T) -> R,
253            {
254                f(self.0)
255            }
256        }
257
258        #[derive(Debug)]
259        pub(crate) struct MutPtr<T: ?Sized>(*mut T);
260
261        impl<T: ?Sized> MutPtr<T> {
262            // Clippy knows that it's Bad and Wrong to construct a mutable reference
263            // from an immutable one...but this function is intended to simulate a raw
264            // pointer, so we have to do that here.
265            #[allow(clippy::mut_from_ref)]
266            #[inline(always)]
267            pub(crate) unsafe fn deref(&self) -> &mut T {
268                &mut *self.0
269            }
270
271            #[inline(always)]
272            pub fn with<F, R>(&self, f: F) -> R
273            where
274                F: FnOnce(*mut T) -> R,
275            {
276                f(self.0)
277            }
278        }
279    }
280
281    pub(crate) mod alloc {
282        #[cfg(test)]
283        use core::{
284            future::Future,
285            pin::Pin,
286            task::{Context, Poll},
287        };
288
289        #[cfg(test)]
290        use std::sync::Arc;
291        #[cfg(test)]
292        pub(in crate::loom) mod track {
293            use std::{
294                cell::RefCell,
295                sync::{
296                    atomic::{AtomicBool, Ordering},
297                    Arc, Mutex, Weak,
298                },
299            };
300
301            #[derive(Clone, Debug, Default)]
302            pub(crate) struct Registry(Arc<Mutex<RegistryInner>>);
303
304            #[derive(Debug, Default)]
305            struct RegistryInner {
306                tracks: Vec<Weak<TrackData>>,
307                next_id: usize,
308            }
309
310            #[derive(Debug)]
311            pub(super) struct TrackData {
312                was_leaked: AtomicBool,
313                type_name: &'static str,
314                location: &'static core::panic::Location<'static>,
315                id: usize,
316            }
317
318            thread_local! {
319                static REGISTRY: RefCell<Option<Registry>> = const { RefCell::new(None) };
320            }
321
322            impl Registry {
323                pub(in crate::loom) fn current() -> Option<Registry> {
324                    REGISTRY.with(|current| current.borrow().clone())
325                }
326
327                pub(in crate::loom) fn set_default(&self) -> impl Drop {
328                    struct Unset(Option<Registry>);
329                    impl Drop for Unset {
330                        fn drop(&mut self) {
331                            let _ =
332                                REGISTRY.try_with(|current| *current.borrow_mut() = self.0.take());
333                        }
334                    }
335
336                    REGISTRY.with(|current| {
337                        let mut current = current.borrow_mut();
338                        let unset = Unset(current.clone());
339                        *current = Some(self.clone());
340                        unset
341                    })
342                }
343
344                #[track_caller]
345                pub(super) fn start_tracking<T>() -> Option<Arc<TrackData>> {
346                    // we don't use `Option::map` here because it creates a
347                    // closure, which breaks `#[track_caller]`, since the caller
348                    // of `insert` becomes the closure, which cannot have a
349                    // `#[track_caller]` attribute on it.
350                    #[allow(clippy::manual_map)]
351                    match Self::current() {
352                        Some(registry) => Some(registry.insert::<T>()),
353                        _ => None,
354                    }
355                }
356
357                #[track_caller]
358                pub(super) fn insert<T>(&self) -> Arc<TrackData> {
359                    let mut inner = self.0.lock().unwrap();
360                    let id = inner.next_id;
361                    inner.next_id += 1;
362                    let location = core::panic::Location::caller();
363                    let type_name = std::any::type_name::<T>();
364                    let data = Arc::new(TrackData {
365                        type_name,
366                        location,
367                        id,
368                        was_leaked: AtomicBool::new(false),
369                    });
370                    let weak = Arc::downgrade(&data);
371                    trace!(
372                        target: "maitake::alloc",
373                        id,
374                        "type" = %type_name,
375                        %location,
376                        "started tracking allocation",
377                    );
378                    inner.tracks.push(weak);
379                    data
380                }
381
382                pub(in crate::loom) fn check(&self) {
383                    let leaked = self
384                        .0
385                        .lock()
386                        .unwrap()
387                        .tracks
388                        .iter()
389                        .filter_map(|weak| {
390                            let data = weak.upgrade()?;
391                            data.was_leaked.store(true, Ordering::SeqCst);
392                            Some(format!(
393                                " - id {}, {} allocated at {}",
394                                data.id, data.type_name, data.location
395                            ))
396                        })
397                        .collect::<Vec<_>>();
398                    if !leaked.is_empty() {
399                        let leaked = leaked.join("\n  ");
400                        panic!("the following allocations were leaked:\n  {leaked}");
401                    }
402                }
403            }
404
405            impl Drop for TrackData {
406                fn drop(&mut self) {
407                    if !self.was_leaked.load(Ordering::SeqCst) {
408                        trace!(
409                            target: "maitake::alloc",
410                            id = self.id,
411                            "type" = %self.type_name,
412                            location = %self.location,
413                            "dropped all references to a tracked allocation",
414                        );
415                    }
416                }
417            }
418        }
419
420        #[cfg(test)]
421        #[derive(Debug)]
422        #[pin_project::pin_project]
423        pub(crate) struct TrackFuture<F> {
424            #[pin]
425            inner: F,
426            track: Option<Arc<track::TrackData>>,
427        }
428
429        #[cfg(test)]
430        impl<F: Future> Future for TrackFuture<F> {
431            type Output = TrackFuture<F::Output>;
432            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
433                let this = self.project();
434                this.inner.poll(cx).map(|inner| TrackFuture {
435                    inner,
436                    track: this.track.clone(),
437                })
438            }
439        }
440
441        #[cfg(test)]
442        impl<F> TrackFuture<F> {
443            /// Wrap a `Future` in a `TrackFuture` that participates in Loom's
444            /// leak checking.
445            #[track_caller]
446            pub(crate) fn new(inner: F) -> Self {
447                let track = track::Registry::start_tracking::<F>();
448                Self { inner, track }
449            }
450
451            /// Stop tracking this future, and return the inner value.
452            pub(crate) fn into_inner(self) -> F {
453                self.inner
454            }
455        }
456
457        #[cfg(test)]
458        #[track_caller]
459        pub(crate) fn track_future<F: Future>(inner: F) -> TrackFuture<F> {
460            TrackFuture::new(inner)
461        }
462
463        // PartialEq impl so that `assert_eq!(..., Ok(...))` works
464        #[cfg(test)]
465        impl<F: PartialEq> PartialEq for TrackFuture<F> {
466            fn eq(&self, other: &Self) -> bool {
467                self.inner == other.inner
468            }
469        }
470
471        /// Track allocations, detecting leaks
472        #[derive(Debug, Default)]
473        pub struct Track<T> {
474            value: T,
475
476            #[cfg(test)]
477            track: Option<Arc<track::TrackData>>,
478        }
479
480        impl<T> Track<T> {
481            /// Track a value for leaks
482            #[inline(always)]
483            #[track_caller]
484            pub fn new(value: T) -> Track<T> {
485                Track {
486                    value,
487
488                    #[cfg(test)]
489                    track: track::Registry::start_tracking::<T>(),
490                }
491            }
492
493            /// Get a reference to the value
494            #[inline(always)]
495            pub fn get_ref(&self) -> &T {
496                &self.value
497            }
498
499            /// Get a mutable reference to the value
500            #[inline(always)]
501            pub fn get_mut(&mut self) -> &mut T {
502                &mut self.value
503            }
504
505            /// Stop tracking the value for leaks
506            #[inline(always)]
507            pub fn into_inner(self) -> T {
508                self.value
509            }
510        }
511    }
512
513    #[cfg(test)]
514    pub(crate) mod future {
515        pub(crate) use tokio_test::block_on;
516    }
517
518    #[cfg(test)]
519    pub(crate) fn traceln(args: std::fmt::Arguments) {
520        eprintln!("{args}");
521    }
522
523    #[cfg(not(test))]
524    pub(crate) fn traceln(_: core::fmt::Arguments) {}
525}
maitake/loom.rs

maitake/
loom.rs
