Repository: vercel/next.js
Pull request: Turbopack: Scope with work queue

### What?

* Use a work queue and worker tasks in the scope implementation to reduce the number of spawned tokio tasks. Only spawn a tokio task per CPU - 1.
* Run work in the current tokio task too to avoid blocking time and use one less task.

Before we only did a blocking wait on the current tokio task, which did require calling block_in_place pretty often.
Now we also do work on the current tokio task too, which requires calling block_in_place less often. (We only call block_in_place when we need to block longer than 1ms)

Closes PACK-5577

Changed files:
filename: turbopack/crates/turbo-tasks/src/parallel.rs
patch: @@ -25,7 +25,7 @@ where
     let f = &f;
     let _results = scope_and_block(len.div_ceil(chunk_size), |scope| {
         for chunk in items.chunks(chunk_size) {
-            scope.spawn(async move {
+            scope.spawn(move || {
                 for item in chunk {
                     f(item);
                 }
@@ -49,7 +49,7 @@ where
     let f = &f;
     let _results = scope_and_block(len.div_ceil(chunk_size), |scope| {
         for chunk in into_chunks(items, chunk_size) {
-            scope.spawn(async move {
+            scope.spawn(move || {
                 // SAFETY: Even when f() panics we drop all items in the chunk.
                 for item in chunk {
                     f(item);
@@ -78,7 +78,7 @@ where
     let f = &f;
     scope_and_block(len.div_ceil(chunk_size), |scope| {
         for chunk in items.chunks(chunk_size) {
-            scope.spawn(async move {
+            scope.spawn(move || {
                 for item in chunk {
                     f(item)?;
                 }
@@ -108,7 +108,7 @@ where
     let f = &f;
     scope_and_block(len.div_ceil(chunk_size), |scope| {
         for chunk in items.chunks_mut(chunk_size) {
-            scope.spawn(async move {
+            scope.spawn(move || {
                 for item in chunk {
                     f(item)?;
                 }
@@ -138,7 +138,7 @@ where
     let f = &f;
     scope_and_block(len.div_ceil(chunk_size), |scope| {
         for chunk in into_chunks(items, chunk_size) {
-            scope.spawn(async move {
+            scope.spawn(move || {
                 for item in chunk {
                     f(item)?;
                 }
@@ -167,7 +167,7 @@ where
     let f = &f;
     scope_and_block(len.div_ceil(chunk_size), |scope| {
         for chunk in items.chunks(chunk_size) {
-            scope.spawn(async move { chunk.iter().map(f).collect::<Vec<_>>() })
+            scope.spawn(move || chunk.iter().map(f).collect::<Vec<_>>())
         }
     })
     .flatten()
@@ -192,7 +192,7 @@ where
     let f = &f;
     scope_and_block(len.div_ceil(chunk_size), |scope| {
         for chunk in into_chunks(items, chunk_size) {
-            scope.spawn(async move { chunk.map(f).collect::<Vec<_>>() })
+            scope.spawn(move || chunk.map(f).collect::<Vec<_>>())
         }
     })
     .flatten()

filename: turbopack/crates/turbo-tasks/src/scope.rs
patch: @@ -2,33 +2,43 @@
 
 use std::{
     any::Any,
+    collections::VecDeque,
     marker::PhantomData,
     panic::{self, AssertUnwindSafe, catch_unwind},
-    pin::Pin,
     sync::{
         Arc,
         atomic::{AtomicUsize, Ordering},
     },
-    thread::{self, Thread},
+    thread::{self, Thread, available_parallelism},
     time::{Duration, Instant},
 };
 
-use futures::FutureExt;
-use parking_lot::Mutex;
+use once_cell::sync::Lazy;
+use parking_lot::{Condvar, Mutex};
 use tokio::{runtime::Handle, task::block_in_place};
-use tracing::{Instrument, Span, info_span};
+use tracing::{Span, info_span};
 
-use crate::{
-    TurboTasksApi,
-    manager::{try_turbo_tasks, turbo_tasks_future_scope},
-};
+use crate::{TurboTasksApi, manager::try_turbo_tasks, turbo_tasks_scope};
+
+/// Number of worker tasks to spawn that process jobs. It's 1 less than the number of cpus as we
+/// also use the current task as worker.
+static WORKER_TASKS: Lazy<usize> = Lazy::new(|| available_parallelism().map_or(0, |n| n.get() - 1));
+
+enum WorkQueueJob {
+    Job(usize, Box<dyn FnOnce() + Send + 'static>),
+    End,
+}
 
 struct ScopeInner {
     main_thread: Thread,
     remaining_tasks: AtomicUsize,
     /// The first panic that occurred in the tasks, by task index.
     /// The usize value is the index of the task.
     panic: Mutex<Option<(Box<dyn Any + Send + 'static>, usize)>>,
+    /// The work queue for spawned jobs that have not yet been picked up by a worker task.
+    work_queue: Mutex<VecDeque<WorkQueueJob>>,
+    /// A condition variable to notify worker tasks of new work or end of work.
+    work_queue_condition_var: Condvar,
 }
 
 impl ScopeInner {
@@ -82,6 +92,57 @@ impl ScopeInner {
             panic::resume_unwind(err);
         }
     }
+
+    fn worker(&self, first_job_index: usize, first_job: Box<dyn FnOnce() + Send + 'static>) {
+        let mut current_job_index = first_job_index;
+        let mut current_job = first_job;
+        loop {
+            let result = catch_unwind(AssertUnwindSafe(current_job));
+            let panic = result.err().map(|e| (e, current_job_index));
+            self.on_task_finished(panic);
+            let Some((index, job)) = self.pick_job_from_work_queue() else {
+                return;
+            };
+            current_job_index = index;
+            current_job = job;
+        }
+    }
+
+    fn pick_job_from_work_queue(&self) -> Option<(usize, Box<dyn FnOnce() + Send + 'static>)> {
+        let mut work_queue = self.work_queue.lock();
+        let job = loop {
+            if let Some(job) = work_queue.pop_front() {
+                break job;
+            } else {
+                self.work_queue_condition_var.wait(&mut work_queue);
+            };
+        };
+        match job {
+            WorkQueueJob::Job(index, job) => {
+                drop(work_queue);
+                Some((index, job))
+            }
+            WorkQueueJob::End => {
+                work_queue.push_front(WorkQueueJob::End);
+                drop(work_queue);
+                self.work_queue_condition_var.notify_all();
+                None
+            }
+        }
+    }
+
+    fn end_and_help_complete(&self) {
+        let job;
+        {
+            let mut work_queue = self.work_queue.lock();
+            job = work_queue.pop_front();
+            work_queue.push_back(WorkQueueJob::End);
+        }
+        self.work_queue_condition_var.notify_all();
+        if let Some(WorkQueueJob::Job(index, job)) = job {
+            self.worker(index, job);
+        }
+    }
 }
 
 /// Scope to allow spawning tasks with a limited lifetime.
@@ -115,6 +176,8 @@ impl<'scope, 'env: 'scope, R: Send + 'env> Scope<'scope, 'env, R> {
                 main_thread: thread::current(),
                 remaining_tasks: AtomicUsize::new(0),
                 panic: Mutex::new(None),
+                work_queue: Mutex::new(VecDeque::new()),
+                work_queue_condition_var: Condvar::new(),
             }),
             handle: Handle::current(),
             turbo_tasks: try_turbo_tasks(),
@@ -126,57 +189,64 @@ impl<'scope, 'env: 'scope, R: Send + 'env> Scope<'scope, 'env, R> {
     /// Spawns a new task in the scope.
     pub fn spawn<F>(&self, f: F)
     where
-        F: Future<Output = R> + Send + 'env,
+        F: FnOnce() -> R + Send + 'env,
     {
         let index = self.index.fetch_add(1, Ordering::Relaxed);
         assert!(index < self.results.len(), "Too many tasks spawned");
         let result_cell: &Mutex<Option<R>> = &self.results[index];
 
-        let f: Box<dyn Future<Output = ()> + Send + 'scope> = Box::new(async move {
-            let result = f.await;
+        let f: Box<dyn FnOnce() + Send + 'scope> = Box::new(|| {
+            let result = f();
             *result_cell.lock() = Some(result);
         });
-        let f: *mut (dyn Future<Output = ()> + Send + 'scope) = Box::into_raw(f);
+        let f: *mut (dyn FnOnce() + Send + 'scope) = Box::into_raw(f);
         // SAFETY: Scope ensures (e. g. in Drop) that spawned tasks is awaited before the
         // lifetime `'env` ends.
         #[allow(
             clippy::unnecessary_cast,
             reason = "Clippy thinks this is unnecessary, but it actually changes the lifetime"
         )]
-        let f = f as *mut (dyn Future<Output = ()> + Send + 'static);
+        let f = f as *mut (dyn FnOnce() + Send + 'static);
         // SAFETY: We just called `Box::into_raw`.
         let f = unsafe { Box::from_raw(f) };
-        // We pin the future in the box in memory to be able to await it.
-        let f = Pin::from(f);
 
         let turbo_tasks = self.turbo_tasks.clone();
         let span = self.span.clone();
 
-        let inner = self.inner.clone();
-        inner.remaining_tasks.fetch_add(1, Ordering::Relaxed);
-        self.handle.spawn(async move {
-            let result = AssertUnwindSafe(
-                async move {
-                    if let Some(turbo_tasks) = turbo_tasks {
-                        // Ensure that the turbo tasks context is maintained across the task.
-                        turbo_tasks_future_scope(turbo_tasks, f).await;
-                    } else {
-                        // If no turbo tasks context is available, just run the future.
-                        f.await;
-                    }
+        self.inner.remaining_tasks.fetch_add(1, Ordering::Relaxed);
+
+        // The first job always goes to the work_queue to be worked on by the main thread.
+        // After that we spawn a new worker for every job until we reach WORKER_TASKS.
+        // After that we queue up jobs in the work_queue again.
+        if (1..=*WORKER_TASKS).contains(&index) {
+            let inner = self.inner.clone();
+            // Spawn a worker task that will process that tasks and potentially more.
+            self.handle.spawn(async move {
+                let _span = span.entered();
+                if let Some(turbo_tasks) = turbo_tasks {
+                    // Ensure that the turbo tasks context is maintained across the worker.
+                    turbo_tasks_scope(turbo_tasks, || {
+                        inner.worker(index, f);
+                    });
+                } else {
+                    // If no turbo tasks context is available, just run the worker.
+                    inner.worker(index, f);
                 }
-                .instrument(span),
-            )
-            .catch_unwind()
-            .await;
-            let panic = result.err().map(|e| (e, index));
-            inner.on_task_finished(panic);
-        });
+            });
+        } else {
+            // Queue the task to be processed by a worker task.
+            self.inner
+                .work_queue
+                .lock()
+                .push_back(WorkQueueJob::Job(index, f));
+            self.inner.work_queue_condition_var.notify_one();
+        }
     }
 }
 
 impl<'scope, 'env: 'scope, R: Send + 'env> Drop for Scope<'scope, 'env, R> {
     fn drop(&mut self) {
+        self.inner.end_and_help_complete();
         self.inner.wait_and_rethrow_panic();
     }
 }
@@ -223,7 +293,7 @@ mod tests {
     async fn test_scope() {
         let results = scope_and_block(1000, |scope| {
             for i in 0..1000 {
-                scope.spawn(async move { i });
+                scope.spawn(move || i);
             }
         });
         results.enumerate().for_each(|(i, result)| {
@@ -235,7 +305,7 @@ mod tests {
     async fn test_empty_scope() {
         let results = scope_and_block(0, |scope| {
             if false {
-                scope.spawn(async move { 42 });
+                scope.spawn(|| 42);
             }
         });
         assert_eq!(results.count(), 0);
@@ -244,7 +314,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
     async fn test_single_task() {
         let results = scope_and_block(1, |scope| {
-            scope.spawn(async move { 42 });
+            scope.spawn(|| 42);
         })
         .collect::<Vec<_>>();
         assert_eq!(results, vec![42]);
@@ -253,7 +323,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
     async fn test_task_finish_before_scope() {
         let results = scope_and_block(1, |scope| {
-            scope.spawn(async move { 42 });
+            scope.spawn(|| 42);
             thread::sleep(std::time::Duration::from_millis(100));
         })
         .collect::<Vec<_>>();
@@ -263,7 +333,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
     async fn test_task_finish_after_scope() {
         let results = scope_and_block(1, |scope| {
-            scope.spawn(async move {
+            scope.spawn(|| {
                 thread::sleep(std::time::Duration::from_millis(100));
                 42
             });
@@ -277,7 +347,7 @@ mod tests {
         let result = catch_unwind(AssertUnwindSafe(|| {
             let _results = scope_and_block(1000, |scope| {
                 for i in 0..500 {
-                    scope.spawn(async move { i });
+                    scope.spawn(move || i);
                 }
                 panic!("Intentional panic");
             });
@@ -295,7 +365,7 @@ mod tests {
         let result = catch_unwind(AssertUnwindSafe(|| {
             let _results = scope_and_block(1000, |scope| {
                 for i in 0..1000 {
-                    scope.spawn(async move {
+                    scope.spawn(move || {
                         if i == 500 {
                             panic!("Intentional panic");
                         } else if i == 501 {