Runtime context

crates/commitlog/Cargo.toml

@@ -41,6 +41,7 @@ zstd-framed.workspace = true
 # For the 'test' feature
 env_logger = { workspace = true, optional = true }
 pretty_assertions.workspace = true
+spacetimedb-runtime.workspace = true
 
 [dev-dependencies]
 # Enable streaming in tests

crates/commitlog/src/stream/reader.rs

@@ -4,10 +4,8 @@ use async_stream::try_stream;
 use bytes::{Buf as _, Bytes};
 use futures::Stream;
 use log::{trace, warn};
-use tokio::{
-    io::{self, AsyncBufRead, AsyncReadExt as _, AsyncSeek, AsyncSeekExt as _},
-    task::spawn_blocking,
-};
+use spacetimedb_runtime::spawn_blocking;
+use tokio::io::{self, AsyncBufRead, AsyncReadExt as _, AsyncSeek, AsyncSeekExt as _};
 use tokio_util::io::SyncIoBridge;
 
 use crate::{
@@ -107,8 +105,7 @@ fn read_segment(
                 }
                 segment.into_inner()
             })
-            .await
-            .unwrap();
+            .await;
         }
 
         let checksum_len = CHECKSUM_LEN[segment_header.checksum_algorithm as usize];

crates/commitlog/src/stream/writer.rs

@@ -5,10 +5,8 @@ use std::{
 
 use futures::TryFutureExt;
 use log::{debug, error, info, trace, warn};
-use tokio::{
-    io::{AsyncBufRead, AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt},
-    task::spawn_blocking,
-};
+use spacetimedb_runtime::{spawn, spawn_blocking};
+use tokio::io::{AsyncBufRead, AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt};
 
 use crate::{
     commit, error,
@@ -217,7 +215,6 @@ where
                     move || create_segment(repo, last_written_tx_range, commitlog_options, header)
                 })
                 .await
-                .unwrap()
                 .map(|(segment, index)| (segment.into_async_writer(), index))?;
                 stream.consume(segment::Header::LEN as _);
 
@@ -378,7 +375,7 @@ where
     fn drop(&mut self) {
         if let Some(current_segment) = self.current_segment.take() {
             trace!("closing current segment on writer drop");
-            tokio::spawn(
+            spawn(
                 current_segment
                     .close()
                     .inspect_err(|e| warn!("error closing segment on drop: {e}")),
@@ -425,8 +422,7 @@ impl<W: AsyncWriteExt + AsyncFsync + Unpin> CurrentSegment<W> {
                     .ok();
                 index
             })
-            .await
-            .unwrap();
+            .await;
             self.offset_index = Some(index);
         }
 

crates/dst/src/engine/workload.rs

@@ -7,6 +7,7 @@ use spacetimedb_sats::ArrayValue;
 
 use super::model::Model;
 use crate::schema::{SchemaPlan, TablePlan, Type};
+use crate::traits::InteractionGen;
 
 pub type Row = ProductValue;
 
@@ -260,6 +261,14 @@ impl WorkloadGen {
     }
 }
 
+impl InteractionGen<Observation> for WorkloadGen {
+    type Interaction = Interaction;
+
+    fn next_interaction(&mut self) -> Self::Interaction {
+        self.next_interaction()
+    }
+}
+
 impl Debug for WorkloadGen {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
         write!(f, "{:?}", self.stats())

crates/dst/src/lib.rs

@@ -3,4 +3,4 @@ pub mod schema;
 pub mod sim;
 pub mod traits;
 
-pub use traits::{Properties, TargetDriver, TestSuite, TestSuiteParts};
+pub use traits::{current_simulation_handle, InteractionGen, Properties, TargetDriver, TestSuite, TestSuiteParts};

crates/dst/src/traits.rs

@@ -16,6 +16,19 @@ pub trait Properties<I, O> {
     fn observe(&mut self, interaction: &I, observation: &O) -> Result<(), Error>;
 }
 
+/// Generates interactions, and can feed observations back to the generator so
+/// its internal state stays in sync with the target.
+pub trait InteractionGen<O>: std::fmt::Debug {
+    type Interaction: std::fmt::Debug;
+
+    fn next_interaction(&mut self) -> Self::Interaction;
+
+    /// Feed an observation back to the generator. Defaults to ignoring it.
+    fn observe(&mut self, _interaction: &Self::Interaction, _observation: &O) -> Result<(), Error> {
+        Ok(())
+    }
+}
+
 pub type TestSuiteParts<S> = (
     <S as TestSuite>::Interactions,
     <S as TestSuite>::Target,
@@ -24,7 +37,10 @@ pub type TestSuiteParts<S> = (
 
 pub trait TestSuite {
     type Interaction: std::fmt::Debug;
-    type Interactions: Iterator<Item = Self::Interaction> + std::fmt::Debug;
+    type Interactions: InteractionGen<
+        <Self::Target as TargetDriver<Self::Interaction>>::Observation,
+        Interaction = Self::Interaction,
+    >;
     type Target: TargetDriver<Self::Interaction>;
     type Properties: Properties<Self::Interaction, <Self::Target as TargetDriver<Self::Interaction>>::Observation>;
 
@@ -39,19 +55,16 @@ pub trait TestSuite {
         async move {
             let (mut interactions, mut target, mut properties) = self.build(rng).await?;
 
-            let result = async {
-                for interaction in interactions.by_ref().take(max_interactions) {
-                    let observation = target.execute(&interaction).await?;
-                    properties.observe(&interaction, &observation)?;
-                }
-
-                Ok(())
+            for _ in 0..max_interactions {
+                let interaction = interactions.next_interaction();
+                let observation = target.execute(&interaction).await?;
+                interactions.observe(&interaction, &observation)?;
+                properties.observe(&interaction, &observation)?;
             }
-            .await;
 
             tracing::info!(interaction_counts = ?interactions, "final interaction counts");
 
-            result
+            Ok(())
         }
     }
 }