Put ExecutionContext inside the Tx context (#1876)

Co-authored-by: joshua-spacetime <[email protected]>

crates/bench/benches/subscription.rs

@@ -1,6 +1,6 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use spacetimedb::error::DBError;
-use spacetimedb::execution_context::ExecutionContext;
+use spacetimedb::execution_context::Workload;
 use spacetimedb::host::module_host::DatabaseTableUpdate;
 use spacetimedb::identity::AuthCtx;
 use spacetimedb::messages::websocket::BsatnFormat;
@@ -52,9 +52,10 @@ fn eval(c: &mut Criterion) {
     let lhs = create_table_footprint(&raw.db).unwrap();
     let rhs = create_table_location(&raw.db).unwrap();
 
+    //TODO: Change this to `Workload::ForTest` once `#[cfg(bench)]` is stabilized.
     let _ = raw
         .db
-        .with_auto_commit(&ExecutionContext::default(), |tx| -> Result<(), DBError> {
+        .with_auto_commit(Workload::Internal, |tx| -> Result<(), DBError> {
             // 1M rows
             for entity_id in 0u64..1_000_000 {
                 let owner = entity_id % 1_000;
@@ -67,7 +68,7 @@ fn eval(c: &mut Criterion) {
 
     let _ = raw
         .db
-        .with_auto_commit(&ExecutionContext::default(), |tx| -> Result<(), DBError> {
+        .with_auto_commit(Workload::Internal, |tx| -> Result<(), DBError> {
             // 1000 chunks, 1200 rows per chunk = 1.2M rows
             for chunk_index in 0u64..1_000 {
                 for i in 0u64..1200 {
@@ -100,13 +101,12 @@ fn eval(c: &mut Criterion) {
 
     let bench_eval = |c: &mut Criterion, name, sql| {
         c.bench_function(name, |b| {
-            let tx = raw.db.begin_tx();
+            let tx = raw.db.begin_tx(Workload::Update);
             let query = compile_read_only_query(&raw.db, &AuthCtx::for_testing(), &tx, sql).unwrap();
             let query: ExecutionSet = query.into();
-            let ctx = &ExecutionContext::subscribe(raw.db.database_identity());
+
             b.iter(|| {
                 drop(black_box(query.eval::<BsatnFormat>(
-                    ctx,
                     &raw.db,
                     &tx,
                     None,
@@ -134,25 +134,19 @@ fn eval(c: &mut Criterion) {
     );
     bench_eval(c, "full-join", &name);
 
-    let ctx_incr = &ExecutionContext::incremental_update(raw.db.database_identity());
-
     // To profile this benchmark for 30s
     // samply record -r 10000000 cargo bench --bench=subscription --profile=profiling -- incr-select --exact --profile-time=30
     c.bench_function("incr-select", |b| {
         // A passthru executed independently of the database.
         let select_lhs = "select * from footprint";
         let select_rhs = "select * from location";
-        let tx = &raw.db.begin_tx();
+        let tx = &raw.db.begin_tx(Workload::Update);
         let query_lhs = compile_read_only_query(&raw.db, &AuthCtx::for_testing(), tx, select_lhs).unwrap();
         let query_rhs = compile_read_only_query(&raw.db, &AuthCtx::for_testing(), tx, select_rhs).unwrap();
         let query = ExecutionSet::from_iter(query_lhs.into_iter().chain(query_rhs));
         let tx = &tx.into();
 
-        b.iter(|| {
-            drop(black_box(
-                query.eval_incr_for_test(ctx_incr, &raw.db, tx, &update, None),
-            ))
-        })
+        b.iter(|| drop(black_box(query.eval_incr_for_test(&raw.db, tx, &update, None))))
     });
 
     // To profile this benchmark for 30s
@@ -165,16 +159,12 @@ fn eval(c: &mut Criterion) {
             from footprint join location on footprint.entity_id = location.entity_id \
             where location.chunk_index = {chunk_index}"
         );
-        let tx = &raw.db.begin_tx();
+        let tx = &raw.db.begin_tx(Workload::Update);
         let query = compile_read_only_query(&raw.db, &AuthCtx::for_testing(), tx, &join).unwrap();
         let query: ExecutionSet = query.into();
         let tx = &tx.into();
 
-        b.iter(|| {
-            drop(black_box(
-                query.eval_incr_for_test(ctx_incr, &raw.db, tx, &update, None),
-            ))
-        });
+        b.iter(|| drop(black_box(query.eval_incr_for_test(&raw.db, tx, &update, None))));
     });
 
     // To profile this benchmark for 30s

crates/bench/src/spacetime_raw.rs

@@ -4,7 +4,7 @@ use crate::{
     ResultBench,
 };
 use spacetimedb::db::relational_db::{tests_utils::TestDB, RelationalDB};
-use spacetimedb::execution_context::ExecutionContext;
+use spacetimedb::execution_context::Workload;
 use spacetimedb_lib::sats::AlgebraicValue;
 use spacetimedb_primitives::{ColId, IndexId, TableId};
 use spacetimedb_schema::{
@@ -40,7 +40,7 @@ impl BenchDatabase for SpacetimeRaw {
 
     fn create_table<T: BenchTable>(&mut self, index_strategy: IndexStrategy) -> ResultBench<Self::TableId> {
         let name = table_name::<T>(index_strategy);
-        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
+        self.db.with_auto_commit(Workload::Internal, |tx| {
             let mut table_schema = TableSchema::from_product_type(T::product_type());
             table_schema.table_name = name.clone().into();
             let table_id = self.db.create_table(tx, table_schema)?;
@@ -84,25 +84,24 @@ impl BenchDatabase for SpacetimeRaw {
     }
 
     fn clear_table(&mut self, table_id: &Self::TableId) -> ResultBench<()> {
-        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
+        self.db.with_auto_commit(Workload::Internal, |tx| {
             self.db.clear_table(tx, *table_id)?;
             Ok(())
         })
     }
 
     fn count_table(&mut self, table_id: &Self::TableId) -> ResultBench<u32> {
-        let ctx = ExecutionContext::default();
-        self.db.with_auto_commit(&ctx, |tx| {
-            Ok(self.db.iter_mut(&ctx, tx, *table_id)?.map(|_| 1u32).sum())
+        self.db.with_auto_commit(Workload::Internal, |tx| {
+            Ok(self.db.iter_mut(tx, *table_id)?.map(|_| 1u32).sum())
         })
     }
 
     fn empty_transaction(&mut self) -> ResultBench<()> {
-        self.db.with_auto_commit(&ExecutionContext::default(), |_tx| Ok(()))
+        self.db.with_auto_commit(Workload::Internal, |_tx| Ok(()))
     }
 
     fn insert_bulk<T: BenchTable>(&mut self, table_id: &Self::TableId, rows: Vec<T>) -> ResultBench<()> {
-        self.db.with_auto_commit(&ExecutionContext::default(), |tx| {
+        self.db.with_auto_commit(Workload::Internal, |tx| {
             for row in rows {
                 self.db.insert(tx, *table_id, row.into_product_value())?;
             }
@@ -111,11 +110,10 @@ impl BenchDatabase for SpacetimeRaw {
     }
 
     fn update_bulk<T: BenchTable>(&mut self, table_id: &Self::TableId, row_count: u32) -> ResultBench<()> {
-        let ctx = ExecutionContext::default();
-        self.db.with_auto_commit(&ctx, |tx| {
+        self.db.with_auto_commit(Workload::Internal, |tx| {
             let rows = self
                 .db
-                .iter_mut(&ctx, tx, *table_id)?
+                .iter_mut(tx, *table_id)?
                 .take(row_count as usize)
                 .map(|row| row.to_product_value())
                 .collect::<Vec<_>>();
@@ -127,7 +125,7 @@ impl BenchDatabase for SpacetimeRaw {
                 // (update_by_{field} -> spacetimedb::query::update_by_field -> (delete_by_col_eq; insert))
                 let id = self
                     .db
-                    .iter_by_col_eq_mut(&ctx, tx, *table_id, 0, &row.elements[0])?
+                    .iter_by_col_eq_mut(tx, *table_id, 0, &row.elements[0])?
                     .next()
                     .expect("failed to find row during update!")
                     .pointer();
@@ -152,9 +150,8 @@ impl BenchDatabase for SpacetimeRaw {
     }
 
     fn iterate(&mut self, table_id: &Self::TableId) -> ResultBench<()> {
-        let ctx = ExecutionContext::default();
-        self.db.with_auto_commit(&ctx, |tx| {
-            for row in self.db.iter_mut(&ctx, tx, *table_id)? {
+        self.db.with_auto_commit(Workload::Internal, |tx| {
+            for row in self.db.iter_mut(tx, *table_id)? {
                 black_box(row);
             }
             Ok(())
@@ -167,9 +164,8 @@ impl BenchDatabase for SpacetimeRaw {
         col_id: impl Into<ColId>,
         value: AlgebraicValue,
     ) -> ResultBench<()> {
-        let ctx = ExecutionContext::default();
-        self.db.with_auto_commit(&ctx, |tx| {
-            for row in self.db.iter_by_col_eq_mut(&ctx, tx, *table_id, col_id, &value)? {
+        self.db.with_auto_commit(Workload::Internal, |tx| {
+            for row in self.db.iter_by_col_eq_mut(tx, *table_id, col_id, &value)? {
                 black_box(row);
             }
             Ok(())

crates/client-api/src/routes/database.rs

@@ -16,6 +16,7 @@ use serde::Deserialize;
 use serde_json::{json, Value};
 use spacetimedb::address::Address;
 use spacetimedb::database_logger::DatabaseLogger;
+use spacetimedb::execution_context::Workload;
 use spacetimedb::host::ReducerCallError;
 use spacetimedb::host::ReducerOutcome;
 use spacetimedb::host::{DescribedEntityType, UpdateDatabaseResult};
@@ -24,7 +25,7 @@ use spacetimedb::identity::Identity;
 use spacetimedb::json::client_api::StmtResultJson;
 use spacetimedb::messages::control_db::{Database, HostType, Replica};
 use spacetimedb::sql;
-use spacetimedb::sql::execute::{ctx_sql, translate_col};
+use spacetimedb::sql::execute::translate_col;
 use spacetimedb_client_api_messages::name::{self, DnsLookupResponse, DomainName, PublishOp, PublishResult};
 use spacetimedb_data_structures::map::HashMap;
 use spacetimedb_lib::address::AddressForUrl;
@@ -557,7 +558,7 @@ where
                         }
                     })?;
 
-                let json = db.with_read_only(&ctx_sql(db), |tx| {
+                let json = db.with_read_only(Workload::Sql, |tx| {
                     results
                         .into_iter()
                         .map(|result| {

crates/core/src/db/datastore/locking_tx_datastore/committed_state.rs

@@ -77,26 +77,25 @@ impl StateView for CommittedState {
         self.get_table(table_id).map(|table| table.row_count)
     }
 
-    fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>> {
+    fn iter(&self, table_id: TableId) -> Result<Iter<'_>> {
         if let Some(table_name) = self.table_name(table_id) {
-            return Ok(Iter::new(ctx, table_id, table_name, None, self));
+            return Ok(Iter::new(table_id, table_name, None, self));
         }
         Err(TableError::IdNotFound(SystemTable::st_table, table_id.0).into())
     }
     /// Returns an iterator,
     /// yielding every row in the table identified by `table_id`,
     /// where the values of `cols` are contained in `range`.
-    fn iter_by_col_range<'a, R: RangeBounds<AlgebraicValue>>(
-        &'a self,
-        ctx: &'a ExecutionContext,
+    fn iter_by_col_range<R: RangeBounds<AlgebraicValue>>(
+        &self,
         table_id: TableId,
         cols: ColList,
         range: R,
-    ) -> Result<IterByColRange<'a, R>> {
+    ) -> Result<IterByColRange<'_, R>> {
         // TODO: Why does this unconditionally return a `Scan` iter,
         // instead of trying to return a `CommittedIndex` iter?
         Ok(IterByColRange::Scan(ScanIterByColRange::new(
-            self.iter(ctx, table_id)?,
+            self.iter(table_id)?,
             cols,
             range,
         )))
@@ -274,7 +273,7 @@ impl CommittedState {
             with_label_values(ST_SEQUENCE_ID, ST_SEQUENCE_NAME).inc();
         }
 
-        self.reset_system_table_schemas(database_identity)?;
+        self.reset_system_table_schemas()?;
 
         Ok(())
     }
@@ -286,12 +285,11 @@ impl CommittedState {
     /// for objects like indexes and constraints
     /// which are computed at insert-time,
     /// and therefore not included in the hardcoded schemas.
-    pub(super) fn reset_system_table_schemas(&mut self, database_identity: Identity) -> Result<()> {
+    pub(super) fn reset_system_table_schemas(&mut self) -> Result<()> {
         // Re-read the schema with the correct ids...
-        let ctx = ExecutionContext::internal(database_identity);
         for schema in system_tables() {
             self.tables.get_mut(&schema.table_id).unwrap().schema =
-                Arc::new(self.schema_for_table_raw(&ctx, schema.table_id)?);
+                Arc::new(self.schema_for_table_raw(schema.table_id)?);
         }
 
         Ok(())
@@ -384,7 +382,7 @@ impl CommittedState {
         // For already built tables, we need to reschema them to account for constraints et al.
         let mut schemas = Vec::with_capacity(self.tables.len());
         for table_id in self.tables.keys().copied() {
-            schemas.push(self.schema_for_table_raw(&ExecutionContext::default(), table_id)?);
+            schemas.push(self.schema_for_table_raw(table_id)?);
         }
         for (table, schema) in self.tables.values_mut().zip(schemas) {
             table.with_mut_schema(|s| *s = schema);
@@ -407,7 +405,7 @@ impl CommittedState {
 
         // Construct their schemas and insert tables for them.
         for table_id in table_ids {
-            let schema = self.schema_for_table(&ExecutionContext::default(), table_id)?;
+            let schema = self.schema_for_table(table_id)?;
             self.tables.insert(table_id, Self::make_table(schema));
         }
         Ok(())
@@ -467,12 +465,7 @@ impl CommittedState {
         // Note that this may change in the future: some analytics and/or
         // timetravel queries may benefit from seeing all inputs, even if
         // the database state did not change.
-        tx_data.inserts().any(|(_, inserted_rows)| !inserted_rows.is_empty())
-            || tx_data.deletes().any(|(_, deleted_rows)| !deleted_rows.is_empty())
-            || matches!(
-                ctx.reducer_context().map(|rcx| rcx.name.strip_prefix("__identity_")),
-                Some(Some("connected__" | "disconnected__"))
-            )
+        tx_data.has_rows_or_connect_disconnect(ctx.reducer_context())
     }
 
     pub(super) fn merge(&mut self, tx_state: TxState, ctx: &ExecutionContext) -> TxData {
@@ -636,8 +629,6 @@ impl CommittedState {
 }
 
 pub struct CommittedIndexIter<'a> {
-    #[allow(dead_code)]
-    ctx: &'a ExecutionContext,
     table_id: TableId,
     tx_state: Option<&'a TxState>,
     #[allow(dead_code)]
@@ -648,14 +639,12 @@ pub struct CommittedIndexIter<'a> {
 
 impl<'a> CommittedIndexIter<'a> {
     pub(super) fn new(
-        ctx: &'a ExecutionContext,
         table_id: TableId,
         tx_state: Option<&'a TxState>,
         committed_state: &'a CommittedState,
         committed_rows: IndexScanIter<'a>,
     ) -> Self {
         Self {
-            ctx,
             table_id,
             tx_state,
             committed_state,