summaryrefslogtreecommitdiff
path: root/cli/util/sync/task_queue.rs
diff options
context:
space:
mode:
Diffstat (limited to 'cli/util/sync/task_queue.rs')
-rw-r--r--cli/util/sync/task_queue.rs266
1 files changed, 266 insertions, 0 deletions
diff --git a/cli/util/sync/task_queue.rs b/cli/util/sync/task_queue.rs
new file mode 100644
index 000000000..6ef747e1a
--- /dev/null
+++ b/cli/util/sync/task_queue.rs
@@ -0,0 +1,266 @@
+// Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
+
+use std::collections::LinkedList;
+use std::sync::Arc;
+
+use deno_core::futures::task::AtomicWaker;
+use deno_core::futures::Future;
+use deno_core::parking_lot::Mutex;
+
+use super::AtomicFlag;
+
+#[derive(Debug, Default)]
+struct TaskQueueTaskItem {
+ is_ready: AtomicFlag,
+ is_future_dropped: AtomicFlag,
+ waker: AtomicWaker,
+}
+
+#[derive(Debug, Default)]
+struct TaskQueueTasks {
+ is_running: bool,
+ items: LinkedList<Arc<TaskQueueTaskItem>>,
+}
+
+/// A queue that executes tasks sequentially one after the other
+/// ensuring order and that no task runs at the same time as another.
+///
+/// Note that this differs from tokio's semaphore in that the order
+/// is acquired synchronously.
+#[derive(Debug, Default)]
+pub struct TaskQueue {
+ tasks: Mutex<TaskQueueTasks>,
+}
+
+impl TaskQueue {
+ /// Acquires a permit where the tasks are executed one at a time
+ /// and in the order that they were acquired.
+ pub fn acquire(&self) -> TaskQueuePermitAcquireFuture {
+ TaskQueuePermitAcquireFuture::new(self)
+ }
+
+ /// Alternate API that acquires a permit internally
+ /// for the duration of the future.
+ #[allow(unused)]
+ pub fn run<'a, R>(
+ &'a self,
+ future: impl Future<Output = R> + 'a,
+ ) -> impl Future<Output = R> + 'a {
+ let acquire_future = self.acquire();
+ async move {
+ let permit = acquire_future.await;
+ let result = future.await;
+ drop(permit); // explicit for clarity
+ result
+ }
+ }
+
+ fn raise_next(&self) {
+ let front_item = {
+ let mut tasks = self.tasks.lock();
+
+ // clear out any wakers for futures that were dropped
+ while let Some(front_waker) = tasks.items.front() {
+ if front_waker.is_future_dropped.is_raised() {
+ tasks.items.pop_front();
+ } else {
+ break;
+ }
+ }
+ let front_item = tasks.items.pop_front();
+ tasks.is_running = front_item.is_some();
+ front_item
+ };
+
+ // wake up the next waker
+ if let Some(front_item) = front_item {
+ front_item.is_ready.raise();
+ front_item.waker.wake();
+ }
+ }
+}
+
+/// A permit that when dropped will allow another task to proceed.
+pub struct TaskQueuePermit<'a>(&'a TaskQueue);
+
+impl<'a> Drop for TaskQueuePermit<'a> {
+ fn drop(&mut self) {
+ self.0.raise_next();
+ }
+}
+
+pub struct TaskQueuePermitAcquireFuture<'a> {
+ task_queue: Option<&'a TaskQueue>,
+ item: Arc<TaskQueueTaskItem>,
+}
+
+impl<'a> TaskQueuePermitAcquireFuture<'a> {
+ pub fn new(task_queue: &'a TaskQueue) -> Self {
+ // acquire the waker position synchronously
+ let mut tasks = task_queue.tasks.lock();
+ let item = if !tasks.is_running {
+ tasks.is_running = true;
+ let item = Arc::new(TaskQueueTaskItem::default());
+ item.is_ready.raise();
+ item
+ } else {
+ let item = Arc::new(TaskQueueTaskItem::default());
+ tasks.items.push_back(item.clone());
+ item
+ };
+ drop(tasks);
+ Self {
+ task_queue: Some(task_queue),
+ item,
+ }
+ }
+}
+
+impl<'a> Drop for TaskQueuePermitAcquireFuture<'a> {
+ fn drop(&mut self) {
+ if let Some(task_queue) = self.task_queue.take() {
+ if self.item.is_ready.is_raised() {
+ task_queue.raise_next();
+ } else {
+ self.item.is_future_dropped.raise();
+ }
+ }
+ }
+}
+
+impl<'a> Future for TaskQueuePermitAcquireFuture<'a> {
+ type Output = TaskQueuePermit<'a>;
+
+ fn poll(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> std::task::Poll<Self::Output> {
+ if self.item.is_ready.is_raised() {
+ std::task::Poll::Ready(TaskQueuePermit(self.task_queue.take().unwrap()))
+ } else {
+ self.item.waker.register(cx.waker());
+ std::task::Poll::Pending
+ }
+ }
+}
+
+#[cfg(test)]
+mod test {
+ use deno_core::futures;
+ use deno_core::parking_lot::Mutex;
+ use std::sync::Arc;
+
+ use super::*;
+
+ #[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.run(async move {
+ deno_core::unsync::spawn_blocking(move || {
+ let mut data = data.lock();
+ assert_eq!(*data, i);
+ *data = i + 1;
+ })
+ .await
+ .unwrap();
+ }));
+ }
+ futures::future::join_all(tasks).await;
+ }
+
+ #[tokio::test]
+ async fn task_queue_run_in_sequence() {
+ let task_queue = TaskQueue::default();
+ let data = Arc::new(Mutex::new(0));
+
+ let first = task_queue.run(async {
+ *data.lock() = 1;
+ });
+ let second = task_queue.run(async {
+ assert_eq!(*data.lock(), 1);
+ *data.lock() = 2;
+ });
+ let _ = tokio::join!(first, second);
+
+ assert_eq!(*data.lock(), 2);
+ }
+
+ #[tokio::test]
+ async fn task_queue_future_dropped_before_poll() {
+ let task_queue = Arc::new(TaskQueue::default());
+
+ // acquire a future, but do not await it
+ let future = task_queue.acquire();
+
+ // this task tries to acquire another permit, but will be blocked by the first permit.
+ let enter_flag = Arc::new(AtomicFlag::default());
+ let delayed_task = deno_core::unsync::spawn({
+ let enter_flag = enter_flag.clone();
+ let task_queue = task_queue.clone();
+ async move {
+ enter_flag.raise();
+ task_queue.acquire().await;
+ true
+ }
+ });
+
+ // ensure the task gets a chance to be scheduled and blocked
+ tokio::task::yield_now().await;
+ assert!(enter_flag.is_raised());
+
+ // now, drop the first future
+ drop(future);
+
+ assert!(delayed_task.await.unwrap());
+ }
+
+ #[tokio::test]
+ async fn task_queue_many_future_dropped_before_poll() {
+ let task_queue = Arc::new(TaskQueue::default());
+
+ // acquire a future, but do not await it
+ let mut futures = Vec::new();
+ for _ in 0..=10_000 {
+ futures.push(task_queue.acquire());
+ }
+
+ // this task tries to acquire another permit, but will be blocked by the first permit.
+ let enter_flag = Arc::new(AtomicFlag::default());
+ let delayed_task = deno_core::unsync::spawn({
+ let task_queue = task_queue.clone();
+ let enter_flag = enter_flag.clone();
+ async move {
+ enter_flag.raise();
+ task_queue.acquire().await;
+ true
+ }
+ });
+
+ // ensure the task gets a chance to be scheduled and blocked
+ tokio::task::yield_now().await;
+ assert!(enter_flag.is_raised());
+
+ // now, drop the futures
+ drop(futures);
+
+ assert!(delayed_task.await.unwrap());
+ }
+
+ #[tokio::test]
+ async fn task_queue_middle_future_dropped_while_permit_acquired() {
+ let task_queue = TaskQueue::default();
+
+ let fut1 = task_queue.acquire();
+ let fut2 = task_queue.acquire();
+ let fut3 = task_queue.acquire();
+
+ // should not hang
+ drop(fut2);
+ drop(fut1.await);
+ drop(fut3.await);
+ }
+}