diff --git a/crates/common/src/context.rs b/crates/common/src/context.rs index a666a59e..fad36f15 100644 --- a/crates/common/src/context.rs +++ b/crates/common/src/context.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use futures::{ AsyncRead, AsyncWrite, future::{self, BoxFuture, Either}, + stream::{self, StreamExt, TryStreamExt}, }; #[cfg(any(test, feature = "test-utils"))] @@ -19,7 +20,12 @@ pub use test::{ test_mt_context, test_mt_context_with_spawn, test_st_context, }; -use crate::{ContextId, executor::Inner, io::Io, mux::Mux}; +use crate::{ContextId, io::Io, mux::Mux, thread_pool::ThreadPool}; + +/// Default maximum number of [`map`](Context::map) items processed +/// concurrently. Both parties must agree on this value, so it is a fixed +/// constant rather than data- or timing-dependent. +pub const DEFAULT_CONCURRENCY_LIMIT: usize = 32; /// A task execution context. /// @@ -42,42 +48,31 @@ enum Mode { Single, Multi { mux: Arc, - executor: Option>, + /// Pool for parallel execution; `None` runs sub-tasks cooperatively + /// on the caller's future. + pool: Option, + /// Maximum number of [`map`](Context::map) items processed + /// concurrently. + concurrency_limit: usize, }, } impl std::fmt::Debug for Context { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mode = match &self.mode { - Mode::Single => "single", - Mode::Multi { - executor: Some(_), .. - } => "multi-threaded", - Mode::Multi { executor: None, .. } => "multi-cooperative", - }; f.debug_struct("Context") .field("id", &self.id) .field("io", &self.io) - .field("mode", &mode) - .finish() + .finish_non_exhaustive() } } impl Context { - /// Creates a new context that uses `mux` to allocate a channel per - /// sub-task. - /// - /// Sub-tasks are executed cooperatively on the calling future. For - /// parallel execution, build an [`Executor`](crate::Executor) and use - /// [`Executor::new_context`](crate::Executor::new_context) instead. - pub fn new(mux: M) -> Result { - Self::with_prefix(mux, ContextId::default()) - } - /// Creates a new context backed by a single I/O channel. /// /// Sub-tasks spawned via [`join`], [`try_join`], [`map`] etc. share the - /// channel and run **sequentially** in the order given. + /// channel and run **sequentially** in the order given. For parallel + /// execution, build a [`Session`](crate::Session) and use + /// [`Session::new_context`](crate::Session::new_context) instead. /// /// [`join`]: Self::join /// [`try_join`]: Self::try_join @@ -98,45 +93,32 @@ impl Context { } } - /// Like [`Context::new`], but namespaces all channels under `prefix` so - /// several sub-protocols can share a mux without colliding. - pub fn with_prefix( - mux: M, - prefix: impl AsRef<[u8]>, - ) -> Result { - let mux: Arc = Arc::new(mux); - let id = ContextId::from_prefix(prefix); - let io = mux.open(id.as_ref()).map_err(ContextError::mux)?; - Ok(Self { - id, - io, - mode: Mode::Multi { - mux, - executor: None, - }, - fork_counter: 0, - }) - } - - pub(crate) fn with_executor( + pub(crate) fn for_session( id: ContextId, io: Io, mux: Arc, - executor: Arc, + pool: Option, + concurrency_limit: usize, ) -> Self { Self { id, io, mode: Mode::Multi { mux, - executor: Some(executor), + pool, + concurrency_limit, }, fork_counter: 0, } } fn child(&self, id: ContextId) -> Result { - let Mode::Multi { mux, executor } = &self.mode else { + let Mode::Multi { + mux, + pool, + concurrency_limit, + } = &self.mode + else { unreachable!("child() called on a single-channel context"); }; let io = mux.open(id.as_ref()).map_err(ContextError::mux)?; @@ -145,7 +127,8 @@ impl Context { io, mode: Mode::Multi { mux: mux.clone(), - executor: executor.clone(), + pool: pool.clone(), + concurrency_limit: *concurrency_limit, }, fork_counter: 0, }) @@ -180,27 +163,50 @@ impl Context { T: Send + 'static, R: Send + 'static, { - if matches!(self.mode, Mode::Single) { - let mut results = Vec::with_capacity(items.len()); - for item in items { - results.push(f(self, item).await); + let (mux, pool, concurrency_limit) = match &self.mode { + Mode::Single => { + let mut results = Vec::with_capacity(items.len()); + for item in items { + results.push(f(self, item).await); + } + return Ok(results); } - return Ok(results); - } + Mode::Multi { + mux, + pool, + concurrency_limit, + } => (mux.clone(), pool.clone(), *concurrency_limit), + }; let parent_id = self.next_fork(); - let executor = self.executor().cloned(); - let mut tasks = Vec::with_capacity(items.len()); - for (i, item) in items.into_iter().enumerate() { - let i = u32::try_from(i).expect("more than u32::MAX items"); - let mut ctx = self.child(parent_id.child(i))?; - let f = f.clone(); - tasks.push(run( - executor.as_ref(), - async move { f(&mut ctx, item).await }, - )); - } - Ok(future::join_all(tasks).await) + + // Each item lazily opens its own channel only once `buffered` polls it, + // so at most `limit` channels are open at any time. Channel IDs stay + // keyed by item index and results are yielded in input order, so the + // bound changes neither the wire protocol nor the output ordering. + stream::iter(items.into_iter().enumerate()) + .map(move |(i, item)| { + let i = u32::try_from(i).expect("more than u32::MAX items"); + let id = parent_id.child(i); + let (mux, pool, f) = (mux.clone(), pool.clone(), f.clone()); + async move { + let io = mux.open(id.as_ref()).map_err(ContextError::mux)?; + let mut ctx = Context { + id, + io, + mode: Mode::Multi { + mux, + pool: pool.clone(), + concurrency_limit, + }, + fork_counter: 0, + }; + Ok(run(pool.as_ref(), async move { f(&mut ctx, item).await }).await) + } + }) + .buffered(concurrency_limit) + .try_collect() + .await } /// Runs `a` and `b` concurrently and returns both results. @@ -218,12 +224,12 @@ impl Context { } let parent_id = self.next_fork(); - let executor = self.executor().cloned(); + let pool = self.pool().cloned(); let mut ctx_a = self.child(parent_id.child(0))?; let mut ctx_b = self.child(parent_id.child(1))?; - let task_a = run(executor.as_ref(), async move { a(&mut ctx_a).await }); - let task_b = run(executor.as_ref(), async move { b(&mut ctx_b).await }); + let task_a = run(pool.as_ref(), async move { a(&mut ctx_a).await }); + let task_b = run(pool.as_ref(), async move { b(&mut ctx_b).await }); Ok(future::join(task_a, task_b).await) } @@ -251,12 +257,12 @@ impl Context { } let parent_id = self.next_fork(); - let executor = self.executor().cloned(); + let pool = self.pool().cloned(); let mut ctx_a = self.child(parent_id.child(0))?; let mut ctx_b = self.child(parent_id.child(1))?; - let task_a = run(executor.as_ref(), async move { a(&mut ctx_a).await }); - let task_b = run(executor.as_ref(), async move { b(&mut ctx_b).await }); + let task_a = run(pool.as_ref(), async move { a(&mut ctx_a).await }); + let task_b = run(pool.as_ref(), async move { b(&mut ctx_b).await }); Ok(future::try_join(task_a, task_b).await) } @@ -287,14 +293,14 @@ impl Context { } let parent_id = self.next_fork(); - let executor = self.executor().cloned(); + let pool = self.pool().cloned(); let mut ctx_a = self.child(parent_id.child(0))?; let mut ctx_b = self.child(parent_id.child(1))?; let mut ctx_c = self.child(parent_id.child(2))?; - let task_a = run(executor.as_ref(), async move { a(&mut ctx_a).await }); - let task_b = run(executor.as_ref(), async move { b(&mut ctx_b).await }); - let task_c = run(executor.as_ref(), async move { c(&mut ctx_c).await }); + let task_a = run(pool.as_ref(), async move { a(&mut ctx_a).await }); + let task_b = run(pool.as_ref(), async move { b(&mut ctx_b).await }); + let task_c = run(pool.as_ref(), async move { c(&mut ctx_c).await }); Ok(future::try_join3(task_a, task_b, task_c).await) } @@ -329,40 +335,37 @@ impl Context { } let parent_id = self.next_fork(); - let executor = self.executor().cloned(); + let pool = self.pool().cloned(); let mut ctx_a = self.child(parent_id.child(0))?; let mut ctx_b = self.child(parent_id.child(1))?; let mut ctx_c = self.child(parent_id.child(2))?; let mut ctx_d = self.child(parent_id.child(3))?; - let task_a = run(executor.as_ref(), async move { a(&mut ctx_a).await }); - let task_b = run(executor.as_ref(), async move { b(&mut ctx_b).await }); - let task_c = run(executor.as_ref(), async move { c(&mut ctx_c).await }); - let task_d = run(executor.as_ref(), async move { d(&mut ctx_d).await }); + let task_a = run(pool.as_ref(), async move { a(&mut ctx_a).await }); + let task_b = run(pool.as_ref(), async move { b(&mut ctx_b).await }); + let task_c = run(pool.as_ref(), async move { c(&mut ctx_c).await }); + let task_d = run(pool.as_ref(), async move { d(&mut ctx_d).await }); Ok(future::try_join4(task_a, task_b, task_c, task_d).await) } - fn executor(&self) -> Option<&Arc> { - if let Mode::Multi { executor, .. } = &self.mode { - executor.as_ref() + fn pool(&self) -> Option<&ThreadPool> { + if let Mode::Multi { pool, .. } = &self.mode { + pool.as_ref() } else { None } } } -/// Spawns `fut` on `executor` if one is provided, otherwise yields the future +/// Spawns `fut` on `pool` if one is provided, otherwise yields the future /// as-is. The output type is identical either way. -fn run( - executor: Option<&Arc>, - fut: F, -) -> impl std::future::Future + Send +fn run(pool: Option<&ThreadPool>, fut: F) -> impl std::future::Future + Send where F: std::future::Future + Send + 'static, F::Output: Send + 'static, { - match executor { - Some(exec) => Either::Left(crate::executor::spawn_on(exec, fut)), + match pool { + Some(pool) => Either::Left(crate::thread_pool::spawn_on(pool, fut)), None => Either::Right(fut), } } diff --git a/crates/common/src/context/test/helpers.rs b/crates/common/src/context/test/helpers.rs index 9a6f1bb1..8c79f5ef 100644 --- a/crates/common/src/context/test/helpers.rs +++ b/crates/common/src/context/test/helpers.rs @@ -4,9 +4,10 @@ use serio::channel::duplex; use crate::{ context::Context, - executor::{Executor, ExecutorBuilder}, io::Io, mux::test_framed_mux, + session::{Session, SessionBuilder}, + thread_pool::ThreadPool, }; /// Creates a pair of single-threaded contexts using memory I/O channels. @@ -19,19 +20,19 @@ pub fn test_st_context(io_buffer: usize) -> (Context, Context) { ) } -/// Creates a pair of multi-threaded executors sharing multiplexed I/O channels. -pub fn test_mt_context(io_buffer: usize) -> (Executor, Executor) { +/// Creates a pair of multi-threaded sessions sharing multiplexed I/O channels. +pub fn test_mt_context(io_buffer: usize) -> (Session, Session) { let (mux_0, mux_1) = test_framed_mux(io_buffer); ( - ExecutorBuilder::default().build(mux_0), - ExecutorBuilder::default().build(mux_1), + SessionBuilder::default().build(mux_0).unwrap(), + SessionBuilder::default().build(mux_1).unwrap(), ) } /// Like [`test_mt_context`], but uses a custom worker spawn callback (e.g. /// `web_spawn::spawn` on wasm). -pub fn test_mt_context_with_spawn(io_buffer: usize, spawn: F) -> (Executor, Executor) +pub fn test_mt_context_with_spawn(io_buffer: usize, spawn: F) -> (Session, Session) where F: Fn(Box) -> Result<(), std::io::Error> + Clone @@ -40,9 +41,11 @@ where + 'static, { let (mux_0, mux_1) = test_framed_mux(io_buffer); + let pool_0 = ThreadPool::builder().spawn(spawn.clone()).build().unwrap(); + let pool_1 = ThreadPool::builder().spawn(spawn).build().unwrap(); ( - ExecutorBuilder::default().spawn(spawn.clone()).build(mux_0), - ExecutorBuilder::default().spawn(spawn).build(mux_1), + SessionBuilder::default().pool(pool_0).build(mux_0).unwrap(), + SessionBuilder::default().pool(pool_1).build(mux_1).unwrap(), ) } diff --git a/crates/common/src/context/test/recording.rs b/crates/common/src/context/test/recording.rs index 7d30c6ff..04e48dd3 100644 --- a/crates/common/src/context/test/recording.rs +++ b/crates/common/src/context/test/recording.rs @@ -11,11 +11,11 @@ use futures::{AsyncRead, AsyncWrite}; use tokio_util::compat::{Compat, TokioAsyncReadCompatExt}; use crate::{ - ContextId, context::Context, - executor::{Executor, ExecutorBuilder}, io::Io, mux::Mux, + session::{Session, SessionBuilder}, + thread_pool::ThreadPool, }; /// A duplex stream that records all bytes written. @@ -68,46 +68,6 @@ impl AsyncWrite for RecordingDuplex { } } -/// A simple mux that wraps a single I/O stream for recording tests. -struct SingleChannelMux { - io: Mutex>, - max_frame_length: Option, -} - -impl SingleChannelMux { - fn new(io: I, max_frame_length: Option) -> Self { - Self { - io: Mutex::new(Some(io)), - max_frame_length, - } - } -} - -impl Mux for SingleChannelMux -where - I: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static, -{ - fn open(&self, id: &[u8]) -> Result { - // Only allow opening the root ID - if id != ContextId::default().as_bytes() { - return Err(std::io::Error::other( - "single channel mux only supports root ID", - )); - } - let io = self - .io - .lock() - .unwrap() - .take() - .ok_or_else(|| std::io::Error::other("channel already opened"))?; - if let Some(limit) = self.max_frame_length { - Ok(Io::from_io_with_limit(io, limit)) - } else { - Ok(Io::from_io(io)) - } - } -} - /// Creates a pair of single-threaded contexts where writes from ctx_1 to ctx_0 /// are recorded. /// @@ -124,12 +84,9 @@ pub fn recording_st_context(io_buffer: usize) -> (Context, Context, Arc (Executor, Executor, Arc>) { +pub fn recording_mt_context(io_buffer: usize) -> (Session, Session, Arc>) { let (mux_0, mux_1, recorded) = recording_test_mux(io_buffer, None); ( - ExecutorBuilder::default().build(mux_0), - ExecutorBuilder::default().build(mux_1), + SessionBuilder::default().build(mux_0).unwrap(), + SessionBuilder::default().build(mux_1).unwrap(), recorded, ) } @@ -425,12 +379,12 @@ pub fn recording_mt_context(io_buffer: usize) -> (Executor, Executor, Arc (Executor, Executor, Arc>) { +) -> (Session, Session, Arc>) { let (mux_0, mux_1, recorded) = recording_test_mux(io_buffer, Some(max_frame_length)); ( - ExecutorBuilder::default().build(mux_0), - ExecutorBuilder::default().build(mux_1), + SessionBuilder::default().build(mux_0).unwrap(), + SessionBuilder::default().build(mux_1).unwrap(), recorded, ) } @@ -438,13 +392,13 @@ pub fn recording_mt_context_with_limit( /// Like [`recording_mt_context_with_limit`], but uses a custom worker spawn /// callback (e.g. `web_spawn::spawn` on wasm) and a fixed concurrency level. /// -/// The same `spawn` callback is used for both executors. +/// The same `spawn` callback is used for both sessions. pub fn recording_mt_context_with_spawn_and_limit( io_buffer: usize, max_frame_length: usize, concurrency: usize, spawn: F, -) -> (Executor, Executor, Arc>) +) -> (Session, Session, Arc>) where F: Fn(Box) -> Result<(), std::io::Error> + Clone @@ -453,13 +407,17 @@ where + 'static, { let (mux_0, mux_1, recorded) = recording_test_mux(io_buffer, Some(max_frame_length)); - let exec_0 = ExecutorBuilder::default() + let pool_0 = ThreadPool::builder() .num_threads(concurrency) .spawn(spawn.clone()) - .build(mux_0); - let exec_1 = ExecutorBuilder::default() + .build() + .unwrap(); + let pool_1 = ThreadPool::builder() .num_threads(concurrency) .spawn(spawn) - .build(mux_1); + .build() + .unwrap(); + let exec_0 = SessionBuilder::default().pool(pool_0).build(mux_0).unwrap(); + let exec_1 = SessionBuilder::default().pool(pool_1).build(mux_1).unwrap(); (exec_0, exec_1, recorded) } diff --git a/crates/common/src/context/test/replay.rs b/crates/common/src/context/test/replay.rs index 89880384..ef010ad2 100644 --- a/crates/common/src/context/test/replay.rs +++ b/crates/common/src/context/test/replay.rs @@ -10,9 +10,10 @@ use futures::{AsyncRead, AsyncWrite}; use crate::{ context::Context, - executor::{Executor, ExecutorBuilder}, io::Io, mux::Mux, + session::{Session, SessionBuilder}, + thread_pool::ThreadPool, }; use super::recording::RecordedMtData; @@ -65,43 +66,6 @@ impl AsyncWrite for ReplayDuplex { } } -/// A simple mux that wraps a single replay stream. -struct SingleReplayMux { - replay: Mutex>, - max_frame_length: Option, -} - -impl SingleReplayMux { - fn new(recorded: Vec, max_frame_length: Option) -> Self { - Self { - replay: Mutex::new(Some(ReplayDuplex::new(recorded))), - max_frame_length, - } - } -} - -impl Mux for SingleReplayMux { - fn open(&self, id: &[u8]) -> Result { - // Only allow opening the root ID - if id != [0] { - return Err(std::io::Error::other( - "single replay mux only supports root ID", - )); - } - let replay = self - .replay - .lock() - .unwrap() - .take() - .ok_or_else(|| std::io::Error::other("channel already opened"))?; - if let Some(limit) = self.max_frame_length { - Ok(Io::from_io_with_limit(replay, limit)) - } else { - Ok(Io::from_io(replay)) - } - } -} - /// Creates a single-threaded context that replays recorded bytes. /// /// The context will read from the recorded bytes and discard all writes. @@ -112,8 +76,10 @@ impl Mux for SingleReplayMux { /// * `recorded` - The recorded bytes to replay. /// * `max_frame_length` - Maximum frame size in bytes. pub fn replay_st_context(recorded: Vec, max_frame_length: usize) -> Context { - let mux = SingleReplayMux::new(recorded, Some(max_frame_length)); - Context::new(mux).unwrap() + Context::from_io(Io::from_io_with_limit( + ReplayDuplex::new(recorded), + max_frame_length, + )) } // ============================================================================ @@ -165,9 +131,9 @@ impl Mux for ReplayTestMux { /// # Arguments /// /// * `recorded` - The recorded data to replay (per-channel). -pub fn replay_mt_context(recorded: RecordedMtData) -> Executor { +pub fn replay_mt_context(recorded: RecordedMtData) -> Session { let mux = ReplayTestMux::new(recorded, None); - ExecutorBuilder::default().build(mux) + SessionBuilder::default().build(mux).unwrap() } /// Creates a multi-threaded context that replays recorded data with a custom @@ -177,9 +143,9 @@ pub fn replay_mt_context(recorded: RecordedMtData) -> Executor { /// /// * `recorded` - The recorded data to replay (per-channel). /// * `max_frame_length` - Maximum frame size in bytes. -pub fn replay_mt_context_with_limit(recorded: RecordedMtData, max_frame_length: usize) -> Executor { +pub fn replay_mt_context_with_limit(recorded: RecordedMtData, max_frame_length: usize) -> Session { let mux = ReplayTestMux::new(recorded, Some(max_frame_length)); - ExecutorBuilder::default().build(mux) + SessionBuilder::default().build(mux).unwrap() } /// Like [`replay_mt_context_with_limit`], but uses a custom worker spawn @@ -189,13 +155,15 @@ pub fn replay_mt_context_with_spawn_and_limit( max_frame_length: usize, concurrency: usize, spawn: F, -) -> Executor +) -> Session where F: Fn(Box) -> Result<(), std::io::Error> + Send + Sync + 'static, { let mux = ReplayTestMux::new(recorded, Some(max_frame_length)); - ExecutorBuilder::default() + let pool = ThreadPool::builder() .num_threads(concurrency) .spawn(spawn) - .build(mux) + .build() + .unwrap(); + SessionBuilder::default().pool(pool).build(mux).unwrap() } diff --git a/crates/common/src/context/test/tests.rs b/crates/common/src/context/test/tests.rs index 33c05fd9..b77afbf8 100644 --- a/crates/common/src/context/test/tests.rs +++ b/crates/common/src/context/test/tests.rs @@ -445,3 +445,55 @@ async fn test_recording_mt_nested_try_join() { recorded_data.channels.len() ); } + +#[tokio::test] +async fn test_map_respects_concurrency_limit() { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + const LIMIT: usize = 3; + + let (mux, _peer) = crate::mux::test_framed_mux(1024); + let session = crate::session::SessionBuilder::default() + .cooperative() + .concurrency_limit(LIMIT) + .build(mux) + .unwrap(); + let mut ctx = session.new_context().unwrap(); + + // Tracks how many item futures are running at once and the peak observed. + let active = Arc::new(AtomicUsize::new(0)); + let max_seen = Arc::new(AtomicUsize::new(0)); + + let items: Vec = (0..12).collect(); + let results = ctx + .map(items, { + let active = active.clone(); + let max_seen = max_seen.clone(); + move |_ctx: &mut Context, item: u32| { + let active = active.clone(); + let max_seen = max_seen.clone(); + Box::pin(async move { + let now = active.fetch_add(1, Ordering::SeqCst) + 1; + max_seen.fetch_max(now, Ordering::SeqCst); + // Yield repeatedly so concurrently-buffered futures overlap + // before any of them completes. + for _ in 0..8 { + tokio::task::yield_now().await; + } + active.fetch_sub(1, Ordering::SeqCst); + item + }) + } + }) + .await + .unwrap(); + + // Results are returned in input order regardless of the bound. + assert_eq!(results, (0..12).collect::>()); + // Concurrency never exceeded the limit, and reached it (more items than the + // limit, so the window is fully utilized). + assert_eq!(max_seen.load(Ordering::SeqCst), LIMIT); +} diff --git a/crates/common/src/executor.rs b/crates/common/src/executor.rs deleted file mode 100644 index 890a52f0..00000000 --- a/crates/common/src/executor.rs +++ /dev/null @@ -1,482 +0,0 @@ -//! Work-stealing async executor. -//! -//! This module provides a work-stealing threadpool executor that integrates -//! with the MPC task model. Each task is assigned a deterministic [`ContextId`] -//! and owns its own I/O channel, allowing tasks to be freely migrated between -//! worker threads while maintaining deterministic execution order for I/O. - -use std::sync::{ - Arc, - atomic::{AtomicBool, AtomicU32, Ordering}, -}; - -use async_task::{Runnable, Task}; -use crossbeam_deque::{Injector, Steal, Stealer, Worker}; -use crossbeam_utils::sync::{Parker, Unparker}; - -use crate::{Context, ContextId, mux::Mux}; - -/// A work-stealing async executor. -#[derive(Debug)] -pub struct Executor { - inner: Arc, -} - -/// Per-worker parking state. -struct WorkerState { - unparker: Unparker, - parked: AtomicBool, -} - -pub(crate) struct Inner { - /// Global task queue for new tasks and cross-thread wakeups. - injector: Injector, - - /// Stealers for each worker's local queue. - stealers: Vec>, - - /// Per-worker parking state, indexed by worker index. - workers: Box<[WorkerState]>, - - /// Shutdown flag. - shutdown: AtomicBool, - - /// Multiplexer for creating I/O channels. - mux: Arc, - - /// Namespace prefix applied to all contexts created by this executor. - prefix: ContextId, - - /// Counter handed out to each new context, ensuring uniqueness. - next_context: AtomicU32, -} - -impl std::fmt::Debug for Inner { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Inner") - .field("workers", &self.workers.len()) - .field("shutdown", &self.shutdown) - .finish_non_exhaustive() - } -} - -/// A worker spawn callback. -/// -/// Receives a worker entry-point and dispatches it on a thread (or -/// platform-equivalent, e.g. `web_spawn::spawn` on wasm). -pub type SpawnFn = - Box) -> Result<(), std::io::Error> + Send + Sync>; - -/// Builder for [`Executor`]. -pub struct ExecutorBuilder { - num_threads: usize, - prefix: ContextId, - spawn: SpawnFn, -} - -impl std::fmt::Debug for ExecutorBuilder { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ExecutorBuilder") - .field("num_threads", &self.num_threads) - .field("prefix", &self.prefix) - .finish_non_exhaustive() - } -} - -impl Default for ExecutorBuilder { - fn default() -> Self { - Self { - num_threads: std::thread::available_parallelism() - .map(|n| n.get()) - .unwrap_or(4), - prefix: ContextId::from_prefix([]), - spawn: Box::new(default_spawn), - } - } -} - -fn default_spawn(f: Box) -> Result<(), std::io::Error> { - std::thread::Builder::new() - .name("mpz-executor-worker".to_string()) - .spawn(f) - .map(drop) -} - -impl ExecutorBuilder { - /// Sets the number of worker threads. - pub fn num_threads(mut self, n: usize) -> Self { - self.num_threads = n; - self - } - - /// Sets a namespace prefix applied to all contexts created by the - /// executor. - /// - /// Useful when several sub-protocols share a mux and need to be kept in - /// disjoint ID spaces. - pub fn prefix(mut self, prefix: impl AsRef<[u8]>) -> Self { - self.prefix = ContextId::from_prefix(prefix); - self - } - - /// Sets a custom worker spawn callback. - /// - /// Defaults to `std::thread::spawn`. Useful on platforms without OS - /// threads (e.g. wasm32 where workers must be created via - /// `web_spawn::spawn`). - pub fn spawn(mut self, spawn: F) -> Self - where - F: Fn(Box) -> Result<(), std::io::Error> - + Send - + Sync - + 'static, - { - self.spawn = Box::new(spawn); - self - } - - /// Builds the executor with the given multiplexer. - pub fn build(self, mux: M) -> Executor { - let injector = Injector::new(); - - // Create local worker queues and their stealers. - let worker_queues: Vec> = - (0..self.num_threads).map(|_| Worker::new_fifo()).collect(); - - let stealers: Vec> = worker_queues.iter().map(|w| w.stealer()).collect(); - - let parkers: Vec = (0..self.num_threads).map(|_| Parker::new()).collect(); - let workers: Box<[WorkerState]> = parkers - .iter() - .map(|p| WorkerState { - unparker: p.unparker().clone(), - parked: AtomicBool::new(false), - }) - .collect(); - - let inner = Arc::new(Inner { - injector, - stealers, - workers, - shutdown: AtomicBool::new(false), - mux: Arc::new(mux), - prefix: self.prefix, - next_context: AtomicU32::new(0), - }); - - // Spawn worker threads via the configured spawn callback. - for (index, (local, parker)) in worker_queues.into_iter().zip(parkers).enumerate() { - let inner = inner.clone(); - (self.spawn)(Box::new(move || worker_loop(inner, local, index, parker))) - .expect("failed to spawn worker thread"); - } - - Executor { inner } - } -} - -/// Worker thread loop. -fn worker_loop(inner: Arc, local: Worker, index: usize, parker: Parker) { - let state = &inner.workers[index]; - - let drain_local = |local: &Worker| { - // Drop any runnables still sitting in this worker's local queue. - // Dropping cancels the corresponding task so awaiters of `Task` - // see cancellation instead of hanging on a worker that has exited. - while local.pop().is_some() {} - }; - - while !inner.shutdown.load(Ordering::Relaxed) { - if let Some(runnable) = find_task(&inner, &local, index) { - // Poll the task once. If it returns Pending, the waker will - // reschedule it; if it completes, we're done with the task. - runnable.run(); - continue; - } - - // Slow path: announce we're about to park, then recheck. - // - // The recheck after setting `parked = true` closes the race against a - // producer that pushed before we announced (and therefore didn't see - // us as a candidate to unpark). - state.parked.store(true, Ordering::SeqCst); - - if let Some(runnable) = find_task(&inner, &local, index) { - state.parked.store(false, Ordering::SeqCst); - runnable.run(); - continue; - } - - if inner.shutdown.load(Ordering::Relaxed) { - state.parked.store(false, Ordering::SeqCst); - break; - } - - // If a producer fires between the recheck above and `park()`, the - // `unpark` token is remembered by the parker and `park()` returns - // immediately — no lost wakeup. - parker.park(); - state.parked.store(false, Ordering::SeqCst); - } - - drain_local(&local); -} - -/// Finds a task to execute using work-stealing. -fn find_task(inner: &Inner, local: &Worker, index: usize) -> Option { - // 1. Local queue (fast path, cache-friendly). - if let Some(runnable) = local.pop() { - return Some(runnable); - } - - // 2. Global injector queue. - loop { - match inner.injector.steal_batch_and_pop(local) { - Steal::Success(runnable) => return Some(runnable), - Steal::Empty => break, - Steal::Retry => continue, - } - } - - // 3. Steal from other workers. - let num_stealers = inner.stealers.len(); - for i in 1..num_stealers { - let victim = (index + i) % num_stealers; - loop { - match inner.stealers[victim].steal_batch_and_pop(local) { - Steal::Success(runnable) => return Some(runnable), - Steal::Empty => break, - Steal::Retry => continue, - } - } - } - - None -} - -impl Executor { - /// Creates a new builder. - pub fn builder() -> ExecutorBuilder { - ExecutorBuilder::default() - } - - /// Shuts down the executor. - /// - /// Sets the shutdown flag, drains the global queue, and unparks every - /// worker. After this returns, no further runnables will be accepted by - /// the scheduler (newly woken tasks are dropped on arrival), and any - /// runnables still queued at the moment of shutdown are dropped. Dropping - /// a [`Runnable`] cancels its task, so awaiters of `Task` propagate - /// cancellation rather than hanging on a worker that has exited. - pub fn shutdown(&self) { - self.inner.shutdown.store(true, Ordering::SeqCst); - - // Drain the injector before unparking workers. Any push that races - // with this drain is handled by the shutdown check in the schedule - // callback (see `spawn_on`), which drops the runnable. - loop { - match self.inner.injector.steal() { - Steal::Success(_) => continue, - Steal::Empty => break, - Steal::Retry => continue, - } - } - - for w in self.inner.workers.iter() { - w.unparker.unpark(); - } - } - - /// Returns `true` if the executor has been shut down. - pub fn is_shutdown(&self) -> bool { - self.inner.shutdown.load(Ordering::SeqCst) - } - - /// Creates a new context. - /// - /// Each context produced by an executor is given a distinct ID under the - /// executor's configured prefix. - pub fn new_context(&self) -> Result { - let index = self.inner.next_context.fetch_add(1, Ordering::Relaxed); - let id = self.inner.prefix.child(index); - let io = self.inner.mux.open(id.as_ref())?; - Ok(Context::with_executor( - id, - io, - self.inner.mux.clone(), - self.inner.clone(), - )) - } -} - -impl Drop for Executor { - fn drop(&mut self) { - self.shutdown(); - } -} - -/// Spawns a future on the given executor inner. -pub(crate) fn spawn_on(inner: &Arc, future: F) -> Task -where - F: std::future::Future + Send + 'static, - F::Output: Send + 'static, -{ - let inner = Arc::clone(inner); - let schedule = move |runnable: Runnable| { - // After shutdown, no worker will run this. Dropping the runnable - // cancels the task so the awaiter doesn't hang. SeqCst pairs with - // the SeqCst store in `Executor::shutdown` to ensure that any push - // that "loses" the race is then drained by shutdown's pass over - // the injector. - if inner.shutdown.load(Ordering::SeqCst) { - drop(runnable); - return; - } - inner.injector.push(runnable); - // Scan for an idle worker and claim it for this notification. The - // `load` is a cheap filter; the `compare_exchange` is what makes the - // claim race-free against other concurrent producers. Stops at the - // first claimed worker — one push, one wake. - for w in inner.workers.iter() { - if w.parked.load(Ordering::SeqCst) - && w.parked - .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst) - .is_ok() - { - w.unparker.unpark(); - break; - } - } - }; - let (runnable, task) = async_task::spawn(future, schedule); - runnable.schedule(); - task -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::mux::test_framed_mux; - use serio::{SinkExt, StreamExt}; - - #[test] - fn test_executor_spawn() { - let (mux_a, _mux_b) = test_framed_mux(1024); - let executor = Executor::builder().num_threads(2).build(mux_a); - - let mut ctx = executor.new_context().unwrap(); - let (a, b) = futures::executor::block_on(ctx.join( - |_ctx| Box::pin(async move { 21 }), - |_ctx| Box::pin(async move { 21 }), - )) - .unwrap(); - - assert_eq!(a + b, 42); - - executor.shutdown(); - } - - #[test] - fn test_executor_map() { - let (mux_a, _mux_b) = test_framed_mux(1024); - let executor = Executor::builder().num_threads(2).build(mux_a); - - let mut ctx = executor.new_context().unwrap(); - - let items = vec![1, 2, 3, 4, 5]; - let results = - futures::executor::block_on(ctx.map(items, |_ctx, x| Box::pin(async move { x * 2 }))); - - assert_eq!(results.unwrap(), vec![2, 4, 6, 8, 10]); - - executor.shutdown(); - } - - #[test] - fn test_executor_join() { - let (mux_a, _mux_b) = test_framed_mux(1024); - let executor = Executor::builder().num_threads(2).build(mux_a); - - let mut ctx = executor.new_context().unwrap(); - - let result = futures::executor::block_on(ctx.join( - |_ctx| Box::pin(async move { 1 + 1 }), - |_ctx| Box::pin(async move { 2 + 2 }), - )); - - assert_eq!(result.unwrap(), (2, 4)); - - executor.shutdown(); - } - - #[test] - fn test_executor_io() { - // Test that I/O works between two executors (simulating two parties). - let (mux_a, mux_b) = test_framed_mux(1024); - - let executor_a = Executor::builder().num_threads(2).build(mux_a); - let executor_b = Executor::builder().num_threads(2).build(mux_b); - - let mut ctx_a = executor_a.new_context().unwrap(); - let mut ctx_b = executor_b.new_context().unwrap(); - - let (_, (val1, val2)) = futures::executor::block_on(futures::future::join( - async { - ctx_a.io_mut().send(42u32).await.unwrap(); - ctx_a.io_mut().send(123u32).await.unwrap(); - }, - async { - let val1: u32 = ctx_b.io_mut().next().await.unwrap().unwrap(); - let val2: u32 = ctx_b.io_mut().next().await.unwrap().unwrap(); - (val1, val2) - }, - )); - - assert_eq!(val1, 42); - assert_eq!(val2, 123); - - executor_a.shutdown(); - executor_b.shutdown(); - } - - #[test] - fn test_executor_map_with_io() { - // Test that map works with I/O between two parties. - let (mux_a, mux_b) = test_framed_mux(1024); - - let executor_a = Executor::builder().num_threads(4).build(mux_a); - let executor_b = Executor::builder().num_threads(4).build(mux_b); - - let mut ctx_a = executor_a.new_context().unwrap(); - let mut ctx_b = executor_b.new_context().unwrap(); - - let items_a = vec![1u32, 2, 3, 4]; - let items_b = vec![10u32, 20, 30, 40]; - - // Party A sends each item, Party B receives and returns sum. - let task_a = ctx_a.map(items_a, |ctx, x| { - Box::pin(async move { - ctx.io_mut().send(x).await.unwrap(); - }) - }); - - let task_b = ctx_b.map(items_b, |ctx, x| { - Box::pin(async move { - let received: u32 = ctx.io_mut().next().await.unwrap().unwrap(); - received + x - }) - }); - - let (results_a, results_b) = - futures::executor::block_on(futures::future::join(task_a, task_b)); - - assert!(results_a.is_ok()); - let results_b = results_b.unwrap(); - - // Each B task should receive the corresponding A value and add it to B's value. - assert_eq!(results_b, vec![11, 22, 33, 44]); - - executor_a.shutdown(); - executor_b.shutdown(); - } -} diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index f00a8a89..0e828aa9 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -16,7 +16,6 @@ )] pub mod context; -pub mod executor; #[cfg(any(test, feature = "future"))] pub mod future; mod id; @@ -24,14 +23,17 @@ mod id; pub mod ideal; pub mod io; pub mod mux; +pub mod session; #[cfg(feature = "sync")] pub mod sync; mod task; +pub mod thread_pool; pub use context::{Context, ContextError}; -pub use executor::{Executor, ExecutorBuilder}; pub use id::ContextId; +pub use session::{Session, SessionBuilder}; pub use task::Task; +pub use thread_pool::{ThreadPool, ThreadPoolBuildError, ThreadPoolBuilder}; use async_trait::async_trait; diff --git a/crates/common/src/session.rs b/crates/common/src/session.rs new file mode 100644 index 00000000..e0064a5b --- /dev/null +++ b/crates/common/src/session.rs @@ -0,0 +1,351 @@ +//! Async session. +//! +//! A [`Session`] hands out [`Context`]s, each with a distinct [`ContextId`] +//! and its own I/O channel from the configured multiplexer. Sub-tasks +//! spawned through a `Context` run on the session's [`ThreadPool`] — the +//! global pool by default, a builder-supplied pool, or, when +//! [`SessionBuilder::cooperative`] is set, the caller's future. + +use std::sync::{ + Arc, + atomic::{AtomicU32, Ordering}, +}; + +use crate::{ + Context, ContextId, + context::DEFAULT_CONCURRENCY_LIMIT, + mux::Mux, + thread_pool::{ThreadPool, ThreadPoolBuildError}, +}; + +/// An async session. +pub struct Session { + /// Pool sub-tasks run on; `None` runs them cooperatively on the caller. + pool: Option, + mux: Arc, + prefix: ContextId, + next_context: AtomicU32, + /// Maximum number of [`Context::map`] items processed concurrently. + concurrency_limit: usize, +} + +impl std::fmt::Debug for Session { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Session") + .field("pool", &self.pool) + .field("prefix", &self.prefix) + .finish_non_exhaustive() + } +} + +#[derive(Default)] +enum PoolMode { + /// Use a specific pool. + Pool(ThreadPool), + /// Use the global pool (resolved at `build` time). + #[default] + Global, + /// Run sub-tasks cooperatively on the caller's future. + Cooperative, +} + +impl std::fmt::Debug for PoolMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Pool(pool) => f.debug_tuple("Pool").field(pool).finish(), + Self::Global => f.write_str("Global"), + Self::Cooperative => f.write_str("Cooperative"), + } + } +} + +/// Builder for [`Session`]. +#[derive(Debug)] +pub struct SessionBuilder { + pool: PoolMode, + prefix: ContextId, + concurrency_limit: usize, +} + +impl Default for SessionBuilder { + fn default() -> Self { + Self { + pool: PoolMode::default(), + prefix: ContextId::default(), + concurrency_limit: DEFAULT_CONCURRENCY_LIMIT, + } + } +} + +impl SessionBuilder { + /// Sets the pool the session will run tasks on. + /// + /// Overrides the default of using the global pool + /// ([`ThreadPool::global`]). + pub fn pool(mut self, pool: ThreadPool) -> Self { + self.pool = PoolMode::Pool(pool); + self + } + + /// Configures the session to run sub-tasks cooperatively on the caller's + /// future rather than on a thread pool. + /// + /// Each sub-task is still given its own I/O channel from the mux, but + /// they are polled by the caller — no threads are spawned. + pub fn cooperative(mut self) -> Self { + self.pool = PoolMode::Cooperative; + self + } + + /// Sets a namespace prefix applied to all contexts created by the + /// session. + /// + /// Useful when several sub-protocols share a mux and need to be kept in + /// disjoint ID spaces. + pub fn prefix(mut self, prefix: impl AsRef<[u8]>) -> Self { + self.prefix = ContextId::from_prefix(prefix); + self + } + + /// Sets the maximum number of [`Context::map`] items processed + /// concurrently, bounding how many sub-channels are open at once. + /// + /// Both parties **must** configure the same limit so they open the same + /// sliding window of channels and stay in lockstep. The limit is clamped + /// to a minimum of `1`. Defaults to [`DEFAULT_CONCURRENCY_LIMIT`]. + pub fn concurrency_limit(mut self, limit: usize) -> Self { + self.concurrency_limit = limit.max(1); + self + } + + /// Builds the session with the given multiplexer. + /// + /// Returns an error if the builder is configured to use the global pool + /// (the default) but the global pool cannot be built — for instance on + /// platforms without OS threads. Callers that want to run without a + /// thread pool in those cases should opt in via + /// [`cooperative`](Self::cooperative). + pub fn build( + self, + mux: M, + ) -> Result { + let pool = match self.pool { + PoolMode::Pool(pool) => Some(pool), + PoolMode::Global => Some(ThreadPool::try_global()?), + PoolMode::Cooperative => None, + }; + Ok(Session { + pool, + mux: Arc::new(mux), + prefix: self.prefix, + next_context: AtomicU32::new(0), + concurrency_limit: self.concurrency_limit, + }) + } +} + +impl Session { + /// Creates a new builder. + pub fn builder() -> SessionBuilder { + SessionBuilder::default() + } + + /// Returns the pool this session runs tasks on, or `None` if it runs + /// sub-tasks cooperatively. + pub fn pool(&self) -> Option<&ThreadPool> { + self.pool.as_ref() + } + + /// Creates a new context. + /// + /// Each context produced by a session is given a distinct ID under the + /// session's configured prefix. + pub fn new_context(&self) -> Result { + let index = self.next_context.fetch_add(1, Ordering::Relaxed); + let id = self.prefix.child(index); + let io = self.mux.open(id.as_ref())?; + Ok(Context::for_session( + id, + io, + self.mux.clone(), + self.pool.clone(), + self.concurrency_limit, + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mux::test_framed_mux; + use serio::{SinkExt, StreamExt}; + + fn test_pool() -> ThreadPool { + ThreadPool::builder().num_threads(2).build().unwrap() + } + + #[test] + fn test_session_spawn() { + let (mux_a, _mux_b) = test_framed_mux(1024); + let session = Session::builder().pool(test_pool()).build(mux_a).unwrap(); + + let mut ctx = session.new_context().unwrap(); + let (a, b) = futures::executor::block_on(ctx.join( + |_ctx| Box::pin(async move { 21 }), + |_ctx| Box::pin(async move { 21 }), + )) + .unwrap(); + + assert_eq!(a + b, 42); + } + + #[test] + fn test_session_map() { + let (mux_a, _mux_b) = test_framed_mux(1024); + let session = Session::builder().pool(test_pool()).build(mux_a).unwrap(); + + let mut ctx = session.new_context().unwrap(); + + let items = vec![1, 2, 3, 4, 5]; + let results = + futures::executor::block_on(ctx.map(items, |_ctx, x| Box::pin(async move { x * 2 }))); + + assert_eq!(results.unwrap(), vec![2, 4, 6, 8, 10]); + } + + #[test] + fn test_session_join() { + let (mux_a, _mux_b) = test_framed_mux(1024); + let session = Session::builder().pool(test_pool()).build(mux_a).unwrap(); + + let mut ctx = session.new_context().unwrap(); + + let result = futures::executor::block_on(ctx.join( + |_ctx| Box::pin(async move { 1 + 1 }), + |_ctx| Box::pin(async move { 2 + 2 }), + )); + + assert_eq!(result.unwrap(), (2, 4)); + } + + #[test] + fn test_session_io() { + let (mux_a, mux_b) = test_framed_mux(1024); + + let session_a = Session::builder().pool(test_pool()).build(mux_a).unwrap(); + let session_b = Session::builder().pool(test_pool()).build(mux_b).unwrap(); + + let mut ctx_a = session_a.new_context().unwrap(); + let mut ctx_b = session_b.new_context().unwrap(); + + let (_, (val1, val2)) = futures::executor::block_on(futures::future::join( + async { + ctx_a.io_mut().send(42u32).await.unwrap(); + ctx_a.io_mut().send(123u32).await.unwrap(); + }, + async { + let val1: u32 = ctx_b.io_mut().next().await.unwrap().unwrap(); + let val2: u32 = ctx_b.io_mut().next().await.unwrap().unwrap(); + (val1, val2) + }, + )); + + assert_eq!(val1, 42); + assert_eq!(val2, 123); + } + + #[test] + fn test_session_map_with_io() { + let (mux_a, mux_b) = test_framed_mux(1024); + + let pool = ThreadPool::builder().num_threads(4).build().unwrap(); + let session_a = Session::builder().pool(pool.clone()).build(mux_a).unwrap(); + let session_b = Session::builder().pool(pool).build(mux_b).unwrap(); + + let mut ctx_a = session_a.new_context().unwrap(); + let mut ctx_b = session_b.new_context().unwrap(); + + let items_a = vec![1u32, 2, 3, 4]; + let items_b = vec![10u32, 20, 30, 40]; + + let task_a = ctx_a.map(items_a, |ctx, x| { + Box::pin(async move { + ctx.io_mut().send(x).await.unwrap(); + }) + }); + + let task_b = ctx_b.map(items_b, |ctx, x| { + Box::pin(async move { + let received: u32 = ctx.io_mut().next().await.unwrap().unwrap(); + received + x + }) + }); + + let (results_a, results_b) = + futures::executor::block_on(futures::future::join(task_a, task_b)); + + assert!(results_a.is_ok()); + let results_b = results_b.unwrap(); + + assert_eq!(results_b, vec![11, 22, 33, 44]); + } + + #[test] + fn test_global_pool_shared() { + let (mux_a, mux_b) = test_framed_mux(1024); + let session_a = Session::builder().prefix(b"a").build(mux_a).unwrap(); + let session_b = Session::builder().prefix(b"b").build(mux_b).unwrap(); + + assert!(!session_a.pool().unwrap().is_shutdown()); + assert!(!session_b.pool().unwrap().is_shutdown()); + + let mut ctx_a = session_a.new_context().unwrap(); + let mut ctx_b = session_b.new_context().unwrap(); + + let (sum_a, sum_b) = futures::executor::block_on(futures::future::join( + async { + let (x, y) = ctx_a + .join( + |_ctx| Box::pin(async move { 10 }), + |_ctx| Box::pin(async move { 20 }), + ) + .await + .unwrap(); + x + y + }, + async { + let (x, y) = ctx_b + .join( + |_ctx| Box::pin(async move { 1 }), + |_ctx| Box::pin(async move { 2 }), + ) + .await + .unwrap(); + x + y + }, + )); + + assert_eq!(sum_a, 30); + assert_eq!(sum_b, 3); + + drop(session_a); + assert!(!session_b.pool().unwrap().is_shutdown()); + } + + #[test] + fn test_cooperative_session() { + let (mux_a, _mux_b) = test_framed_mux(1024); + let session = Session::builder().cooperative().build(mux_a).unwrap(); + + assert!(session.pool().is_none()); + + let mut ctx = session.new_context().unwrap(); + let (a, b) = futures::executor::block_on(ctx.join( + |_ctx| Box::pin(async move { 21 }), + |_ctx| Box::pin(async move { 21 }), + )) + .unwrap(); + + assert_eq!(a + b, 42); + } +} diff --git a/crates/common/src/thread_pool.rs b/crates/common/src/thread_pool.rs new file mode 100644 index 00000000..7be47783 --- /dev/null +++ b/crates/common/src/thread_pool.rs @@ -0,0 +1,415 @@ +//! Async thread pool. +//! +//! A [`ThreadPool`] is a shared, [`Clone`]able handle to a pool of worker +//! threads that run async tasks. Build a dedicated pool with +//! [`ThreadPool::builder`], or use the process-wide [`ThreadPool::global`] +//! pool. + +use std::sync::{ + Arc, OnceLock, + atomic::{AtomicBool, Ordering}, +}; + +use async_task::{Runnable, Task}; +use crossbeam_deque::{Injector, Steal, Stealer, Worker}; +use crossbeam_utils::sync::{Parker, Unparker}; + +/// A shared handle to a thread pool. +/// +/// `ThreadPool` is cheap to clone — clones share the same underlying pool. +/// When the last `ThreadPool` handle is dropped, the pool is shut down. +#[derive(Clone)] +pub struct ThreadPool { + inner: Arc, + // Holding a guard alongside the inner Arc lets us distinguish "last user + // handle dropped" from "last worker reference dropped". When the final + // guard drops, the pool is shut down; workers then exit and release their + // own `inner` refs. + _guard: Arc, +} + +impl std::fmt::Debug for ThreadPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ThreadPool") + .field("num_threads", &self.inner.workers.len()) + .field("is_shutdown", &self.is_shutdown()) + .finish_non_exhaustive() + } +} + +struct Guard { + inner: Arc, +} + +impl Drop for Guard { + fn drop(&mut self) { + self.inner.shutdown(); + } +} + +struct WorkerState { + unparker: Unparker, + parked: AtomicBool, +} + +struct Inner { + injector: Injector, + stealers: Vec>, + workers: Box<[WorkerState]>, + shutdown: AtomicBool, +} + +impl Inner { + fn shutdown(&self) { + self.shutdown.store(true, Ordering::SeqCst); + + // Drain the injector before unparking workers. Any push that races + // with this drain is handled by the shutdown check in the schedule + // callback (see `spawn_on`), which drops the runnable. + loop { + match self.injector.steal() { + Steal::Success(_) => continue, + Steal::Empty => break, + Steal::Retry => continue, + } + } + + for w in self.workers.iter() { + w.unparker.unpark(); + } + } +} + +/// A worker spawn callback. +/// +/// Receives a worker entry-point and dispatches it on a thread (or +/// platform-equivalent, e.g. `web_spawn::spawn` on wasm). +pub type SpawnFn = + Box) -> Result<(), std::io::Error> + Send + Sync>; + +/// Builder for [`ThreadPool`]. +pub struct ThreadPoolBuilder { + num_threads: usize, + spawn: SpawnFn, +} + +impl std::fmt::Debug for ThreadPoolBuilder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ThreadPoolBuilder") + .field("num_threads", &self.num_threads) + .finish_non_exhaustive() + } +} + +impl Default for ThreadPoolBuilder { + fn default() -> Self { + Self { + num_threads: default_num_threads(), + spawn: Box::new(default_spawn), + } + } +} + +fn default_num_threads() -> usize { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(8) +} + +fn default_spawn(f: Box) -> Result<(), std::io::Error> { + std::thread::Builder::new() + .name("mpz-pool-worker".to_string()) + .spawn(f) + .map(drop) +} + +impl ThreadPoolBuilder { + /// Sets the number of worker threads. + pub fn num_threads(mut self, n: usize) -> Self { + self.num_threads = n; + self + } + + /// Sets a custom worker spawn callback. + pub fn spawn(mut self, spawn: F) -> Self + where + F: Fn(Box) -> Result<(), std::io::Error> + + Send + + Sync + + 'static, + { + self.spawn = Box::new(spawn); + self + } + + /// Builds the pool. + /// + /// Returns an error if the [spawn callback](Self::spawn) fails to start a + /// worker. + pub fn build(self) -> Result { + ThreadPool::new(self.num_threads, &self.spawn) + } + + /// Builds the pool and installs it as the process-wide global pool, the + /// one returned by [`ThreadPool::global`]. + /// + /// Returns an error if the global pool has already been initialized. + pub fn build_global(self) -> Result<(), ThreadPoolBuildError> { + if GLOBAL_POOL.get().is_some() { + return Err(ThreadPoolBuildError::AlreadyInitialized); + } + let pool = self.build()?; + GLOBAL_POOL + .set(pool) + .map_err(|_| ThreadPoolBuildError::AlreadyInitialized) + } +} + +/// Errors that can occur while building a [`ThreadPool`]. +#[derive(Debug, thiserror::Error)] +pub enum ThreadPoolBuildError { + /// The process-wide global pool has already been initialized. + #[error("global thread pool is already initialized")] + AlreadyInitialized, + /// A worker thread could not be started. + #[error("failed to start a worker thread")] + Spawn(#[source] std::io::Error), +} + +/// Process-wide global pool, initialized at most once. +static GLOBAL_POOL: OnceLock = OnceLock::new(); + +impl ThreadPool { + /// Creates a new builder. + pub fn builder() -> ThreadPoolBuilder { + ThreadPoolBuilder::default() + } + + /// Returns a handle to the process-wide global pool. + /// + /// The pool is initialized with default settings on the first call. The + /// global pool is never shut down implicitly. + /// + /// Panics if the default pool fails to build. Use [`try_global`] for a + /// fallible variant — useful on platforms where the default spawn + /// callback cannot start threads (e.g. wasm32). + /// + /// [`try_global`]: Self::try_global + pub fn global() -> ThreadPool { + Self::try_global().expect("default global thread pool should build") + } + + /// Like [`global`](Self::global), but returns an error instead of + /// panicking when the default pool cannot be built. + pub fn try_global() -> Result { + if let Some(pool) = GLOBAL_POOL.get() { + return Ok(pool.clone()); + } + let pool = ThreadPool::builder().build()?; + // `set` returns `Err` if another thread won the init race; either way, + // `GLOBAL_POOL.get()` is `Some` once we get here. + let _ = GLOBAL_POOL.set(pool); + Ok(GLOBAL_POOL + .get() + .expect("global pool is initialized") + .clone()) + } + + /// Shuts down the pool. + /// + /// After this returns, no further tasks will be accepted, and any tasks + /// still pending are cancelled — awaiters propagate cancellation rather + /// than hanging. + pub fn shutdown(&self) { + self.inner.shutdown(); + } + + /// Returns `true` if the pool has been shut down. + pub fn is_shutdown(&self) -> bool { + self.inner.shutdown.load(Ordering::SeqCst) + } + + fn new(num_threads: usize, spawn: &SpawnFn) -> Result { + let injector = Injector::new(); + + let worker_queues: Vec> = + (0..num_threads).map(|_| Worker::new_fifo()).collect(); + + let stealers: Vec> = worker_queues.iter().map(|w| w.stealer()).collect(); + + let parkers: Vec = (0..num_threads).map(|_| Parker::new()).collect(); + let workers: Box<[WorkerState]> = parkers + .iter() + .map(|p| WorkerState { + unparker: p.unparker().clone(), + parked: AtomicBool::new(false), + }) + .collect(); + + let inner = Arc::new(Inner { + injector, + stealers, + workers, + shutdown: AtomicBool::new(false), + }); + + for (index, (local, parker)) in worker_queues.into_iter().zip(parkers).enumerate() { + let worker_inner = inner.clone(); + if let Err(err) = spawn(Box::new(move || { + worker_loop(worker_inner, local, index, parker) + })) { + // A spawn failed partway through. Signal shutdown so any + // workers we already started exit, then surface the error. + inner.shutdown(); + return Err(ThreadPoolBuildError::Spawn(err)); + } + } + + Ok(ThreadPool { + _guard: Arc::new(Guard { + inner: inner.clone(), + }), + inner, + }) + } +} + +fn worker_loop(inner: Arc, local: Worker, index: usize, parker: Parker) { + let state = &inner.workers[index]; + + let drain_local = |local: &Worker| { + // Drop any runnables still sitting in this worker's local queue. + // Dropping cancels the corresponding task so awaiters of `Task` + // see cancellation instead of hanging on a worker that has exited. + while local.pop().is_some() {} + }; + + while !inner.shutdown.load(Ordering::Relaxed) { + if let Some(runnable) = find_task(&inner, &local, index) { + runnable.run(); + continue; + } + + // Slow path: announce we're about to park, then recheck. + // + // The recheck after setting `parked = true` closes the race against a + // producer that pushed before we announced (and therefore didn't see + // us as a candidate to unpark). + state.parked.store(true, Ordering::SeqCst); + + if let Some(runnable) = find_task(&inner, &local, index) { + state.parked.store(false, Ordering::SeqCst); + runnable.run(); + continue; + } + + if inner.shutdown.load(Ordering::Relaxed) { + state.parked.store(false, Ordering::SeqCst); + break; + } + + // If a producer fires between the recheck above and `park()`, the + // `unpark` token is remembered by the parker and `park()` returns + // immediately — no lost wakeup. + parker.park(); + state.parked.store(false, Ordering::SeqCst); + } + + drain_local(&local); +} + +fn find_task(inner: &Inner, local: &Worker, index: usize) -> Option { + if let Some(runnable) = local.pop() { + return Some(runnable); + } + + loop { + match inner.injector.steal_batch_and_pop(local) { + Steal::Success(runnable) => return Some(runnable), + Steal::Empty => break, + Steal::Retry => continue, + } + } + + let num_stealers = inner.stealers.len(); + for i in 1..num_stealers { + let victim = (index + i) % num_stealers; + loop { + match inner.stealers[victim].steal_batch_and_pop(local) { + Steal::Success(runnable) => return Some(runnable), + Steal::Empty => break, + Steal::Retry => continue, + } + } + } + + None +} + +/// Spawns a future on the given pool. +pub(crate) fn spawn_on(pool: &ThreadPool, future: F) -> Task +where + F: std::future::Future + Send + 'static, + F::Output: Send + 'static, +{ + let inner = pool.inner.clone(); + let schedule = move |runnable: Runnable| { + // After shutdown, no worker will run this. Dropping the runnable + // cancels the task so the awaiter doesn't hang. SeqCst pairs with + // the SeqCst store in `Inner::shutdown` to ensure that any push + // that "loses" the race is then drained by shutdown's pass over + // the injector. + if inner.shutdown.load(Ordering::SeqCst) { + drop(runnable); + return; + } + inner.injector.push(runnable); + // Scan for an idle worker and claim it for this notification. The + // `load` is a cheap filter; the `compare_exchange` is what makes the + // claim race-free against other concurrent producers. Stops at the + // first claimed worker — one push, one wake. + for w in inner.workers.iter() { + if w.parked.load(Ordering::SeqCst) + && w.parked + .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + { + w.unparker.unpark(); + break; + } + } + }; + let (runnable, task) = async_task::spawn(future, schedule); + runnable.schedule(); + task +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_pool_drop_shuts_down() { + let pool = ThreadPool::builder().num_threads(2).build().unwrap(); + let weak = Arc::downgrade(&pool.inner); + drop(pool); + // After the last handle is dropped, workers observe shutdown and + // release their refs. Give them a moment to wind down. + for _ in 0..100 { + if weak.strong_count() == 0 { + break; + } + std::thread::sleep(std::time::Duration::from_millis(10)); + } + assert_eq!(weak.strong_count(), 0); + } + + #[test] + fn test_build_global_rejects_second_call() { + // The global pool may have been initialized by another test in this + // binary. Either way, a subsequent `build_global` must fail. + let _ = ThreadPool::global(); + let err = ThreadPool::builder().num_threads(1).build_global(); + assert!(err.is_err()); + } +} diff --git a/crates/garble/benches/evaluator.rs b/crates/garble/benches/evaluator.rs index 9e8f6ba5..ad850e9c 100644 --- a/crates/garble/benches/evaluator.rs +++ b/crates/garble/benches/evaluator.rs @@ -11,7 +11,7 @@ use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_m use futures::executor::block_on; use mpz_circuits::{AES128, Circuit}; use mpz_common::{ - Executor, + Session, context::{ RecordedMtData, recording_mt_context_with_limit, recording_st_context_with_limit, replay_mt_context_with_limit, replay_st_context, @@ -168,8 +168,8 @@ async fn run_evaluator_with_replay( /// Runs the full garble protocol with MT contexts. /// Records garbler->evaluator messages. async fn run_protocol_record_garbler_mt( - exec_gb: &mut Executor, - exec_ev: &mut Executor, + exec_gb: &mut Session, + exec_ev: &mut Session, circuit: Arc, circuit_count: usize, seed: u64, @@ -264,7 +264,7 @@ fn record_for_evaluator_mt( /// Runs MT evaluator only with replay context. async fn run_evaluator_with_replay_mt( - exec: &mut Executor, + exec: &mut Session, circuit: Arc, circuit_count: usize, ) { diff --git a/crates/garble/benches/garbler.rs b/crates/garble/benches/garbler.rs index 0a2cfb69..0b8f322c 100644 --- a/crates/garble/benches/garbler.rs +++ b/crates/garble/benches/garbler.rs @@ -10,7 +10,7 @@ use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_m use futures::executor::block_on; use mpz_circuits::{AES128, Circuit}; use mpz_common::{ - Executor, + Session, context::{ RecordedMtData, recording_mt_context_with_limit, recording_st_context_with_limit, replay_mt_context_with_limit, replay_st_context, @@ -170,8 +170,8 @@ async fn run_garbler_with_replay( /// Runs the full garble protocol with MT contexts. /// Records evaluator->garbler messages. async fn run_protocol_record_evaluator_mt( - exec_gb: &mut Executor, - exec_ev: &mut Executor, + exec_gb: &mut Session, + exec_ev: &mut Session, circuit: Arc, circuit_count: usize, seed: u64, @@ -268,7 +268,7 @@ fn record_for_garbler_mt( /// Runs MT garbler only with replay context. async fn run_garbler_with_replay_mt( - exec: &mut Executor, + exec: &mut Session, circuit: Arc, circuit_count: usize, delta: Delta, diff --git a/crates/ot/benches/ferret_receiver.rs b/crates/ot/benches/ferret_receiver.rs index 43ad75c6..b3688bcd 100644 --- a/crates/ot/benches/ferret_receiver.rs +++ b/crates/ot/benches/ferret_receiver.rs @@ -8,7 +8,7 @@ use criterion::{Criterion, Throughput, criterion_group, criterion_main}; use futures::executor::block_on; use mpz_common::{ - Executor, Flush, + Flush, Session, context::{ RecordedMtData, recording_mt_context_with_limit, recording_st_context_with_limit, replay_mt_context_with_limit, replay_st_context, @@ -141,8 +141,8 @@ struct RecordedDataMt { /// Runs the full Ferret protocol with MT contexts. /// Records sender->receiver messages. async fn run_protocol_record_sender_mt( - exec_sender: &mut Executor, - exec_receiver: &mut Executor, + exec_sender: &mut Session, + exec_receiver: &mut Session, config: FerretConfig, delta: Block, cot_seed: Block, @@ -213,7 +213,7 @@ fn record_for_receiver_mt(seed: u64) -> RecordedDataMt { } /// Runs MT receiver only with replay context. -async fn run_receiver_with_replay_mt(exec: &mut Executor, data: &RecordedDataMt) { +async fn run_receiver_with_replay_mt(exec: &mut Session, data: &RecordedDataMt) { let cot_recv = IdealRCOTReceiver::from_seed(data.cot_recv_seed); let config = bench_config(); let mut receiver = Receiver::new(config, data.receiver_seed, cot_recv); diff --git a/crates/ot/benches/ferret_sender.rs b/crates/ot/benches/ferret_sender.rs index 85c14f62..3ae88241 100644 --- a/crates/ot/benches/ferret_sender.rs +++ b/crates/ot/benches/ferret_sender.rs @@ -8,7 +8,7 @@ use criterion::{Criterion, Throughput, criterion_group, criterion_main}; use futures::executor::block_on; use mpz_common::{ - Executor, Flush, + Flush, Session, context::{ RecordedMtData, recording_mt_context_with_limit, recording_st_context_with_limit, replay_mt_context_with_limit, replay_st_context, @@ -146,8 +146,8 @@ struct RecordedDataMt { /// Runs the full Ferret protocol with MT contexts. /// Records receiver->sender messages. async fn run_protocol_record_receiver_mt( - exec_sender: &mut Executor, - exec_receiver: &mut Executor, + exec_sender: &mut Session, + exec_receiver: &mut Session, config: FerretConfig, delta: Block, cot_seed: Block, @@ -219,7 +219,7 @@ fn record_for_sender_mt(seed: u64) -> RecordedDataMt { } /// Runs MT sender only with replay context. -async fn run_sender_with_replay_mt(exec: &mut Executor, data: &RecordedDataMt) { +async fn run_sender_with_replay_mt(exec: &mut Session, data: &RecordedDataMt) { let cot_send = IdealRCOTSender::new(data.cot_seed, data.delta); let config = bench_config(); let mut sender = Sender::new(config, data.sender_seed, cot_send); diff --git a/crates/wasm-bench/src/garble/evaluator.rs b/crates/wasm-bench/src/garble/evaluator.rs index f14ddaa2..d7664380 100644 --- a/crates/wasm-bench/src/garble/evaluator.rs +++ b/crates/wasm-bench/src/garble/evaluator.rs @@ -9,7 +9,7 @@ use wasm_bindgen::prelude::*; #[cfg(target_arch = "wasm32")] use mpz_circuits::AES128; #[cfg(target_arch = "wasm32")] -use mpz_common::Executor; +use mpz_common::Session; use mpz_common::context::{ RecordedMtData, recording_mt_context_with_spawn_and_limit, replay_mt_context_with_spawn_and_limit, @@ -45,8 +45,8 @@ fn max_frame_length(circuit: &mpz_circuits::Circuit, circuit_count: usize) -> us #[cfg(target_arch = "wasm32")] async fn run_protocol_record_garbler( - exec_gb: &mut Executor, - exec_ev: &mut Executor, + exec_gb: &mut Session, + exec_ev: &mut Session, circuit_count: usize, seed: u64, ) { @@ -142,7 +142,7 @@ async fn record_for_evaluator( } #[cfg(target_arch = "wasm32")] -async fn run_evaluator_with_replay(exec: &mut Executor, circuit_count: usize) { +async fn run_evaluator_with_replay(exec: &mut Session, circuit_count: usize) { let (_, cot_recv) = ideal_cot([0u8; 16].into()); let mut ev = Evaluator::new(cot_recv); diff --git a/crates/wasm-bench/src/garble/garbler.rs b/crates/wasm-bench/src/garble/garbler.rs index ba7f3032..e0e5147c 100644 --- a/crates/wasm-bench/src/garble/garbler.rs +++ b/crates/wasm-bench/src/garble/garbler.rs @@ -9,7 +9,7 @@ use wasm_bindgen::prelude::*; #[cfg(target_arch = "wasm32")] use mpz_circuits::AES128; #[cfg(target_arch = "wasm32")] -use mpz_common::Executor; +use mpz_common::Session; use mpz_common::context::{ RecordedMtData, recording_mt_context_with_spawn_and_limit, replay_mt_context_with_spawn_and_limit, @@ -45,8 +45,8 @@ fn max_frame_length(circuit: &mpz_circuits::Circuit, circuit_count: usize) -> us #[cfg(target_arch = "wasm32")] async fn run_protocol_record_evaluator( - exec_gb: &mut Executor, - exec_ev: &mut Executor, + exec_gb: &mut Session, + exec_ev: &mut Session, circuit_count: usize, seed: u64, ) { @@ -137,7 +137,7 @@ async fn record_for_garbler(circuit_count: usize, seed: u64, concurrency: usize) } #[cfg(target_arch = "wasm32")] -async fn run_garbler_with_replay(exec: &mut Executor, circuit_count: usize, delta: Delta) { +async fn run_garbler_with_replay(exec: &mut Session, circuit_count: usize, delta: Delta) { let (cot_send, _) = ideal_cot(delta.into_inner()); let mut gb = Garbler::new(cot_send, [0u8; 16], delta); diff --git a/crates/wasm-bench/src/ot/ferret.rs b/crates/wasm-bench/src/ot/ferret.rs index 1abba9fc..cfebb4ad 100644 --- a/crates/wasm-bench/src/ot/ferret.rs +++ b/crates/wasm-bench/src/ot/ferret.rs @@ -49,7 +49,7 @@ fn bench_config() -> FerretConfig { // ============================================================================ #[cfg(target_arch = "wasm32")] -use mpz_common::Executor; +use mpz_common::Session; use mpz_common::context::{ RecordedMtData, recording_mt_context_with_spawn_and_limit, replay_mt_context_with_spawn_and_limit, @@ -73,8 +73,8 @@ struct RecordedDataMt { #[cfg(target_arch = "wasm32")] #[allow(clippy::too_many_arguments)] async fn run_protocol_record_receiver_mt( - exec_sender: &mut Executor, - exec_receiver: &mut Executor, + exec_sender: &mut Session, + exec_receiver: &mut Session, config: FerretConfig, delta: Block, cot_seed: Block, @@ -150,7 +150,7 @@ async fn record_for_sender_mt(seed: u64, concurrency: usize, ot_count: usize) -> /// Runs MT sender only with replay context. #[cfg(target_arch = "wasm32")] -async fn run_sender_with_replay_mt(exec: &mut Executor, data: &RecordedDataMt, ot_count: usize) { +async fn run_sender_with_replay_mt(exec: &mut Session, data: &RecordedDataMt, ot_count: usize) { let (cot_send, _) = ideal_rcot(data.cot_seed, data.delta); let config = bench_config(); let mut sender = Sender::new(config, data.sender_seed, cot_send); diff --git a/crates/wasm-bench/src/zk/prover.rs b/crates/wasm-bench/src/zk/prover.rs index c51546f4..990b5823 100644 --- a/crates/wasm-bench/src/zk/prover.rs +++ b/crates/wasm-bench/src/zk/prover.rs @@ -7,7 +7,7 @@ use wasm_bindgen::prelude::*; use mpz_circuits::AES128; #[cfg(target_arch = "wasm32")] -use mpz_common::Executor; +use mpz_common::Session; use mpz_common::context::{ RecordedMtData, recording_mt_context_with_spawn_and_limit, replay_mt_context_with_spawn_and_limit, @@ -32,8 +32,8 @@ fn max_frame_length(circuit: &mpz_circuits::Circuit, circuit_count: usize) -> us /// Records verifier->prover messages. #[cfg(target_arch = "wasm32")] async fn run_protocol_record_verifier( - exec_p: &mut Executor, - exec_v: &mut Executor, + exec_p: &mut Session, + exec_v: &mut Session, seed: u64, circuit_count: usize, ) { @@ -134,7 +134,7 @@ async fn record_for_prover(seed: u64, circuit_count: usize, concurrency: usize) /// Runs prover only with replay context. #[cfg(target_arch = "wasm32")] -async fn run_prover_with_replay(exec: &mut Executor, circuit_count: usize) { +async fn run_prover_with_replay(exec: &mut Session, circuit_count: usize) { let (_, ot_recv) = ideal_rcot([0u8; 16].into(), [0u8; 16].into()); let prover_config = ProverConfig::builder().build().unwrap(); let mut prover = Prover::new(prover_config, ot_recv); diff --git a/crates/wasm-bench/src/zk/verifier.rs b/crates/wasm-bench/src/zk/verifier.rs index 3128a9f1..f96df7c6 100644 --- a/crates/wasm-bench/src/zk/verifier.rs +++ b/crates/wasm-bench/src/zk/verifier.rs @@ -9,7 +9,7 @@ use std::sync::{Arc, Mutex}; use mpz_circuits::AES128; #[cfg(target_arch = "wasm32")] -use mpz_common::Executor; +use mpz_common::Session; use mpz_common::context::{ RecordedMtData, recording_mt_context_with_spawn_and_limit, replay_mt_context_with_spawn_and_limit, @@ -35,8 +35,8 @@ fn max_frame_length(circuit: &mpz_circuits::Circuit, circuit_count: usize) -> us /// Records prover->verifier messages. #[cfg(target_arch = "wasm32")] async fn run_protocol_record_prover( - exec_p: &mut Executor, - exec_v: &mut Executor, + exec_p: &mut Session, + exec_v: &mut Session, seed: u64, circuit_count: usize, ) { @@ -150,7 +150,7 @@ async fn record_for_verifier( /// Runs verifier only with replay context. #[cfg(target_arch = "wasm32")] async fn run_verifier_with_replay( - exec: &mut Executor, + exec: &mut Session, circuit_count: usize, delta: Delta, ot_seed: Block, diff --git a/crates/zk/benches/prover.rs b/crates/zk/benches/prover.rs index 7e91299f..e04f4074 100644 --- a/crates/zk/benches/prover.rs +++ b/crates/zk/benches/prover.rs @@ -8,7 +8,7 @@ use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_m use futures::executor::block_on; use mpz_circuits::AES128; use mpz_common::{ - Executor, + Session, context::{ RecordedMtData, recording_mt_context_with_limit, recording_st_context_with_limit, replay_mt_context_with_limit, replay_st_context, @@ -165,8 +165,8 @@ async fn run_prover_with_replay(ctx: &mut mpz_common::Context, circuit_count: us /// Runs the full ZK protocol with MT contexts. async fn run_protocol_record_verifier_mt( - exec_p: &mut Executor, - exec_v: &mut Executor, + exec_p: &mut Session, + exec_v: &mut Session, circuit_count: usize, seed: u64, ) { @@ -256,7 +256,7 @@ fn record_for_prover_mt(circuit_count: usize, seed: u64) -> RecordedMtData { } /// Runs MT prover only with replay context. -async fn run_prover_with_replay_mt(exec: &mut Executor, circuit_count: usize) { +async fn run_prover_with_replay_mt(exec: &mut Session, circuit_count: usize) { let (_, ot_recv) = ideal_rcot([0u8; 16].into(), [0u8; 16].into()); let prover_config = ProverConfig::builder().build().unwrap(); let mut prover = Prover::new(prover_config, ot_recv); diff --git a/crates/zk/benches/verifier.rs b/crates/zk/benches/verifier.rs index 1bfa658f..4ebcc291 100644 --- a/crates/zk/benches/verifier.rs +++ b/crates/zk/benches/verifier.rs @@ -9,7 +9,7 @@ use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_m use futures::executor::block_on; use mpz_circuits::AES128; use mpz_common::{ - Executor, + Session, context::{ RecordedMtData, recording_mt_context_with_limit, recording_st_context_with_limit, replay_mt_context_with_limit, replay_st_context, @@ -176,8 +176,8 @@ async fn run_verifier_with_replay( /// Runs the full ZK protocol with MT contexts. async fn run_protocol_record_prover_mt( - exec_p: &mut Executor, - exec_v: &mut Executor, + exec_p: &mut Session, + exec_v: &mut Session, circuit_count: usize, seed: u64, ) { @@ -273,7 +273,7 @@ fn record_for_verifier_mt(circuit_count: usize, seed: u64) -> (RecordedMtData, B /// Runs MT verifier only with replay context. async fn run_verifier_with_replay_mt( - exec: &mut Executor, + exec: &mut Session, circuit_count: usize, delta: Delta, ot_seed: Block,