1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
use futures::task::AtomicWaker;
use futures::Future;
use parking_lot::Mutex;
use std::collections::LinkedList;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
#[derive(Debug, Default)]
struct TaskQueueTaskWaker {
is_ready: AtomicBool,
waker: AtomicWaker,
}
#[derive(Debug, Default)]
struct TaskQueueTasks {
is_running: bool,
wakers: LinkedList<Arc<TaskQueueTaskWaker>>,
}
#[derive(Debug, Default)]
pub struct TaskQueue {
tasks: Mutex<TaskQueueTasks>,
}
impl TaskQueue {
pub async fn acquire(&self) -> TaskQueuePermit {
let acquire = TaskQueuePermitAcquire::new(self);
acquire.await;
TaskQueuePermit(self)
}
pub async fn queue<R>(&self, future: impl Future<Output = R>) -> R {
let _permit = self.acquire().await;
future.await
}
}
pub struct TaskQueuePermit<'a>(&'a TaskQueue);
impl<'a> Drop for TaskQueuePermit<'a> {
fn drop(&mut self) {
let next_item = {
let mut tasks = self.0.tasks.lock();
let next_item = tasks.wakers.pop_front();
tasks.is_running = next_item.is_some();
next_item
};
if let Some(next_item) = next_item {
next_item.is_ready.store(true, Ordering::SeqCst);
next_item.waker.wake();
}
}
}
struct TaskQueuePermitAcquire<'a> {
task_queue: &'a TaskQueue,
initialized: AtomicBool,
waker: Arc<TaskQueueTaskWaker>,
}
impl<'a> TaskQueuePermitAcquire<'a> {
pub fn new(task_queue: &'a TaskQueue) -> Self {
Self {
task_queue,
initialized: Default::default(),
waker: Default::default(),
}
}
}
impl<'a> Future for TaskQueuePermitAcquire<'a> {
type Output = ();
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
self.waker.waker.register(cx.waker());
if !self.initialized.swap(true, Ordering::SeqCst) {
let mut tasks = self.task_queue.tasks.lock();
if !tasks.is_running {
tasks.is_running = true;
return std::task::Poll::Ready(());
}
tasks.wakers.push_back(self.waker.clone());
return std::task::Poll::Pending;
}
if self.waker.is_ready.load(Ordering::SeqCst) {
std::task::Poll::Ready(())
} else {
std::task::Poll::Pending
}
}
}
#[cfg(test)]
mod tests {
use parking_lot::Mutex;
use std::sync::Arc;
use super::TaskQueue;
#[tokio::test]
async fn task_queue_runs_one_after_other() {
let task_queue = TaskQueue::default();
let mut tasks = Vec::new();
let data = Arc::new(Mutex::new(0));
for i in 0..100 {
let data = data.clone();
tasks.push(task_queue.queue(async move {
tokio::task::spawn_blocking(move || {
let mut data = data.lock();
if *data != i {
panic!("Value was not equal.");
}
*data = i + 1;
})
.await
.unwrap();
}));
}
futures::future::join_all(tasks).await;
}
}