cuckoo_runtime/
lib.rs

1use std::{
2    sync::Arc,
3    thread::{self, ThreadId},
4};
5
6use async_task::{Runnable, Task};
7use crossbeam::deque::{Injector, Stealer, Worker};
8use dashmap::DashMap;
9
10use pin_project_lite::pin_project;
11
12thread_local! {
13    static WORK_QUEUE: Worker<Runnable> = Worker::new_fifo();
14}
15
16pub struct Runtime {
17    injector: Arc<Injector<Runnable>>,
18    stealers: Arc<DashMap<ThreadId, Stealer<Runnable>>>,
19}
20
21#[derive(Clone)]
22pub struct Handle {
23    injector: Arc<Injector<Runnable>>,
24    stealers: Arc<DashMap<ThreadId, Stealer<Runnable>>>,
25}
26
27pin_project! {
28    pub struct JoinHandle<T> {
29        #[pin]
30        task: Task<T>,
31    }
32}
33
34impl<T> Future for JoinHandle<T> {
35    type Output = T;
36
37    fn poll(
38        self: std::pin::Pin<&mut Self>,
39        cx: &mut std::task::Context<'_>,
40    ) -> std::task::Poll<Self::Output> {
41        self.project().task.poll(cx)
42    }
43}
44
45impl Runtime {
46    /// Creates a new runtime with a fixed number of background threads.
47    ///
48    /// Background threads execute tasks that are spawned in the background, allowing foreground threads to keep doing other work.
49    pub fn new(background_threads: usize) -> Self {
50        let injector: Arc<Injector<Runnable>> = Arc::new(Injector::new());
51        for _ in 0..background_threads {
52            std::thread::spawn({
53                let injector = injector.clone();
54                move || {
55                    loop {
56                        if let Some(task) = injector.steal().success() {
57                            task.run();
58                        }
59                    }
60                }
61            });
62        }
63        Self {
64            injector,
65            stealers: Arc::new(DashMap::new()),
66        }
67    }
68
69    pub fn handle(&self) -> Handle {
70        Handle {
71            injector: Arc::clone(&self.injector),
72            stealers: Arc::clone(&self.stealers),
73        }
74    }
75
76    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
77    where
78        F: Future + Send + 'static,
79        F::Output: Send + Sync,
80    {
81        let injector = self.injector.clone();
82        let schedule = move |runnable| injector.push(runnable);
83        let (runnable, task) = async_task::spawn(future, schedule);
84        runnable.schedule();
85        JoinHandle { task }
86    }
87
88    pub fn block_on<F>(&self, future: F) -> F::Output
89    where
90        F: Future + Send + 'static,
91        F::Output: Send + Sync,
92    {
93        self.handle().block_on(future)
94    }
95}
96
97impl Handle {
98    pub fn block_on<F>(&self, future: F) -> F::Output
99    where
100        F: Future + Send + 'static,
101        F::Output: Send + Sync,
102    {
103        let schedule = {
104            let stealers = self.stealers.clone();
105            move |runnable| {
106                WORK_QUEUE.with(|q| {
107                    // Make sure this thread's stealer is populated
108                    stealers
109                        .entry(thread::current().id())
110                        .or_insert(q.stealer());
111                    q.push(runnable);
112                })
113            }
114        };
115        let (runnable, task) = async_task::spawn(future, schedule);
116        runnable.schedule();
117
118        while !task.is_finished() {
119            if let Some(stolen_task) = self.find_task() {
120                stolen_task.run();
121            }
122        }
123
124        futures::executor::block_on(task)
125    }
126
127    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
128    where
129        F: Future + Send + 'static,
130        F::Output: Send + Sync,
131    {
132        let injector = self.injector.clone();
133        let schedule = move |runnable| injector.push(runnable);
134        let (runnable, task) = async_task::spawn(future, schedule);
135        runnable.schedule();
136        JoinHandle { task }
137    }
138
139    fn find_task(&self) -> Option<Runnable> {
140        WORK_QUEUE.with(|local| {
141            // Pop a task from the local queue, if not empty.
142            local.pop().or_else(|| {
143                std::iter::repeat_with(|| {
144                    self.injector
145                        .steal_batch_with_limit_and_pop(local, 1)
146                        .or_else(|| self.stealers.iter().map(|s| s.steal()).collect())
147                })
148                .find(|s| !s.is_retry())
149                .and_then(|s| s.success())
150            })
151        })
152    }
153}
154
155#[cfg(test)]
156mod tests {
157
158    use rstest::rstest;
159
160    use super::*;
161
162    #[rstest]
163    #[case(0)]
164    #[case(1)]
165    fn test_basic(#[case] thread_count: usize) {
166        let rt = Runtime::new(thread_count);
167        let jh: JoinHandle<()> = rt.spawn(async move { println!("spawned task!") });
168        let handle = rt.handle();
169        let h2 = handle.clone();
170
171        handle.block_on(async move {
172            _ = h2.spawn(async move {
173                if thread_count == 0 {
174                    panic!("This should never run");
175                }
176            });
177            println!("blocking here!");
178            println!("Awaited the second future");
179        });
180        println!("Finished block_on block");
181        rt.block_on(jh);
182    }
183}
cuckoo_runtime/lib.rs

cuckoo_runtime/
lib.rs
