jonaskruckenberg.de
Next Generation WASM Microkernel Operating System
1// Copyright 2025. Jonas Kruckenberg
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8use crate::loom::sync::Arc;
9use crate::loom::sync::{Condvar, Mutex as StdMutex};
10use core::mem::ManuallyDrop;
11use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
12use futures::pin_mut;
13use futures::task::WakerRef;
14use util::loom_const_fn;
15
16#[derive(Debug)]
17pub struct ThreadNotify {
18 mutex: StdMutex<bool>,
19 condvar: Condvar,
20}
21
22impl ThreadNotify {
23 loom_const_fn! {
24 const fn new() -> Self {
25 Self {
26 mutex: StdMutex::new(false),
27 condvar: Condvar::new()
28 }
29 }
30 }
31
32 #[tracing::instrument(level = "trace", skip(self))]
33 fn wait(&self) {
34 let mut notified = self.mutex.lock().unwrap();
35 while !*notified {
36 notified = self.condvar.wait(notified).unwrap();
37 }
38 *notified = false;
39 }
40
41 fn notify(&self) {
42 let mut signaled = self.mutex.lock().unwrap();
43 *signaled = true;
44 self.condvar.notify_one();
45 }
46}
47
48fn waker_ref(wake: &Arc<ThreadNotify>) -> WakerRef<'_> {
49 static WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
50 clone_arc_raw,
51 wake_arc_raw,
52 wake_by_ref_arc_raw,
53 drop_arc_raw,
54 );
55
56 unsafe fn clone_arc_raw(data: *const ()) -> RawWaker {
57 unsafe { Arc::increment_strong_count(data.cast::<ThreadNotify>()) }
58 RawWaker::new(data, &WAKER_VTABLE)
59 }
60
61 unsafe fn wake_arc_raw(data: *const ()) {
62 let arc = unsafe { Arc::from_raw(data.cast::<ThreadNotify>()) };
63 ThreadNotify::notify(&arc);
64 }
65
66 // used by `waker_ref`
67 unsafe fn wake_by_ref_arc_raw(data: *const ()) {
68 // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
69 let arc = ManuallyDrop::new(unsafe { Arc::from_raw(data.cast::<ThreadNotify>()) });
70 ThreadNotify::notify(&arc);
71 }
72
73 unsafe fn drop_arc_raw(data: *const ()) {
74 drop(unsafe { Arc::from_raw(data.cast::<ThreadNotify>()) })
75 }
76
77 // simply copy the pointer instead of using Arc::into_raw,
78 // as we don't actually keep a refcount by using ManuallyDrop.<
79 let ptr = Arc::as_ptr(wake).cast::<()>();
80
81 let waker = ManuallyDrop::new(unsafe { Waker::from_raw(RawWaker::new(ptr, &WAKER_VTABLE)) });
82 WakerRef::new_unowned(waker)
83}
84
85pub fn block_on<F: Future>(f: F) -> F::Output {
86 pin_mut!(f);
87
88 crate::loom::thread_local! {
89 static CURRENT_THREAD_NOTIFY: Arc<ThreadNotify> = Arc::new(ThreadNotify::new());
90 }
91
92 CURRENT_THREAD_NOTIFY.with(|thread_notify| {
93 let waker = waker_ref(&thread_notify);
94 let mut cx = Context::from_waker(&waker);
95 loop {
96 if let Poll::Ready(t) = f.as_mut().poll(&mut cx) {
97 return t;
98 }
99
100 thread_notify.wait();
101 }
102 })
103}