Implement coroutine tasks and nested future waits

Author: samuel-williams-shopify

agent-context.md

@@ -7,8 +7,9 @@ layout.
 ## Workspace
 
 - `socketry` is the public facade and re-exports the common concurrency API.
-- `socketry-concurrent` owns the stackful fiber implementation and the
-  cooperative scheduler.
+- `socketry-concurrent` owns the coroutine implementation and the cooperative
+  future scheduler. `Scheduler::spawn` accepts futures; `wait` lets ordinary
+  synchronous functions suspend on futures using the current task's stack.
 - Each workspace package is published independently to crates.io.
 
 ## Implementation boundaries
@@ -17,7 +18,13 @@ layout.
   revision and attribution are recorded in
   `crates/concurrent/vendor/cruby/upstream.md` and
   `crates/concurrent/vendor/cruby/license.md`.
-- Suspended fibers are thread-affine. Do not migrate their stacks between
+- Suspended tasks are thread-affine. Do not migrate their stacks between
   worker threads.
 - Future wakers may enqueue work from another thread; the owning scheduler
-  resumes its fibers.
+  resumes its tasks.
+- A nested `wait` preserves an in-progress `Future::poll` call. Resume that
+  stack before polling the outer future again. Preserve outer wakeups consumed
+  while nested waits are active.
+- Both nested waits and ordinary future suspension currently stay on the
+  owning thread. Migration between completed polls is a possible future
+  extension and is not implemented.

crates/concurrent/Cargo.toml

@@ -2,11 +2,11 @@
 name = "socketry-concurrent"
 version.workspace = true
 edition.workspace = true
-description = "Stackful fibers and a small cooperative scheduler"
+description = "A cooperative future scheduler with nested stackful waits"
 license.workspace = true
 repository.workspace = true
 readme = "readme.md"
-include = ["Cargo.toml", "build.rs", "license.md", "readme.md", "src/**", "vendor/**"]
+include = ["Cargo.toml", "build.rs", "license.md", "readme.md", "src/**", "examples/**", "vendor/**"]
 
 [lib]
 name = "socketry_concurrent"

crates/concurrent/examples/nested_wait.rs

@@ -0,0 +1,55 @@
+use socketry_concurrent::{Scheduler, wait};
+use std::cell::RefCell;
+use std::future::poll_fn;
+use std::rc::Rc;
+use std::task::{Poll, Waker};
+
+#[derive(Default)]
+struct Message {
+    value: Option<String>,
+    reader: Option<Waker>,
+}
+
+// An ordinary synchronous function can wait for a future at any call depth.
+fn read_message(message: &RefCell<Message>) -> String {
+    wait(poll_fn(|context| {
+        let mut message = message.borrow_mut();
+        if let Some(value) = message.value.take() {
+            Poll::Ready(value)
+        } else {
+            message.reader = Some(context.waker().clone());
+            Poll::Pending
+        }
+    }))
+}
+
+fn process_message(message: &RefCell<Message>) {
+    println!("consumer: waiting");
+    let value = read_message(message);
+    println!("consumer: received {value}");
+}
+
+fn main() -> std::io::Result<()> {
+    let mut scheduler = Scheduler::new(256 * 1024);
+    let message = Rc::new(RefCell::new(Message::default()));
+
+    let consumer_message = Rc::clone(&message);
+    scheduler.spawn(async move {
+        process_message(&consumer_message);
+    })?;
+
+    scheduler.spawn(async move {
+        println!("producer: sending");
+        let reader = {
+            let mut message = message.borrow_mut();
+            message.value = Some(String::from("hello"));
+            message.reader.take()
+        };
+        if let Some(reader) = reader {
+            reader.wake();
+        }
+    })?;
+
+    scheduler.run();
+    Ok(())
+}

crates/concurrent/readme.md

@@ -1,41 +1,81 @@
 # socketry-concurrent
 
-Stackful fibers and cooperative scheduling for Socketry's Rust packages. Use
+Futures with nested stackful waits for Socketry's Rust packages. Use
 the `socketry` package for the one-stop public entry point, or depend on this
 package directly to use the concurrency implementation on its own.
 
-The API starts with three pieces:
+The library provides three building blocks:
 
-- Stack reserves memory with inaccessible guard pages at both ends.
-- Fiber owns a stack and switches between its caller and a closure.
-- Pool keeps multiple stacks available for reuse.
+- `Stack` reserves memory with inaccessible guard pages at both ends.
+- `Pool` keeps multiple stacks available for reuse.
+- `Scheduler` accepts futures and polls each task on its own coroutine stack.
 
-Scheduler is an optional, single-threaded executor. It can run Rust futures and
-also exposes block_current and task unblock operations for stackful code.
-Future wakers may run on any thread; the fiber itself always resumes on the
-thread that owns its stack. Arbitrary Rust locals on a suspended stack are not
-tracked by the type system, so moving a suspended fiber between workers is not
-safe. The scheduler does not migrate fibers.
+`wait(future)` lets an ordinary function wait for an asynchronous result. When
+the future is pending, the task's stack is suspended and the executor can run
+other tasks. The function returns the future's output after it completes. Calls
+can be nested, including inside another future's `poll`, without making the
+calling functions async.
+
+`Task::current` provides a thread-affine reference for blocking and task-to-task
+transfer. `Scheduler::spawn` returns a separate, thread-safe `TaskHandle` for
+waking the task. Future wakers may run on any thread; each task always resumes
+on the thread that owns its stack.
 
 ## Example
 
-    use socketry_concurrent::Scheduler;
+    use socketry_concurrent::{Scheduler, wait};
+    use std::future::poll_fn;
+    use std::task::Poll;
+
+    fn answer() -> usize {
+        let mut first_poll = true;
+        wait(poll_fn(|context| {
+            if first_poll {
+                first_poll = false;
+                context.waker().wake_by_ref();
+                Poll::Pending
+            } else {
+                Poll::Ready(42)
+            }
+        }))
+    }
 
     fn main() -> std::io::Result<()> {
         let mut scheduler = Scheduler::new(256 * 1024);
         scheduler.spawn(async {
-            // Await ordinary Rust futures here.
+            assert_eq!(answer(), 42);
         })?;
         scheduler.run();
         Ok(())
     }
 
-Fiber::yield_now returns control to the caller. Resuming that fiber continues
-after the yield. Dropping a suspended fiber resumes it with a private
-cancellation panic so Rust unwinds the stack and runs local destructors.
+`Scheduler::current().unwrap().wait(future)` is also available. Both forms
+require a current scheduler task. A future can borrow local data and does not
+need to implement Send or Unpin. Each wait pins the future on the task stack
+and reuses the task's wake signal, without allocating a separate future or
+waker. A future from another runtime still needs that runtime's I/O, timer,
+or other services to be running.
+
+Run the producer/consumer example with:
+
+    cargo run --package socketry-concurrent --example nested_wait
+
+Use `Task::current().unwrap().block()` or `Scheduler::block_current()` to park a
+task. Another thread can make it runnable through its `TaskHandle`. Dropping a
+scheduler unwinds suspended task stacks and runs their local destructors.
+
+## Polling and nested waits
+
+An ordinary `.await` can return `Poll::Pending` from the task's future. A nested
+`wait`, however, can suspend while that same `poll` call is still executing.
+The scheduler resumes the saved stack at that wait; it does not call the outer
+future's `poll` again until the previous invocation has returned. Wakeups for
+outer futures are preserved while inner waits run.
 
-Use Scheduler::spawn_fiber for synchronous stackful tasks. Such a task can call
-Scheduler::block_current and be made runnable again through its TaskHandle.
+This implementation uses one scheduler thread for both cases. A future
+executor could move a Send future between completed polls, but a stack
+suspended inside a poll must remain on its original worker. The current
+scheduler implements neither work stealing nor cross-thread migration.
 
 ## Native context switches
 
@@ -53,14 +93,14 @@ vendored tree for reference.
 
 The x86-64 Linux build includes CRuby's CET shadow-stack switch path. It checks
 whether shadow stacks are enabled at runtime and allocates a shadow stack for
-each fiber only when needed.
+each task only when needed.
 
 ## Sanitizers
 
 The `address-sanitizer` and `thread-sanitizer` Cargo features enable the
 compiler runtime's fiber-switch hooks. Pair one feature with the matching Rust
 sanitizer flag on nightly; the hooks let each runtime follow the custom stacks
-used by Fiber.
+used by scheduled tasks.
 
     RUSTFLAGS="-Zsanitizer=address" cargo +nightly test -Zbuild-std --target aarch64-apple-darwin --package socketry-concurrent --features address-sanitizer
 
@@ -78,10 +118,10 @@ embedded header.
 
 ## Current boundaries
 
-- Fiber and Scheduler are thread-affine and are not Send.
+- Task contexts and Scheduler are thread-affine and are not Send.
 - Scheduler tasks may contain non-Send futures because they are polled only on
   the scheduler's owning thread.
 - TaskHandle::unblock and future wakers are thread-safe and only enqueue work;
-  they never resume a fiber on the waking thread.
+  they never resume a task on the waking thread.
 - This first version does not implement work stealing or cross-thread stack
   migration.

crates/concurrent/src/context.rs

@@ -294,7 +294,7 @@ impl Drop for Context {
         #[cfg(coroutine_thread_sanitizer)]
         if self.thread_sanitizer_fiber_owned != 0 && !self.thread_sanitizer_fiber.is_null() {
             // SAFETY: this context owns the fiber handle and is no longer
-            // running when its containing Fiber is dropped.
+            // running when its containing task coroutine is dropped.
             unsafe {
                 thread_sanitizer_destroy_fiber(self.thread_sanitizer_fiber);
             }