scoped blockign

Author: kariy

crates/sync/stage/src/blocks/mod.rs

@@ -13,6 +13,7 @@ use katana_primitives::transaction::{Tx, TxWithHash};
 use katana_primitives::Felt;
 use katana_provider::api::block::{BlockHashProvider, BlockWriter};
 use katana_provider::{MutableProvider, ProviderError, ProviderFactory};
+use katana_tasks::TaskSpawner;
 use num_traits::ToPrimitive;
 use starknet::core::types::ResourcePrice;
 use tracing::{error, info_span, Instrument};
@@ -28,19 +29,20 @@ pub use downloader::{BatchBlockDownloader, BlockDownloader};
 pub struct Blocks<PF, B> {
     provider: PF,
     downloader: B,
+    task_spawner: TaskSpawner,
 }
 
 impl<PF, B> Blocks<PF, B> {
     /// Create a new [`Blocks`] stage.
-    pub fn new(provider: PF, downloader: B) -> Self {
-        Self { provider, downloader }
+    pub fn new(provider: PF, downloader: B, task_spawner: TaskSpawner) -> Self {
+        Self { provider, downloader, task_spawner }
     }
 
     /// Validates that the downloaded blocks form a valid chain.
     ///
     /// This method checks the chain invariant: block N's parent hash must be block N-1's hash.
     /// For the first block in the list (if not block 0), it fetches the parent hash from storage.
-    fn validate_chain_invariant(&self, blocks: &[StateUpdateWithBlock]) -> Result<(), Error>
+    async fn validate_chain_invariant(&self, blocks: &[StateUpdateWithBlock]) -> Result<(), Error>
     where
         PF: ProviderFactory,
         <PF as ProviderFactory>::Provider: BlockHashProvider,
@@ -71,7 +73,27 @@ impl<PF, B> Blocks<PF, B> {
             }
         }
 
-        // Validate the rest of the blocks in the list
+        // self.task_spawner.cpu_bound().spawn(|| {
+        //     // Validate the rest of the blocks in the list
+        //     for window in blocks.windows(2) {
+        //         let prev_block = &window[0].block;
+        //         let curr_block = &window[1].block;
+
+        //         let prev_hash = prev_block.block_hash.unwrap_or_default();
+        //         let curr_block_num = curr_block.block_number.unwrap_or_default();
+
+        //         if curr_block.parent_block_hash != prev_hash {
+        //             return Err(Error::ChainInvariantViolation {
+        //                 block_num: curr_block_num,
+        //                 parent_hash: curr_block.parent_block_hash,
+        //                 expected_hash: prev_hash,
+        //             });
+        //         }
+        //     }
+
+        //     Ok(())
+        // });
+
         for window in blocks.windows(2) {
             let prev_block = &window[0].block;
             let curr_block = &window[1].block;
@@ -112,7 +134,10 @@ where
                 .await
                 .map_err(Error::Gateway)?;
 
-            self.validate_chain_invariant(&blocks)?;
+            self.task_spawner
+                .scope(|s| s.spawn(self.validate_chain_invariant(&blocks)))
+                .await
+                .unwrap()?;
 
             let span = info_span!(target: "stage", "blocks.insert", from = %input.from(), to = %input.to());
             let _enter = span.enter();

crates/tasks/src/blocking.rs

@@ -86,6 +86,14 @@ impl CpuBlockingTaskPool {
         });
         CpuBlockingJoinHandle { inner: rx }
     }
+
+    pub(crate) fn scope<'scope, F, R>(&self, func: F) -> R
+    where
+        F: for<'s> FnOnce(&'s rayon::Scope<'scope>) -> R + Send + 'scope,
+        R: Send + 'scope,
+    {
+        self.pool.scope(func)
+    }
 }
 
 #[derive(Debug)]

crates/tasks/src/scope.rs

@@ -2,7 +2,7 @@ use core::future::Future;
 use core::marker::PhantomData;
 use core::pin::Pin;
 use std::panic::{self, AssertUnwindSafe};
-use std::sync::Arc;
+use std::sync::{mpsc, Arc};
 use std::task::Poll;
 
 use futures::stream::FuturesUnordered;
@@ -50,7 +50,29 @@ impl<'scope> ScopedTaskSpawner<'scope> {
         ScopedJoinHandle { receiver, _marker: PhantomData }
     }
 
+    /// Returns a scoped spawner for CPU-bound work.
+    pub fn cpu_bound(&self) -> ScopedCpuTaskSpawner<'scope> {
+        ScopedCpuTaskSpawner {
+            inner: self.inner.clone(),
+            sender: self.sender.clone(),
+            _marker: PhantomData,
+        }
+    }
+
     /// Spawn a blocking task that can borrow data with lifetime `'scope`.
+    ///
+    /// # Runtime requirement
+    ///
+    /// This uses [`tokio::task::block_in_place`] under the hood, which only works on
+    /// Tokio's multi-threaded runtime. On a current-thread runtime this will panic. If
+    /// you need to support a current-thread runtime, prefer:
+    ///
+    /// 1. Moving/cloning the data so the closure can be `'static` and using the non-scoped
+    ///    `TaskSpawner::spawn_blocking`, or
+    /// 2. Running the work inline (accepting that it will block the current thread).
+    ///
+    /// The scoped variant exists to let you borrow non-`'static` data; that capability
+    /// currently depends on `block_in_place`.
     pub fn spawn_blocking<F, R>(&self, func: F) -> ScopedJoinHandle<'scope, R>
     where
         F: FnOnce() -> R + Send + 'scope,
@@ -83,6 +105,67 @@ impl<'scope> ScopedTaskSpawner<'scope> {
     }
 }
 
+/// Spawner for CPU-bound scoped tasks executed on the dedicated blocking pool.
+#[derive(Debug, Clone)]
+pub struct ScopedCpuTaskSpawner<'scope> {
+    pub(crate) inner: Arc<TaskManagerInner>,
+    pub(crate) sender: UnboundedSender<BoxScopedFuture<'scope>>,
+    pub(crate) _marker: PhantomData<&'scope ()>,
+}
+
+impl<'scope> ScopedCpuTaskSpawner<'scope> {
+    /// Spawn a CPU-bound task that can borrow data with lifetime `'scope`.
+    ///
+    /// # Runtime requirement
+    ///
+    /// Uses [`tokio::task::block_in_place`] to enter the blocking pool. This only works on
+    /// Tokio's multithreaded runtime. On a current-thread runtime this will panic. If you
+    /// must support a current-thread runtime, either move/clone data into a `'static` closure
+    /// and use the non-scoped `CPUBoundTaskSpawner`, or run the work inline knowing it will
+    /// block the thread.
+    pub fn spawn<F, R>(&self, func: F) -> ScopedJoinHandle<'scope, R>
+    where
+        F: FnOnce() -> R + Send + 'scope,
+        R: Send + 'scope,
+    {
+        let (tx, rx) = oneshot::channel();
+        let mut receiver = rx;
+        let inner = self.inner.clone();
+
+        let task = async move {
+            let result = block_in_place(|| {
+                if inner.on_cancel.is_cancelled() {
+                    return Err(JoinError::Cancelled);
+                }
+
+                let (res_tx, res_rx) = mpsc::channel();
+                inner.blocking_pool.scope(|scope| {
+                    scope.spawn(|_| {
+                        let _ = res_tx.send(panic::catch_unwind(AssertUnwindSafe(func)));
+                    });
+                });
+
+                match res_rx.recv() {
+                    Ok(Ok(value)) => Ok(value),
+                    Ok(Err(panic)) => Err(JoinError::Panic(panic)),
+                    Err(..) => Err(JoinError::Cancelled),
+                }
+            });
+
+            let _ = tx.send(result);
+        };
+
+        let task = self.inner.tracker.track_future(task);
+        let task = Box::pin(task);
+
+        if self.sender.send(task).is_err() {
+            receiver.close();
+        }
+
+        ScopedJoinHandle { receiver, _marker: PhantomData }
+    }
+}
+
 pub(crate) struct TaskScope<'scope, R>
 where
     R: Send + 'scope,

crates/tasks/src/spawner.rs

@@ -58,21 +58,26 @@ impl TaskSpawner {
     /// Runs a scoped block in which tasks may borrow non-`'static` data but are guaranteed to be
     /// completed before this method returns.
     ///
-    /// # Examples
+    /// Unlike [`TaskSpawner::spawn`], scoped tasks don't need to be `'static`. Without a scope, the
+    /// compiler prevents you from capturing stack data because the task could outlive the
+    /// borrow:
     ///
-    /// What happens if you try to capture non-`'static` data without a scope?
     /// ```compile_fail
     /// # use katana_tasks::TaskManager;
     /// # #[tokio::main] async fn main() {
     /// let manager = TaskManager::current();
     /// let spawner = manager.task_spawner();
-    /// let mut prefix = String::from("hi ");
+    /// let prefix = String::from("Hello world!");
     ///
     /// // ERROR: `spawn` requires futures to be `'static`.
-    /// spawner.spawn(async { prefix.push_str("there") });
+    /// spawner.spawn(async { println!("{prefix}") });
     /// # }
     /// ```
     ///
+    /// A scope keeps everything within the lifetime of the call, so you can borrow safely.
+    ///
+    /// # Examples
+    ///
     /// Borrow stack data in an async task:
     /// ```
     /// # use katana_tasks::TaskManager;
@@ -119,7 +124,7 @@ impl TaskSpawner {
     /// # }
     /// ```
     ///
-    /// Use scoped blocking work with borrowed data:
+    /// Use scoped blocking work with borrowed data (requires a multithreaded Tokio runtime):
     /// ```
     /// # use katana_tasks::TaskManager;
     /// # use std::sync::{Arc, Mutex};
@@ -422,36 +427,70 @@ mod tests {
         let spawner = manager.task_spawner();
 
         let text = String::from("scoped");
-        let expected_len = text.len();
 
-        spawner
+        let returned_ref = spawner
             .scope(|scope| {
                 let text_ref: &String = &text;
+                scope.spawn(async move { text_ref })
+            })
+            .await
+            .unwrap();
+
+        // original value is still valid after scope returns
+        assert_eq!(text, *returned_ref);
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn scoped_spawn_blocking_allows_borrowed_data() {
+        let manager = TaskManager::current();
+        let spawner = manager.task_spawner();
+
+        let counter = AtomicUsize::new(0);
+
+        spawner
+            .scope(|scope| {
+                let counter_ref = &counter;
 
                 async move {
-                    let handle = scope.spawn(async move { text_ref.len() });
-                    assert_eq!(handle.await.unwrap(), expected_len);
+                    let handle = scope.spawn_blocking(move || {
+                        counter_ref.fetch_add(1, Ordering::SeqCst);
+                        counter_ref.load(Ordering::SeqCst)
+                    });
+
+                    assert_eq!(handle.await.unwrap(), 1);
                 }
             })
             .await;
 
-        // original value is still valid after scope returns
-        assert_eq!(text.len(), expected_len);
+        assert_eq!(counter.load(Ordering::SeqCst), 1);
+
+        let mut prefix = String::from("hello ");
+        spawner
+            .scope(|scope| {
+                scope.spawn(async {
+                    prefix.push_str("world!");
+                })
+            })
+            .await
+            .unwrap();
+
+        assert_eq!(prefix.as_str(), "hello world!")
     }
 
-    #[tokio::test]
-    async fn scoped_spawn_blocking_allows_borrowed_data() {
+    #[tokio::test(flavor = "multi_thread")]
+    async fn scoped_cpu_bound_allows_borrowed_data() {
         let manager = TaskManager::current();
         let spawner = manager.task_spawner();
 
         let counter = AtomicUsize::new(0);
 
         spawner
             .scope(|scope| {
+                let cpu = scope.cpu_bound();
                 let counter_ref = &counter;
 
                 async move {
-                    let handle = scope.spawn_blocking(move || {
+                    let handle = cpu.spawn(move || {
                         counter_ref.fetch_add(1, Ordering::SeqCst);
                         counter_ref.load(Ordering::SeqCst)
                     });