[runtime/tokio] support for colocated tasks

Cargo.toml

@@ -74,6 +74,7 @@ suspicious_op_assign_impl = "allow"
 [workspace.dependencies]
 anyhow = { version = "1.0.99", default-features = false }
 arbitrary = "1.4.1"
+async-task = "4.7.1"
 aws-config = "1.8.12"
 aws-lc-rs = "1.15.2"
 aws-sdk-ec2 = "1.200.0"

runtime/Cargo.toml

@@ -15,6 +15,7 @@ workspace = true
 
 [dependencies]
 arbitrary = { workspace = true, optional = true }
+async-task.workspace = true
 bytes.workspace = true
 cfg-if.workspace = true
 commonware-codec.workspace = true

runtime/src/deterministic.rs

@@ -885,6 +885,18 @@ impl Tasks {
 type Network = MeteredNetwork<AuditedNetwork<DeterministicNetwork>>;
 type Storage = MeteredStorage<AuditedStorage<FaultyStorage<MemStorage>>>;
 
+/// Marks a dedicated subtree as closed before the parent task's result is made
+/// observable through its handle.
+struct DedicatedGuard {
+    tree: Arc<Tree>,
+}
+
+impl Drop for DedicatedGuard {
+    fn drop(&mut self) {
+        self.tree.abort();
+    }
+}
+
 /// Implementation of [crate::Spawner], [crate::Clock],
 /// [crate::Network], and [crate::Storage] for the `deterministic`
 /// runtime.
@@ -900,6 +912,7 @@ pub struct Context {
     tree: Arc<Tree>,
     execution: Execution,
     traced: bool,
+    dedicated: Option<Arc<Tree>>,
 }
 
 impl Clone for Context {
@@ -918,6 +931,7 @@ impl Clone for Context {
             tree: child,
             execution: Execution::default(),
             traced: false,
+            dedicated: self.dedicated.clone(),
         }
     }
 }
@@ -999,6 +1013,7 @@ impl Context {
                 tree: Tree::root(),
                 execution: Execution::default(),
                 traced: false,
+                dedicated: None,
             },
             executor,
             panicked,
@@ -1071,6 +1086,7 @@ impl Context {
                 tree: Tree::root(),
                 execution: Execution::default(),
                 traced: false,
+                dedicated: None,
             },
             executor,
             panicked,
@@ -1138,6 +1154,11 @@ impl crate::Spawner for Context {
         self
     }
 
+    fn colocated(mut self) -> Self {
+        self.execution = Execution::Colocated;
+        self
+    }
+
     fn shared(mut self, blocking: bool) -> Self {
         self.execution = Execution::Shared(blocking);
         self
@@ -1154,14 +1175,38 @@ impl crate::Spawner for Context {
 
         // Track supervision before resetting configuration
         let parent = Arc::clone(&self.tree);
+        let past = self.execution;
         let traced = self.traced;
+        let inherited_dedicated = self.dedicated.clone();
         self.execution = Execution::default();
         self.traced = false;
+        if matches!(past, Execution::Colocated) {
+            let dedicated = inherited_dedicated
+                .as_ref()
+                .expect("`colocated()` requires a running dedicated ancestor");
+            assert!(
+                !dedicated.is_aborted(),
+                "`colocated()` requires a running dedicated ancestor"
+            );
+        }
         let (child, aborted) = Tree::child(&parent);
         if aborted {
             return Handle::closed(metric);
         }
+        let child_dedicated = match past {
+            // Dedicated creates a new execution domain for the spawned child.
+            // In the deterministic runtime, the child context node itself is
+            // sufficient to model that domain because the supervision tree
+            // already tracks whether the dedicated root task is still alive.
+            Execution::Dedicated => Some(child.clone()),
+            // Colocated reuses the nearest dedicated ancestor, which must
+            // already have been validated before creating the child node.
+            Execution::Colocated => inherited_dedicated,
+            // Shared clears the dedicated branch assignment.
+            Execution::Shared(_) => None,
+        };
         self.tree = child;
+        self.dedicated = child_dedicated.clone();
 
         // Spawn the task (we don't care about Model)
         let executor = self.executor();
@@ -1174,6 +1219,15 @@ impl crate::Spawner for Context {
         } else {
             Either::Right(f(self))
         };
+        let future = if matches!(past, Execution::Dedicated) {
+            let dedicated = child_dedicated.expect("dedicated tree missing");
+            Either::Left(async move {
+                let _guard = DedicatedGuard { tree: dedicated };
+                future.await
+            })
+        } else {
+            Either::Right(future)
+        };
         let (f, handle) = Handle::init(
             future,
             metric,