Turbopack: add turbo_tasks::spawn

Author: sokra

turbopack/crates/turbo-tasks-backend/src/database/turbo/mod.rs

@@ -1,15 +1,12 @@
 use std::{cmp::max, path::PathBuf, sync::Arc, thread::available_parallelism};
 
-use anyhow::Result;
+use anyhow::{Ok, Result};
 use parking_lot::Mutex;
-use tokio::{
-    runtime::Handle,
-    spawn,
-    task::{JoinHandle, block_in_place},
-};
+use tokio::{runtime::Handle, task::block_in_place};
 use turbo_persistence::{
     ArcSlice, CompactConfig, KeyBase, StoreKey, TurboPersistence, ValueBuffer,
 };
+use turbo_tasks::{JoinHandle, message_queue::TimingEvent, spawn, turbo_tasks};
 
 use crate::database::{
     key_value_database::{KeySpace, KeyValueDatabase},
@@ -237,5 +234,5 @@ impl<'l> From<WriteBuffer<'l>> for ValueBuffer<'l> {
 }
 
 fn join(handle: JoinHandle<Result<()>>) -> Result<()> {
-    block_in_place(|| Handle::current().block_on(handle))?
+    block_in_place(|| Handle::current().block_on(handle))
 }

turbopack/crates/turbo-tasks-malloc/src/lib.rs

@@ -3,6 +3,7 @@ mod counter;
 use std::{
     alloc::{GlobalAlloc, Layout},
     marker::PhantomData,
+    ops::{Add, AddAssign},
 };
 
 use self::counter::{add, flush, get, remove, update};
@@ -16,12 +17,45 @@ pub struct AllocationInfo {
 }
 
 impl AllocationInfo {
+    pub const ZERO: Self = Self {
+        allocations: 0,
+        deallocations: 0,
+        allocation_count: 0,
+        deallocation_count: 0,
+    };
+
     pub fn is_empty(&self) -> bool {
         self.allocations == 0
             && self.deallocations == 0
             && self.allocation_count == 0
             && self.deallocation_count == 0
     }
+
+    pub fn memory_usage(&self) -> usize {
+        self.allocations.saturating_sub(self.deallocations)
+    }
+}
+
+impl Add<Self> for AllocationInfo {
+    type Output = Self;
+
+    fn add(self, other: Self) -> Self {
+        Self {
+            allocations: self.allocations + other.allocations,
+            deallocations: self.deallocations + other.deallocations,
+            allocation_count: self.allocation_count + other.allocation_count,
+            deallocation_count: self.deallocation_count + other.deallocation_count,
+        }
+    }
+}
+
+impl AddAssign<Self> for AllocationInfo {
+    fn add_assign(&mut self, other: Self) {
+        self.allocations += other.allocations;
+        self.deallocations += other.deallocations;
+        self.allocation_count += other.allocation_count;
+        self.deallocation_count += other.deallocation_count;
+    }
 }
 
 #[derive(Default, Clone, Debug)]

turbopack/crates/turbo-tasks/src/capture_future.rs

@@ -31,8 +31,7 @@ pin_project! {
         #[pin]
         future: F,
         duration: Duration,
-        allocations: usize,
-        deallocations: usize,
+        allocations: AllocationInfo,
     }
 }
 
@@ -41,8 +40,7 @@ impl<T, F: Future<Output = T>> CaptureFuture<T, F> {
         Self {
             future,
             duration: Duration::ZERO,
-            allocations: 0,
-            deallocations: 0,
+            allocations: AllocationInfo::ZERO,
         }
     }
 }
@@ -77,14 +75,25 @@ pub struct TurboTasksPanic {
     pub location: Option<String>,
 }
 
+impl TurboTasksPanic {
+    pub fn into_panic(self) -> Box<dyn std::any::Any + Send> {
+        Box::new(format!(
+            "{} at {}",
+            self.message,
+            self.location
+                .unwrap_or_else(|| "unknown location".to_string())
+        ))
+    }
+}
+
 impl Display for TurboTasksPanic {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", self.message)
     }
 }
 
 impl<T, F: Future<Output = T>> Future for CaptureFuture<T, F> {
-    type Output = (Result<T, TurboTasksPanic>, Duration, usize);
+    type Output = (Result<T, TurboTasksPanic>, Duration, AllocationInfo);
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let this = self.project();
@@ -127,17 +136,10 @@ impl<T, F: Future<Output = T>> Future for CaptureFuture<T, F> {
         let elapsed = start.elapsed();
         let allocations = start_allocations.until_now();
         *this.duration += elapsed + data.duration;
-        *this.allocations += allocations.allocations + data.allocations;
-        *this.deallocations += allocations.deallocations + data.deallocations;
+        *this.allocations += allocations;
         match result {
-            Err(err) => {
-                let memory_usage = this.allocations.saturating_sub(*this.deallocations);
-                Poll::Ready((Err(err), *this.duration, memory_usage))
-            }
-            Ok(Poll::Ready(r)) => {
-                let memory_usage = this.allocations.saturating_sub(*this.deallocations);
-                Poll::Ready((Ok(r), *this.duration, memory_usage))
-            }
+            Err(err) => Poll::Ready((Err(err), *this.duration, this.allocations.clone())),
+            Ok(Poll::Ready(r)) => Poll::Ready((Ok(r), *this.duration, this.allocations.clone())),
             Ok(Poll::Pending) => Poll::Pending,
         }
     }

turbopack/crates/turbo-tasks/src/lib.rs

@@ -76,6 +76,7 @@ pub mod registry;
 pub mod scope;
 mod serialization_invalidation;
 pub mod small_duration;
+mod spawn;
 mod state;
 pub mod task;
 mod task_execution_reason;
@@ -110,8 +111,7 @@ pub use manager::{
     CurrentCellRef, ReadConsistency, TaskPersistence, TurboTasks, TurboTasksApi,
     TurboTasksBackendApi, TurboTasksBackendApiExt, TurboTasksCallApi, Unused, UpdateInfo,
     dynamic_call, emit, mark_finished, mark_root, mark_session_dependent, mark_stateful,
-    prevent_gc, run_once, run_once_with_reason, spawn_blocking, spawn_thread, trait_call,
-    turbo_tasks, turbo_tasks_scope,
+    prevent_gc, run_once, run_once_with_reason, trait_call, turbo_tasks, turbo_tasks_scope,
 };
 pub use output::OutputContent;
 pub use raw_vc::{CellId, RawVc, ReadRawVcFuture, ResolveTypeError};
@@ -120,6 +120,7 @@ pub use read_ref::ReadRef;
 use rustc_hash::FxHasher;
 pub use serialization_invalidation::SerializationInvalidator;
 pub use shrink_to_fit::ShrinkToFit;
+pub use spawn::{JoinHandle, spawn, spawn_blocking, spawn_thread};
 pub use state::{State, TransientState};
 pub use task::{SharedReference, TypedSharedReference, task_input::TaskInput};
 pub use task_execution_reason::TaskExecutionReason;

turbopack/crates/turbo-tasks/src/manager.rs

@@ -8,7 +8,6 @@ use std::{
         Arc, Mutex, RwLock, Weak,
         atomic::{AtomicBool, AtomicUsize, Ordering},
     },
-    thread,
     time::{Duration, Instant},
 };
 
@@ -17,10 +16,9 @@ use auto_hash_map::AutoMap;
 use rustc_hash::FxHasher;
 use serde::{Deserialize, Serialize};
 use smallvec::SmallVec;
-use tokio::{runtime::Handle, select, sync::mpsc::Receiver, task_local};
+use tokio::{select, sync::mpsc::Receiver, task_local};
 use tokio_util::task::TaskTracker;
-use tracing::{Instrument, Level, Span, info_span, instrument, trace_span};
-use turbo_tasks_malloc::TurboMalloc;
+use tracing::{Instrument, Level, instrument, trace_span};
 
 use crate::{
     Completion, InvalidationReason, InvalidationReasonSet, OutputContent, ReadCellOptions,
@@ -30,7 +28,7 @@ use crate::{
         Backend, CachedTaskType, CellContent, TaskCollectiblesMap, TaskExecutionSpec,
         TransientTaskType, TurboTasksExecutionError, TypedCellContent,
     },
-    capture_future::{self, CaptureFuture},
+    capture_future::CaptureFuture,
     event::{Event, EventListener},
     id::{BackendJobId, ExecutionId, LocalTaskId, TRANSIENT_TASK_BIT, TraitTypeId},
     id_factory::IdFactoryWithReuse,
@@ -718,7 +716,7 @@ impl<B: Backend + 'static> TurboTasks<B> {
                     };
 
                     async {
-                        let (result, duration, memory_usage) = CaptureFuture::new(future).await;
+                        let (result, duration, alloc_info) = CaptureFuture::new(future).await;
 
                         // wait for all spawned local tasks using `local` to finish
                         let ltt = CURRENT_TASK_STATE
@@ -742,7 +740,7 @@ impl<B: Backend + 'static> TurboTasks<B> {
                         let schedule_again = this.backend.task_execution_completed(
                             task_id,
                             duration,
-                            memory_usage,
+                            alloc_info.memory_usage(),
                             &cell_counters,
                             stateful,
                             has_invalidator,
@@ -1794,35 +1792,6 @@ pub fn emit<T: VcValueTrait + ?Sized>(collectible: ResolvedVc<T>) {
     })
 }
 
-pub async fn spawn_blocking<T: Send + 'static>(func: impl FnOnce() -> T + Send + 'static) -> T {
-    let turbo_tasks = turbo_tasks();
-    let span = Span::current();
-    let (result, duration, alloc_info) = tokio::task::spawn_blocking(|| {
-        let _guard = span.entered();
-        let start = Instant::now();
-        let start_allocations = TurboMalloc::allocation_counters();
-        let r = turbo_tasks_scope(turbo_tasks, func);
-        (r, start.elapsed(), start_allocations.until_now())
-    })
-    .await
-    .unwrap();
-    capture_future::add_duration(duration);
-    capture_future::add_allocation_info(alloc_info);
-    result
-}
-
-pub fn spawn_thread(func: impl FnOnce() + Send + 'static) {
-    let handle = Handle::current();
-    let span = info_span!("thread").or_current();
-    thread::spawn(move || {
-        let span = span.entered();
-        let guard = handle.enter();
-        func();
-        drop(guard);
-        drop(span);
-    });
-}
-
 pub(crate) async fn read_task_output(
     this: &dyn TurboTasksApi,
     id: TaskId,

turbopack/crates/turbo-tasks/src/spawn.rs

@@ -0,0 +1,89 @@
+use std::{
+    panic::resume_unwind,
+    pin::Pin,
+    task::{Context, Poll},
+    thread,
+    time::{Duration, Instant},
+};
+
+use anyhow::Result;
+use futures::{FutureExt, ready};
+use tokio::runtime::Handle;
+use tracing::{Instrument, Span, info_span};
+use turbo_tasks_malloc::{AllocationInfo, TurboMalloc};
+
+use crate::{
+    TurboTasksPanic,
+    capture_future::{self, CaptureFuture},
+    manager::turbo_tasks_future_scope,
+    turbo_tasks, turbo_tasks_scope,
+};
+
+pub struct JoinHandle<T> {
+    join_handle: tokio::task::JoinHandle<(Result<T, TurboTasksPanic>, Duration, AllocationInfo)>,
+}
+
+impl<T> Future for JoinHandle<T> {
+    type Output = T;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+        match ready!(this.join_handle.poll_unpin(cx)) {
+            Ok((res, duration, alloc_info)) => {
+                capture_future::add_duration(duration);
+                capture_future::add_allocation_info(alloc_info);
+                match res {
+                    Ok(res) => Poll::Ready(res),
+                    Err(e) => resume_unwind(e.into_panic()),
+                }
+            }
+            Err(e) => resume_unwind(e.into_panic()),
+        }
+    }
+}
+
+/// Spawns a future as separate task and returns a JoinHandle which can be used to await the result.
+/// The future has access to the current TurboTasks context and runs in the same tracing span.
+/// Allocations and cpu time is accounted to the current turbo-tasks function.
+pub fn spawn<T: Send + 'static>(future: impl Future<Output = T> + Send + 'static) -> JoinHandle<T> {
+    let turbo_tasks = turbo_tasks();
+    let span = Span::current();
+    let join_handle = tokio::task::spawn(
+        turbo_tasks_future_scope(turbo_tasks, CaptureFuture::new(future)).instrument(span),
+    );
+    JoinHandle { join_handle }
+}
+
+/// Spawns a blocking function in a separate task using the blocking pool and returns a JoinHandle
+/// which can be used to await the result. The function has access to the current TurboTasks context
+/// and runs in the same tracing span.
+/// Allocations and cpu time is accounted to the current turbo-tasks function.
+pub fn spawn_blocking<T: Send + 'static>(
+    func: impl FnOnce() -> T + Send + 'static,
+) -> JoinHandle<T> {
+    let turbo_tasks = turbo_tasks();
+    let span = Span::current();
+    let join_handle = tokio::task::spawn_blocking(|| {
+        let _guard = span.entered();
+        let start = Instant::now();
+        let start_allocations = TurboMalloc::allocation_counters();
+        let r = turbo_tasks_scope(turbo_tasks, func);
+        (Ok(r), start.elapsed(), start_allocations.until_now())
+    });
+    JoinHandle { join_handle }
+}
+
+/// Spawns a thread which runs in background. It has access to the current TurboTasks context, but
+/// is not accounted towards the current turbo-tasks function.
+pub fn spawn_thread(func: impl FnOnce() + Send + 'static) {
+    let handle = Handle::current();
+    let span = info_span!("thread").or_current();
+    let turbo_tasks = turbo_tasks();
+    thread::spawn(move || {
+        let _span = span.entered();
+        turbo_tasks_scope(turbo_tasks, || {
+            let _guard = handle.enter();
+            func();
+        })
+    });
+}