From edb089637a16623e2417a1954a934d4194c6a38a Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Sun, 12 Jul 2026 14:15:01 +0530 Subject: [PATCH 01/10] migrations in dst --- .../src/locking_tx_datastore/replay.rs | 2 +- crates/dst/src/engine.rs | 327 +++++++++++++++++- crates/dst/src/engine/migrations.rs | 246 +++++++++++++ crates/dst/src/engine/model.rs | 24 +- crates/dst/src/engine/workload.rs | 157 ++++++++- crates/dst/src/schema.rs | 279 +++++++++++---- 6 files changed, 953 insertions(+), 82 deletions(-) create mode 100644 crates/dst/src/engine/migrations.rs diff --git a/crates/datastore/src/locking_tx_datastore/replay.rs b/crates/datastore/src/locking_tx_datastore/replay.rs index d803360551d..d439553c01f 100644 --- a/crates/datastore/src/locking_tx_datastore/replay.rs +++ b/crates/datastore/src/locking_tx_datastore/replay.rs @@ -1085,7 +1085,7 @@ impl StateView for ReplayCommittedState<'_> { /// then falling back to [`CommittedState::iter_by_col_eq`]. fn find_st_table_row(&self, table_id: TableId) -> Result { if let Some(row_ptr) = self.replay_table_updated.get(&table_id) { - let (table, blob_store, _) = self.state.get_table_and_blob_store(table_id)?; + let (table, blob_store, _) = self.state.get_table_and_blob_store(ST_TABLE_ID)?; // SAFETY: `row_ptr` is stored in `self.replay_table_updated`, // meaning it was inserted into `st_table` by `replay_insert` // and has not yet been deleted by `replay_delete_by_rel`. diff --git a/crates/dst/src/engine.rs b/crates/dst/src/engine.rs index 451b7417e41..006c87c6b57 100644 --- a/crates/dst/src/engine.rs +++ b/crates/dst/src/engine.rs @@ -6,21 +6,27 @@ use spacetimedb_datastore::traits::{IsolationLevel, TxData}; use spacetimedb_engine::error::{DBError, DatastoreError, IndexError}; use spacetimedb_engine::persistence::{DiskSizeFn, Durability as EngineDurability, Persistence}; use spacetimedb_engine::relational_db::{MutTx, RelationalDB}; +use spacetimedb_engine::update::{update_database, UpdateLogger}; +use spacetimedb_lib::db::auth::StAccess; +use spacetimedb_lib::identity::AuthCtx; use spacetimedb_lib::{Identity, RawModuleDef}; use spacetimedb_primitives::TableId; use spacetimedb_runtime::sim::{Rng, Runtime as SimRuntime}; use spacetimedb_runtime::Handle; -use spacetimedb_schema::def::ModuleDef; +use spacetimedb_schema::auto_migrate::{ponder_migrate, AutoMigrateStep, MigratePlan}; +use spacetimedb_schema::def::{IndexAlgorithm as SchemaIndexAlgorithm, ModuleDef}; use spacetimedb_schema::schema::{Schema, TableSchema}; use spacetimedb_table::page_pool::PagePool; +mod migrations; mod model; mod properties; mod workload; +use self::migrations::{ExpectedStep, Migration}; use self::workload::{ - normalize_rows, row_to_bytes, CommitDelta, CountState, InsertOutcome, Interaction, Observation, TableDelta, - TableRowCount, + normalize_rows, row_to_bytes, ColumnState, CommitDelta, CountState, IndexAlgorithmState, IndexState, InsertOutcome, + Interaction, Observation, SchemaState, SequenceState, TableDelta, TableRowCount, TableSchemaState, }; use crate::engine::model::Model; @@ -36,6 +42,7 @@ pub struct EngineTarget { active_mut_tx: Option, commitlog: InMemoryCommitlog, runtime_handle: Handle, + schema: SchemaPlan, } impl EngineTarget { @@ -54,6 +61,7 @@ impl EngineTarget { active_mut_tx: None, commitlog, runtime_handle, + schema, }) } @@ -88,11 +96,14 @@ impl EngineTarget { } } - fn install_schema(db: &RelationalDB, schema: &SchemaPlan) -> anyhow::Result<()> { + fn module_def(schema: &SchemaPlan) -> anyhow::Result { let raw = to_raw_def(schema); let raw_module_def = RawModuleDef::V10(raw); - let module_def = - ModuleDef::try_from(raw_module_def).map_err(|e| anyhow::anyhow!("schema validation failed: {e}"))?; + ModuleDef::try_from(raw_module_def).map_err(|e| anyhow::anyhow!("schema validation failed: {e}")) + } + + fn install_schema(db: &RelationalDB, schema: &SchemaPlan) -> anyhow::Result<()> { + let module_def = Self::module_def(schema)?; db.with_auto_commit(Workload::Internal, |tx| -> Result<(), DBError> { for table_def in module_def.tables() { @@ -138,6 +149,7 @@ impl EngineTarget { .ok_or_else(|| anyhow::anyhow!("database is not open"))?; let tx = db.begin_tx(Workload::Internal); let mut row_counts = Vec::with_capacity(self.table_ids.len()); + let mut schema_tables = Vec::with_capacity(self.table_ids.len()); for (table, table_id) in self.table_ids.iter().enumerate() { let count = match db.iter(&tx, *table_id) { @@ -148,10 +160,57 @@ impl EngineTarget { } }; row_counts.push(TableRowCount { table, count }); + + let schema = match db.schema_for_table(&tx, *table_id) { + Ok(schema) => schema, + Err(err) => { + let _ = db.release_tx(tx); + return Err(err.into()); + } + }; + let mut indexes = schema + .indexes + .iter() + .map(|index| IndexState { + columns: index_columns(&index.index_algorithm), + algorithm: index_algorithm_state(&index.index_algorithm), + }) + .collect::>(); + indexes.sort(); + + let mut sequences = schema + .sequences + .iter() + .map(|sequence| SequenceState { + column: sequence.col_pos.0 as usize, + }) + .collect::>(); + sequences.sort(); + + schema_tables.push(TableSchemaState { + table, + name: schema.table_name.to_string(), + is_public: schema.table_access == StAccess::Public, + is_event: schema.is_event, + primary_key: schema.primary_key.map(|col| col.0 as usize), + columns: schema + .columns + .iter() + .map(|column| ColumnState { + name: column.col_name.to_string(), + ty: column.col_type.clone(), + }) + .collect(), + indexes, + sequences, + }); } let _ = db.release_tx(tx); - Ok(CountState { row_counts }) + Ok(CountState { + row_counts, + schema: SchemaState { tables: schema_tables }, + }) } fn is_unique_constraint_violation(error: &DBError) -> bool { @@ -187,6 +246,55 @@ impl EngineTarget { CommitDelta { tables } } + fn migrate(&mut self, migration: &Migration) -> anyhow::Result<()> { + anyhow::ensure!( + self.active_mut_tx.is_none(), + "migration while mutable transaction is active" + ); + + let new_schema = migration.apply_to(&self.schema)?; + let old_module_def = Self::module_def(&self.schema)?; + let new_module_def = Self::module_def(&new_schema)?; + let plan = ponder_migrate(&old_module_def, &new_module_def)?; + self.ensure_expected_plan(migration, &plan)?; + + let db = self + .db + .as_ref() + .ok_or_else(|| anyhow::anyhow!("database is not open"))?; + let mut tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::Internal); + let _ = update_database(db, &mut tx, AuthCtx::for_testing(), plan, &DstUpdateLogger)?; + let Some((_tx_offset, _tx_data, _tx_metrics, _reducer)) = db.commit_tx(tx)? else { + anyhow::bail!("migration commit produced no transaction data"); + }; + + self.schema = new_schema; + self.table_ids = Self::load_table_ids(db, &self.schema)?; + Ok(()) + } + + fn ensure_expected_plan(&self, migration: &Migration, plan: &MigratePlan<'_>) -> anyhow::Result<()> { + let MigratePlan::Auto(plan) = plan else { + anyhow::bail!("engine DST generated a manual migration plan"); + }; + + let mut actual = plan + .steps + .iter() + .map(expected_step_from_auto_step) + .collect::>>()?; + actual.sort(); + + let mut expected = migration.expected_steps(); + expected.sort(); + + anyhow::ensure!( + actual == expected, + "engine DST generated unexpected migration steps: actual={actual:?}, expected={expected:?}" + ); + Ok(()) + } + pub fn execute(&mut self, interaction: &Interaction) -> anyhow::Result { tracing::debug!(?interaction, "executing interaction"); @@ -253,6 +361,10 @@ impl EngineTarget { delta: self.commit_delta_from_tx_data(&tx_data), }) } + Interaction::Migrate(migration) => { + self.migrate(migration)?; + Ok(Observation::Migrated) + } Interaction::Replay => { let _ = self.active_mut_tx.take(); self.reopen_from_commitlog()?; @@ -271,6 +383,41 @@ impl EngineTarget { } } +fn expected_step_from_auto_step(step: &AutoMigrateStep<'_>) -> anyhow::Result { + match step { + AutoMigrateStep::AddColumns(_) => Ok(ExpectedStep::AddColumns), + AutoMigrateStep::AddIndex(_) => Ok(ExpectedStep::AddIndex), + AutoMigrateStep::RemoveIndex(_) => Ok(ExpectedStep::RemoveIndex), + AutoMigrateStep::ChangeAccess(_) => Ok(ExpectedStep::ChangeAccess), + AutoMigrateStep::ChangePrimaryKey(_) => Ok(ExpectedStep::ChangePrimaryKey), + AutoMigrateStep::ChangeColumns(_) => Ok(ExpectedStep::ChangeColumns), + AutoMigrateStep::ReschemaEventTable(_) => Ok(ExpectedStep::ReschemaEventTable), + AutoMigrateStep::DisconnectAllUsers => Ok(ExpectedStep::DisconnectAllUsers), + step => anyhow::bail!("engine DST generated unsupported migration step: {step:?}"), + } +} + +fn index_columns(algorithm: &SchemaIndexAlgorithm) -> Vec { + algorithm.columns().iter().map(|col| col.0 as usize).collect() +} + +fn index_algorithm_state(algorithm: &SchemaIndexAlgorithm) -> IndexAlgorithmState { + match algorithm { + SchemaIndexAlgorithm::BTree(_) => IndexAlgorithmState::BTree, + SchemaIndexAlgorithm::Hash(_) => IndexAlgorithmState::Hash, + SchemaIndexAlgorithm::Direct(_) => IndexAlgorithmState::Direct, + _ => IndexAlgorithmState::Unknown, + } +} + +struct DstUpdateLogger; + +impl UpdateLogger for DstUpdateLogger { + fn info(&self, msg: &str) { + tracing::debug!(%msg, "engine DST migration update"); + } +} + impl TargetDriver for EngineTarget { type Observation = Observation; @@ -302,3 +449,169 @@ impl TestSuite for EngineTest { Ok((interactions, target, properties)) } } + +#[cfg(test)] +mod tests { + use spacetimedb_lib::AlgebraicValue; + use spacetimedb_sats::product; + + use super::migrations::{Migration, MigrationOp}; + use super::*; + use crate::schema::{ColumnPlan, IndexAlgorithm, IndexPlan, TablePlan, Type, UniqueConstraintPlan}; + + fn migration_replay_schema() -> SchemaPlan { + SchemaPlan { + tables: vec![TablePlan { + name: "items".into(), + columns: vec![ + ColumnPlan { + name: "id".into(), + ty: Type::U64, + }, + ColumnPlan { + name: "kind".into(), + ty: Type::Sum { variants: 1 }, + }, + ], + primary_key: Some(0), + indexes: vec![IndexPlan { + columns: vec![0], + algorithm: IndexAlgorithm::BTree, + }], + unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], + sequences: vec![], + is_public: true, + is_event: false, + }], + } + } + + fn add_column_replay_schema() -> SchemaPlan { + SchemaPlan { + tables: vec![TablePlan { + name: "items".into(), + columns: vec![ + ColumnPlan { + name: "id".into(), + ty: Type::U64, + }, + ColumnPlan { + name: "score".into(), + ty: Type::U64, + }, + ], + primary_key: Some(0), + indexes: vec![IndexPlan { + columns: vec![0], + algorithm: IndexAlgorithm::BTree, + }], + unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], + sequences: vec![], + is_public: true, + is_event: false, + }], + } + } + + fn change_index_replay_schema() -> SchemaPlan { + SchemaPlan { + tables: vec![TablePlan { + name: "items".into(), + columns: vec![ + ColumnPlan { + name: "id".into(), + ty: Type::U64, + }, + ColumnPlan { + name: "score".into(), + ty: Type::U64, + }, + ], + primary_key: Some(0), + indexes: vec![ + IndexPlan { + columns: vec![0], + algorithm: IndexAlgorithm::BTree, + }, + IndexPlan { + columns: vec![1], + algorithm: IndexAlgorithm::BTree, + }, + ], + unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], + sequences: vec![], + is_public: true, + is_event: false, + }], + } + } + + fn insert_u64_rows(target: &mut EngineTarget) -> anyhow::Result<()> { + target.execute(&Interaction::BeginMutTx)?; + for id in 0..128u64 { + target.execute(&Interaction::Insert { + table: 0, + row: product![id, id * 10], + })?; + } + target.execute(&Interaction::CommitTx)?; + Ok(()) + } + + #[test] + fn engine_dst_smoke_runs_random_workload() -> anyhow::Result<()> { + let mut runtime = SimRuntime::new(0); + runtime.block_on(EngineTest.run(Rng::new(0), 1_000))?; + Ok(()) + } + + #[test] + fn add_column_migration_replays_with_existing_rows() -> anyhow::Result<()> { + let mut target = EngineTarget::init(add_column_replay_schema(), 0)?; + insert_u64_rows(&mut target)?; + + target.execute(&Interaction::Migrate(Migration { + table: 0, + ops: vec![MigrationOp::ChangeAccess, MigrationOp::AddColumn { ty: Type::U64 }], + }))?; + target.execute(&Interaction::Replay)?; + + Ok(()) + } + + #[test] + fn change_index_migration_replays_with_existing_rows() -> anyhow::Result<()> { + let mut target = EngineTarget::init(change_index_replay_schema(), 0)?; + insert_u64_rows(&mut target)?; + + target.execute(&Interaction::Migrate(Migration { + table: 0, + ops: vec![MigrationOp::ChangeAccess, MigrationOp::ChangeIndex { index: 1 }], + }))?; + target.execute(&Interaction::Replay)?; + + Ok(()) + } + + #[test] + fn migration_that_updates_st_table_and_st_column_replays() -> anyhow::Result<()> { + let mut target = EngineTarget::init(migration_replay_schema(), 0)?; + + target.execute(&Interaction::BeginMutTx)?; + for id in 0..128u64 { + target.execute(&Interaction::Insert { + table: 0, + row: product![id, AlgebraicValue::sum(0, AlgebraicValue::U8(1))], + })?; + } + target.execute(&Interaction::CommitTx)?; + + target.execute(&Interaction::Migrate(Migration { + table: 0, + ops: vec![MigrationOp::ChangeAccess, MigrationOp::ChangeColumnType { column: 1 }], + }))?; + target.execute(&Interaction::Replay)?; + + Ok(()) + } +} diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs new file mode 100644 index 00000000000..90015e2e557 --- /dev/null +++ b/crates/dst/src/engine/migrations.rs @@ -0,0 +1,246 @@ +use spacetimedb_lib::AlgebraicValue; +use spacetimedb_runtime::sim::Rng; + +use crate::schema::{ColumnPlan, IndexAlgorithm, SchemaDecisions, SchemaNames, SchemaPlan, TablePlan, Type}; + +const MAX_SUM_VARIANTS: u8 = 32; +const MAX_EVENT_COLUMNS: usize = 32; +const MAX_TABLE_COLUMNS: usize = 32; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Migration { + pub table: usize, + pub ops: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MigrationOp { + ChangeAccess, + AddColumn { ty: Type }, + ChangePrimaryKey { column: Option }, + ChangeIndex { index: usize }, + ChangeColumnType { column: usize }, + ReschemaEventTable, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum ExpectedStep { + AddColumns, + AddIndex, + RemoveIndex, + ChangeAccess, + ChangePrimaryKey, + ChangeColumns, + ReschemaEventTable, + DisconnectAllUsers, +} + +impl Migration { + pub fn choose(schema: &SchemaPlan, rng: &Rng, table_row_count: impl Fn(usize) -> usize) -> Option { + let candidates = Self::candidates(schema, table_row_count); + SchemaDecisions::choose_index(rng, candidates.len()).map(|idx| candidates[idx].clone()) + } + + pub fn candidates(schema: &SchemaPlan, table_row_count: impl Fn(usize) -> usize) -> Vec { + let mut candidates = Vec::new(); + + for (table, table_plan) in schema.tables.iter().enumerate() { + if table_plan.is_event { + if table_row_count(table) == 0 && table_plan.columns.len() < MAX_EVENT_COLUMNS { + candidates.push(Self::new( + table, + [MigrationOp::ChangeAccess, MigrationOp::ReschemaEventTable], + )); + } + continue; + } + + if table_plan.columns.len() < MAX_TABLE_COLUMNS { + candidates.extend( + Type::ALL + .iter() + .copied() + .map(|ty| Self::new(table, [MigrationOp::ChangeAccess, MigrationOp::AddColumn { ty }])), + ); + } + + if let Some(column) = first_widenable_sum_column(table_plan) { + candidates.push(Self::new( + table, + [MigrationOp::ChangeAccess, MigrationOp::ChangeColumnType { column }], + )); + + if table_plan.primary_key.is_some() && table_plan.sequences.is_empty() { + candidates.push(Self::new( + table, + [ + MigrationOp::ChangePrimaryKey { column: None }, + MigrationOp::ChangeColumnType { column }, + ], + )); + } + } + + if let Some(index) = first_changeable_index(table_plan) { + candidates.push(Self::new( + table, + [MigrationOp::ChangeAccess, MigrationOp::ChangeIndex { index }], + )); + } + } + + candidates + } + + pub fn apply_to(&self, schema: &SchemaPlan) -> anyhow::Result { + let mut next = schema.clone(); + let table = next + .tables + .get_mut(self.table) + .ok_or_else(|| anyhow::anyhow!("migration references missing table {}", self.table))?; + + for op in &self.ops { + apply_op(table, *op)?; + } + + Ok(next) + } + + pub fn added_column_defaults(&self) -> Vec { + self.ops + .iter() + .filter_map(|op| match op { + MigrationOp::AddColumn { ty } => Some(ty.default_value()), + _ => None, + }) + .collect() + } + + pub fn expected_steps(&self) -> Vec { + let mut expected = Vec::new(); + + for op in &self.ops { + match op { + MigrationOp::ChangeAccess => expected.push(ExpectedStep::ChangeAccess), + MigrationOp::AddColumn { .. } => { + expected.push(ExpectedStep::AddColumns); + expected.push(ExpectedStep::DisconnectAllUsers); + } + MigrationOp::ChangePrimaryKey { .. } => expected.push(ExpectedStep::ChangePrimaryKey), + MigrationOp::ChangeIndex { .. } => { + expected.push(ExpectedStep::RemoveIndex); + expected.push(ExpectedStep::AddIndex); + } + MigrationOp::ChangeColumnType { .. } => expected.push(ExpectedStep::ChangeColumns), + MigrationOp::ReschemaEventTable => { + expected.push(ExpectedStep::ReschemaEventTable); + expected.push(ExpectedStep::DisconnectAllUsers); + } + } + } + + if expected.contains(&ExpectedStep::AddColumns) { + expected.retain(|step| *step != ExpectedStep::ChangeColumns); + } + + expected.sort(); + expected.dedup(); + expected + } + + fn new(table: usize, ops: impl Into>) -> Self { + Self { table, ops: ops.into() } + } +} + +fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { + match op { + MigrationOp::ChangeAccess => { + table.is_public = !table.is_public; + } + MigrationOp::AddColumn { ty } => { + anyhow::ensure!(!table.is_event, "add-column migration selected event table"); + anyhow::ensure!( + table.columns.len() < MAX_TABLE_COLUMNS, + "table already has the configured maximum number of columns" + ); + table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "added_col"), + ty, + }); + } + MigrationOp::ChangePrimaryKey { column } => { + if let Some(column) = column { + anyhow::ensure!( + column < table.columns.len(), + "primary-key migration references missing column" + ); + } + table.primary_key = column; + } + MigrationOp::ChangeIndex { index } => { + let Some(index_plan) = table.indexes.get_mut(index) else { + anyhow::bail!("index migration references missing index {index}"); + }; + index_plan.algorithm = match index_plan.algorithm { + IndexAlgorithm::BTree => IndexAlgorithm::Hash, + IndexAlgorithm::Hash => IndexAlgorithm::BTree, + }; + } + MigrationOp::ChangeColumnType { column } => { + widen_sum_column(table, Some(column))?; + } + MigrationOp::ReschemaEventTable => { + anyhow::ensure!(table.is_event, "event-table migration selected non-event table"); + anyhow::ensure!( + table.columns.len() < MAX_EVENT_COLUMNS, + "event table already has the configured maximum number of columns" + ); + table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "reschema_payload"), + ty: Type::U64, + }); + } + } + + Ok(()) +} + +fn first_widenable_sum_column(table: &TablePlan) -> Option { + table.columns.iter().position(|column| match column.ty { + Type::Sum { variants } => variants < MAX_SUM_VARIANTS, + _ => false, + }) +} + +fn first_changeable_index(table: &TablePlan) -> Option { + table + .indexes + .iter() + .position(|index| !is_required_index(table, &index.columns)) +} + +fn is_required_index(table: &TablePlan, columns: &[usize]) -> bool { + table.primary_key.is_some_and(|primary_key| columns == [primary_key]) + || table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == columns) +} + +fn widen_sum_column(table: &mut TablePlan, column: Option) -> anyhow::Result<()> { + let column = column.ok_or_else(|| anyhow::anyhow!("sum-widening migration missing column"))?; + let Some(column_plan) = table.columns.get_mut(column) else { + anyhow::bail!("sum-widening migration references missing column {column}"); + }; + + let Type::Sum { variants } = &mut column_plan.ty else { + anyhow::bail!("sum-widening migration selected a non-sum column"); + }; + anyhow::ensure!( + *variants < MAX_SUM_VARIANTS, + "sum column already has the configured maximum number of variants" + ); + *variants += 1; + Ok(()) +} diff --git a/crates/dst/src/engine/model.rs b/crates/dst/src/engine/model.rs index 5d6913212e0..881329baf2d 100644 --- a/crates/dst/src/engine/model.rs +++ b/crates/dst/src/engine/model.rs @@ -1,5 +1,6 @@ use super::workload::{ - normalize_rows, CommitDelta, CountState, InsertOutcome, Interaction, Observation, Row, TableDelta, TableRowCount, + normalize_rows, schema_state_for_plan, CommitDelta, CountState, InsertOutcome, Interaction, Observation, Row, + TableDelta, TableRowCount, }; use crate::schema::SchemaPlan; @@ -151,6 +152,21 @@ impl Model { let delta = self.commit_pending(pending_tx); Observation::Committed { delta } } + Interaction::Migrate(migration) => { + debug_assert!(self.pending_tx.is_none()); + let added_defaults = migration.added_column_defaults(); + self.schema = migration + .apply_to(&self.schema) + .expect("generated migrations must be valid for the model schema"); + if !added_defaults.is_empty() { + for row in &mut self.committed_tables[migration.table].rows { + let mut elements = row.elements.to_vec(); + elements.extend(added_defaults.iter().cloned()); + row.elements = elements.into_boxed_slice(); + } + } + Observation::Migrated + } Interaction::Replay => { self.pending_tx = None; Observation::Replayed { @@ -233,7 +249,10 @@ impl Model { count: self.visible_count(table), }) .collect(); - CountState { row_counts } + CountState { + row_counts, + schema: schema_state_for_plan(&self.schema), + } } } @@ -260,6 +279,7 @@ mod tests { unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], sequences: vec![], is_public: true, + is_event: false, }], } } diff --git a/crates/dst/src/engine/workload.rs b/crates/dst/src/engine/workload.rs index bc9e4fcc36e..9076985a9f0 100644 --- a/crates/dst/src/engine/workload.rs +++ b/crates/dst/src/engine/workload.rs @@ -1,12 +1,13 @@ use std::fmt::{Debug, Error, Formatter}; use spacetimedb_lib::bsatn::to_vec; -use spacetimedb_lib::{AlgebraicValue, ProductValue}; +use spacetimedb_lib::{AlgebraicType, AlgebraicValue, ProductValue}; use spacetimedb_runtime::sim::Rng; use spacetimedb_sats::ArrayValue; +use super::migrations::Migration; use super::model::Model; -use crate::schema::{SchemaPlan, TablePlan, Type}; +use crate::schema::{IndexAlgorithm, SchemaPlan, TablePlan, Type}; pub type Row = ProductValue; @@ -16,6 +17,7 @@ pub enum Interaction { Insert { table: usize, row: Row }, Delete { table: usize, row: Row }, CommitTx, + Migrate(Migration), Replay, } @@ -26,6 +28,7 @@ pub struct InteractionCounts { pub insert: usize, pub delete: usize, pub commit_tx: usize, + pub migrate: usize, pub replay: usize, } @@ -38,6 +41,7 @@ impl InteractionCounts { Interaction::Insert { .. } => self.insert += 1, Interaction::Delete { .. } => self.delete += 1, Interaction::CommitTx => self.commit_tx += 1, + Interaction::Migrate(_) => self.migrate += 1, Interaction::Replay => self.replay += 1, } } @@ -49,6 +53,7 @@ pub enum Observation { Inserted { outcome: InsertOutcome }, Deleted, Committed { delta: CommitDelta }, + Migrated, Replayed { state: CountState }, } @@ -63,6 +68,7 @@ pub struct InteractionWeights { pub insert: u64, pub delete: u64, pub commit_tx: u64, + pub migrate: u64, pub replay: u64, } @@ -71,7 +77,8 @@ impl Default for InteractionWeights { Self { insert: 50, delete: 20, - commit_tx: 29, + commit_tx: 28, + migrate: 1, replay: 1, } } @@ -82,6 +89,7 @@ enum InteractionChoice { Insert, Delete, CommitTx, + Migrate, Replay, } @@ -121,6 +129,10 @@ impl WorkloadGen { let bytes: Vec = (0..len).map(|_| self.rng.next_u64() as u8).collect(); AlgebraicValue::Array(ArrayValue::U8(bytes.into())) } + Type::Sum { variants } => { + let tag = self.rng.index(variants as usize) as u8; + AlgebraicValue::sum(tag, AlgebraicValue::U8(self.rng.next_u64() as u8)) + } } } @@ -153,7 +165,8 @@ impl WorkloadGen { .auto_inc_table_and_column() .map(|(table_idx, _)| table_idx); - (0..self.schema().tables.len()).find(|&table_idx| Some(table_idx) != auto_inc_table) + (0..self.schema().tables.len()) + .find(|&table_idx| Some(table_idx) != auto_inc_table && !self.schema().tables[table_idx].is_event) } pub fn next_interaction(&mut self) -> Interaction { @@ -170,6 +183,10 @@ impl WorkloadGen { if !self.model.in_mut_tx() { return match choice { InteractionChoice::Replay => Interaction::Replay, + InteractionChoice::Migrate => self + .gen_migration() + .map(Interaction::Migrate) + .unwrap_or(Interaction::Replay), // Insert/Delete/CommitTx are not legal outside a mutable tx. // Treat those weighted choices as pressure to start one. @@ -182,6 +199,8 @@ impl WorkloadGen { match choice { InteractionChoice::Replay => Interaction::Replay, + InteractionChoice::Migrate => Interaction::CommitTx, + InteractionChoice::Insert => { let table = self.insert_table_idx(); @@ -215,11 +234,18 @@ impl WorkloadGen { fn pick_interaction_choice(&mut self) -> InteractionChoice { let weights = self.weights; - match self.pick_weighted(&[weights.insert, weights.delete, weights.commit_tx, weights.replay]) { + match self.pick_weighted(&[ + weights.insert, + weights.delete, + weights.commit_tx, + weights.migrate, + weights.replay, + ]) { 0 => InteractionChoice::Insert, 1 => InteractionChoice::Delete, 2 => InteractionChoice::CommitTx, - 3 => InteractionChoice::Replay, + 3 => InteractionChoice::Migrate, + 4 => InteractionChoice::Replay, _ => unreachable!(), } } @@ -246,14 +272,35 @@ impl WorkloadGen { let auto_inc_table_idx = self .schema() .auto_inc_table_and_column() - .map(|(table_idx, _)| table_idx); + .map(|(table_idx, _)| table_idx) + .filter(|&table_idx| !self.schema().tables[table_idx].is_event); + let data_tables = self.data_table_indices(); match auto_inc_table_idx { Some(table_idx) if !self.rng.next_u64().is_multiple_of(3) => table_idx, - _ => self.rng.index(self.schema().tables.len()), + _ => data_tables[self.rng.index(data_tables.len())], } } + fn data_table_indices(&self) -> Vec { + let data_tables: Vec<_> = self + .schema() + .tables + .iter() + .enumerate() + .filter_map(|(table_idx, table)| (!table.is_event).then_some(table_idx)) + .collect(); + assert!( + !data_tables.is_empty(), + "engine DST schema must include a non-event table" + ); + data_tables + } + + fn gen_migration(&self) -> Option { + Migration::choose(self.schema(), &self.rng, |table| self.model.row_count(table)) + } + fn deletable_table_idx(&self) -> Option { self.non_auto_inc_table_idx() .filter(|&table_idx| self.model.row_count(table_idx) > 0) @@ -286,6 +333,7 @@ pub fn normalize_rows(mut rows: Vec) -> Vec { #[derive(Debug, Clone, PartialEq, Eq)] pub struct CountState { pub row_counts: Vec, + pub schema: SchemaState, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -294,6 +342,99 @@ pub struct TableRowCount { pub count: u64, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SchemaState { + pub tables: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableSchemaState { + pub table: usize, + pub name: String, + pub is_public: bool, + pub is_event: bool, + pub primary_key: Option, + pub columns: Vec, + pub indexes: Vec, + pub sequences: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ColumnState { + pub name: String, + pub ty: AlgebraicType, +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct IndexState { + pub columns: Vec, + pub algorithm: IndexAlgorithmState, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum IndexAlgorithmState { + BTree, + Hash, + Direct, + Unknown, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub struct SequenceState { + pub column: usize, +} + +pub fn schema_state_for_plan(schema: &SchemaPlan) -> SchemaState { + SchemaState { + tables: schema + .tables + .iter() + .enumerate() + .map(|(table, table_plan)| { + let mut indexes = table_plan + .indexes + .iter() + .map(|index| IndexState { + columns: index.columns.clone(), + algorithm: match index.algorithm { + IndexAlgorithm::BTree => IndexAlgorithmState::BTree, + IndexAlgorithm::Hash => IndexAlgorithmState::Hash, + }, + }) + .collect::>(); + indexes.sort(); + + let mut sequences = table_plan + .sequences + .iter() + .map(|sequence| SequenceState { + column: sequence.column, + }) + .collect::>(); + sequences.sort(); + + TableSchemaState { + table, + name: table_plan.name.clone(), + is_public: table_plan.is_public, + is_event: table_plan.is_event, + primary_key: table_plan.primary_key, + columns: table_plan + .columns + .iter() + .map(|column| ColumnState { + name: column.name.clone(), + ty: column.ty.to_algebraic(), + }) + .collect(), + indexes, + sequences, + } + }) + .collect(), + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct CommitDelta { pub tables: Vec, diff --git a/crates/dst/src/schema.rs b/crates/dst/src/schema.rs index 641281db3c3..830f9102bcf 100644 --- a/crates/dst/src/schema.rs +++ b/crates/dst/src/schema.rs @@ -2,12 +2,15 @@ use spacetimedb_lib::db::raw_def::v10::*; use spacetimedb_lib::db::raw_def::v9::{RawIndexAlgorithm, TableAccess, TableType}; use spacetimedb_primitives::{ColId, ColList}; use spacetimedb_runtime::sim::Rng; -use spacetimedb_sats::{AlgebraicType, ArrayType, ProductType, ProductTypeElement}; +use spacetimedb_sats::{ + AlgebraicType, AlgebraicValue, ArrayType, ArrayValue, ProductType, ProductTypeElement, SumType, SumTypeVariant, +}; pub fn default_schema(rng: Rng) -> SchemaPlan { let profile = SchemaProfile::default(); let mut plan = SchemaGenerator::new(rng, profile).gen_schema(); plan.ensure_auto_inc_table(); + plan.ensure_engine_migration_coverage(); plan } @@ -18,6 +21,7 @@ pub enum Type { U64, String, Bytes, + Sum { variants: u8 }, } impl Type { @@ -32,6 +36,26 @@ impl Type { Type::Bytes => AlgebraicType::Array(ArrayType { elem_ty: Box::new(AlgebraicType::U8), }), + Type::Sum { variants } => { + debug_assert!(variants > 0); + AlgebraicType::Sum(SumType::new( + (0..variants) + .map(|variant| SumTypeVariant::new_named(AlgebraicType::U8, format!("variant_{variant}"))) + .collect::>() + .into_boxed_slice(), + )) + } + } + } + + pub fn default_value(self) -> AlgebraicValue { + match self { + Type::Bool => AlgebraicValue::Bool(false), + Type::I64 => AlgebraicValue::I64(0), + Type::U64 => AlgebraicValue::U64(0), + Type::String => AlgebraicValue::String("".into()), + Type::Bytes => AlgebraicValue::Array(ArrayValue::U8(Vec::new().into())), + Type::Sum { .. } => AlgebraicValue::sum(0, AlgebraicValue::U8(0)), } } @@ -40,6 +64,110 @@ impl Type { } } +pub struct SchemaDecisions; + +impl SchemaDecisions { + pub fn range(rng: &Rng, (lo, hi): (usize, usize)) -> usize { + if lo >= hi { + return lo; + } + lo + (rng.next_u64() as usize % (hi - lo + 1)) + } + + pub fn index(rng: &Rng, len: usize) -> usize { + rng.index(len) + } + + pub fn choose_index(rng: &Rng, len: usize) -> Option { + (len > 0).then(|| Self::index(rng, len)) + } + + pub fn sample_probability(rng: &Rng, probability: f64) -> bool { + rng.sample_probability(probability) + } + + pub fn gen_type(rng: &Rng) -> Type { + Type::ALL[Self::index(rng, Type::ALL.len())] + } + + pub fn gen_table_name(rng: &Rng, tables: &[TablePlan]) -> String { + loop { + let name = format!("tbl_{}", Self::gen_ident(rng)); + if tables.iter().all(|table| table.name != name) { + return name; + } + } + } + + pub fn gen_column_name(rng: &Rng, seen: &[String]) -> String { + loop { + let name = Self::gen_ident(rng); + if !seen.contains(&name) { + return name; + } + } + } + + fn gen_ident(rng: &Rng) -> String { + const CHARS: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789_"; + const FIRST: &[u8] = b"abcdefghijklmnopqrstuvwxyz_"; + let len = 4 + (rng.next_u64() as usize % 12); + let mut s = String::with_capacity(len); + s.push(FIRST[Self::index(rng, FIRST.len())] as char); + for _ in 1..len { + s.push(CHARS[Self::index(rng, CHARS.len())] as char); + } + s + } +} + +pub struct SchemaNames; + +impl SchemaNames { + pub fn fresh_column_name(table: &TablePlan, base: &str) -> String { + if table.columns.iter().all(|column| column.name != base) { + return base.into(); + } + + for suffix in 0.. { + let candidate = format!("{base}_{suffix}"); + if table.columns.iter().all(|column| column.name != candidate) { + return candidate; + } + } + + unreachable!("unbounded suffix search must find a unique column name") + } + + pub fn fresh_table_name(tables: &[TablePlan], base: &str) -> String { + if tables.iter().all(|table| table.name != base) { + return base.into(); + } + + for suffix in 0.. { + let candidate = format!("{base}_{suffix}"); + if tables.iter().all(|table| table.name != candidate) { + return candidate; + } + } + + unreachable!("unbounded suffix search must find a unique table name") + } + + pub fn index_name(table: &TablePlan, index: &IndexPlan) -> String { + format!( + "{}_{}_idx", + table.name, + index + .columns + .iter() + .map(|&c| table.columns[c].name.as_str()) + .collect::>() + .join("_") + ) + } +} + // Schema plan — the canonical source of truth. // This Schema should be able to translate to valid `RawModuleDefV10`. #[derive(Debug, Clone)] @@ -48,6 +176,11 @@ pub struct SchemaPlan { } impl SchemaPlan { + pub fn ensure_engine_migration_coverage(&mut self) { + self.ensure_widenable_sum_column(); + self.ensure_event_table(); + } + pub fn auto_inc_table_and_column(&self) -> Option<(usize, usize)> { self.tables .iter() @@ -86,6 +219,47 @@ impl SchemaPlan { } table.sequences = vec![SequencePlan::new(0, Type::U64).expect("u64 is integral")]; } + + fn ensure_widenable_sum_column(&mut self) { + if self + .tables + .iter() + .any(|table| !table.is_event && table.columns.iter().any(|column| matches!(column.ty, Type::Sum { .. }))) + { + return; + } + + let table = self + .tables + .iter_mut() + .find(|table| !table.is_event) + .expect("schema must contain at least one non-event table"); + table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "dst_sum"), + ty: Type::Sum { variants: 1 }, + }); + } + + fn ensure_event_table(&mut self) { + if self.tables.iter().any(|table| table.is_event) { + return; + } + + let name = SchemaNames::fresh_table_name(&self.tables, "dst_events"); + self.tables.push(TablePlan { + name, + columns: vec![ColumnPlan { + name: "payload".into(), + ty: Type::U64, + }], + primary_key: None, + indexes: vec![], + unique_constraints: vec![], + sequences: vec![], + is_public: true, + is_event: true, + }); + } } #[derive(Debug, Clone)] @@ -97,6 +271,7 @@ pub struct TablePlan { pub unique_constraints: Vec, pub sequences: Vec, pub is_public: bool, + pub is_event: bool, } #[derive(Debug, Clone)] @@ -165,9 +340,10 @@ fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { .collect(), }; - let mut tbl = builder.build_table_with_new_type(table.name.clone(), product_type, true); + let mut tbl = builder.build_table_with_new_type_for_tests(table.name.clone(), product_type, true); tbl = tbl.with_type(TableType::User); + tbl = tbl.with_event(table.is_event); tbl = tbl.with_access(if table.is_public { TableAccess::Public } else { @@ -193,18 +369,7 @@ fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { IndexAlgorithm::Hash => RawIndexAlgorithm::Hash { columns: col_list }, }; - let source_name = format!( - "{}_{}_idx", - table.name, - index - .columns - .iter() - .map(|&c| table.columns[c].name.as_str()) - .collect::>() - .join("_") - ); - - tbl = tbl.with_index_no_accessor_name(algorithm, source_name); + tbl = tbl.with_index_no_accessor_name(algorithm, SchemaNames::index_name(table, index)); } // Sequences — all of them. @@ -212,6 +377,12 @@ fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { tbl = tbl.with_column_sequence(ColId(seq.column as u16)); } + // AddColumns needs defaults when existing rows are present. Supplying stable + // defaults for all columns lets the engine keep only the newly-added tail. + for (col_id, column) in table.columns.iter().enumerate() { + tbl = tbl.with_default_column_value(ColId(col_id as u16), column.ty.default_value()); + } + tbl.finish(); } @@ -253,60 +424,30 @@ impl SchemaGenerator { Self { rng, profile } } - fn range(&self, (lo, hi): (usize, usize)) -> usize { - if lo >= hi { - return lo; - } - lo + (self.rng.next_u64() as usize % (hi - lo + 1)) - } - - fn gen_type(&self) -> Type { - Type::ALL[self.rng.index(Type::ALL.len())] - } - fn gen_columns(&self) -> Vec { - let n = self.range(self.profile.columns); + let n = SchemaDecisions::range(&self.rng, self.profile.columns); let mut names = Vec::with_capacity(n); let mut seen = Vec::with_capacity(n); for _ in 0..n { - let name = self.gen_column_name(&seen); + let name = SchemaDecisions::gen_column_name(&self.rng, &seen); seen.push(name.clone()); names.push(ColumnPlan { name, - ty: self.gen_type(), + ty: SchemaDecisions::gen_type(&self.rng), }); } names } - fn gen_ident(&self) -> String { - const CHARS: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789_"; - const FIRST: &[u8] = b"abcdefghijklmnopqrstuvwxyz_"; - let len = 4 + (self.rng.next_u64() as usize % 12); - let mut s = String::with_capacity(len); - s.push(FIRST[self.rng.index(FIRST.len())] as char); - for _ in 1..len { - s.push(CHARS[self.rng.index(CHARS.len())] as char); - } - s - } - - fn gen_column_name(&self, seen: &[String]) -> String { - loop { - let name = self.gen_ident(); - if !seen.contains(&name) { - return name; - } - } - } - fn gen_unique_constraints(&self, columns: &[ColumnPlan], pk: &Option) -> Vec { - let n = self.range(self.profile.unique_constraints); + let n = SchemaDecisions::range(&self.rng, self.profile.unique_constraints); let mut seen: Vec> = Vec::new(); let mut result = Vec::new(); for _ in 0..n { - let num_cols = 1 + self.rng.index(columns.len().min(3)); - let mut cols: Vec = (0..num_cols).map(|_| self.rng.index(columns.len())).collect(); + let num_cols = 1 + SchemaDecisions::index(&self.rng, columns.len().min(3)); + let mut cols: Vec = (0..num_cols) + .map(|_| SchemaDecisions::index(&self.rng, columns.len())) + .collect(); cols.sort(); cols.dedup(); if !cols.is_empty() && !seen.contains(&cols) { @@ -355,17 +496,19 @@ impl SchemaGenerator { } // Additional random indexes. - let n = self.range(self.profile.indexes); + let n = SchemaDecisions::range(&self.rng, self.profile.indexes); for _ in 0..n { - let num_cols = 1 + self.rng.index(columns.len().min(3)); - let mut cols: Vec = (0..num_cols).map(|_| self.rng.index(columns.len())).collect(); + let num_cols = 1 + SchemaDecisions::index(&self.rng, columns.len().min(3)); + let mut cols: Vec = (0..num_cols) + .map(|_| SchemaDecisions::index(&self.rng, columns.len())) + .collect(); cols.sort(); cols.dedup(); if cols.is_empty() || seen_cols.contains(&cols) { continue; } seen_cols.push(cols.clone()); - let algorithm = if self.rng.sample_probability(self.profile.btree_prob) { + let algorithm = if SchemaDecisions::sample_probability(&self.rng, self.profile.btree_prob) { IndexAlgorithm::BTree } else { IndexAlgorithm::Hash @@ -379,11 +522,12 @@ impl SchemaGenerator { indexes } - fn gen_table(&self, _table_index: usize) -> TablePlan { + fn gen_table(&self, existing_tables: &[TablePlan]) -> TablePlan { let columns = self.gen_columns(); - let primary_key = if self.rng.sample_probability(self.profile.pk_prob) && !columns.is_empty() { - Some(self.rng.index(columns.len())) + let primary_key = if SchemaDecisions::sample_probability(&self.rng, self.profile.pk_prob) && !columns.is_empty() + { + Some(SchemaDecisions::index(&self.rng, columns.len())) } else { None }; @@ -391,7 +535,9 @@ impl SchemaGenerator { let unique_constraints = self.gen_unique_constraints(&columns, &primary_key); let sequences = if let Some(pk) = primary_key { - if columns[pk].ty.is_integral() && self.rng.sample_probability(self.profile.auto_inc_prob) { + if columns[pk].ty.is_integral() + && SchemaDecisions::sample_probability(&self.rng, self.profile.auto_inc_prob) + { SequencePlan::new(pk, columns[pk].ty).into_iter().collect() } else { vec![] @@ -402,7 +548,7 @@ impl SchemaGenerator { let indexes = self.gen_indexes(&columns, &unique_constraints, &primary_key); - let name = format!("tbl_{}", self.gen_ident()); + let name = SchemaDecisions::gen_table_name(&self.rng, existing_tables); TablePlan { name, @@ -411,13 +557,17 @@ impl SchemaGenerator { indexes, unique_constraints, sequences, - is_public: !self.rng.sample_probability(self.profile.private_prob), + is_public: !SchemaDecisions::sample_probability(&self.rng, self.profile.private_prob), + is_event: false, } } pub fn gen_schema(&self) -> SchemaPlan { - let table_count = self.range(self.profile.table_count); - let tables = (0..table_count).map(|i| self.gen_table(i)).collect(); + let table_count = SchemaDecisions::range(&self.rng, self.profile.table_count); + let mut tables = Vec::with_capacity(table_count); + for _ in 0..table_count { + tables.push(self.gen_table(&tables)); + } SchemaPlan { tables } } } @@ -453,6 +603,7 @@ mod tests { unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], sequences: vec![SequencePlan::new(0, Type::U64).unwrap()], is_public: true, + is_event: false, }], }; From 591f15934eed3f250c4182c8cef0db320f81160e Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Mon, 13 Jul 2026 14:28:45 +0530 Subject: [PATCH 02/10] cleanup --- crates/dst/Cargo.toml | 1 + crates/dst/src/engine.rs | 31 ++- crates/dst/src/engine/generation.rs | 263 ++++++++++++++++++++ crates/dst/src/engine/migrations.rs | 366 ++++++++++++++++++++++++++-- crates/dst/src/engine/model.rs | 85 ++++++- crates/dst/src/engine/workload.rs | 151 +++++------- crates/dst/src/schema.rs | 190 ++++++++++++++- crates/dst/src/traits.rs | 14 +- 8 files changed, 970 insertions(+), 131 deletions(-) create mode 100644 crates/dst/src/engine/generation.rs diff --git a/crates/dst/Cargo.toml b/crates/dst/Cargo.toml index 7d47c9a1b47..0fa65aa24d7 100644 --- a/crates/dst/Cargo.toml +++ b/crates/dst/Cargo.toml @@ -11,6 +11,7 @@ fallocate = ["spacetimedb-commitlog/fallocate"] [dependencies] anyhow.workspace = true clap.workspace = true +futures.workspace = true spacetimedb-datastore.workspace = true spacetimedb-commitlog.workspace = true spacetimedb-durability.workspace = true diff --git a/crates/dst/src/engine.rs b/crates/dst/src/engine.rs index 006c87c6b57..ac1b02b1087 100644 --- a/crates/dst/src/engine.rs +++ b/crates/dst/src/engine.rs @@ -18,6 +18,7 @@ use spacetimedb_schema::def::{IndexAlgorithm as SchemaIndexAlgorithm, ModuleDef} use spacetimedb_schema::schema::{Schema, TableSchema}; use spacetimedb_table::page_pool::PagePool; +mod generation; mod migrations; mod model; mod properties; @@ -26,7 +27,8 @@ mod workload; use self::migrations::{ExpectedStep, Migration}; use self::workload::{ normalize_rows, row_to_bytes, ColumnState, CommitDelta, CountState, IndexAlgorithmState, IndexState, InsertOutcome, - Interaction, Observation, SchemaState, SequenceState, TableDelta, TableRowCount, TableSchemaState, + Interaction, Observation, SchemaState, SequenceState, TableDelta, TableRowCount, TableRows, TableSchemaState, + UniqueConstraintState, }; use crate::engine::model::Model; @@ -149,17 +151,20 @@ impl EngineTarget { .ok_or_else(|| anyhow::anyhow!("database is not open"))?; let tx = db.begin_tx(Workload::Internal); let mut row_counts = Vec::with_capacity(self.table_ids.len()); + let mut table_rows = Vec::with_capacity(self.table_ids.len()); let mut schema_tables = Vec::with_capacity(self.table_ids.len()); for (table, table_id) in self.table_ids.iter().enumerate() { - let count = match db.iter(&tx, *table_id) { - Ok(iter) => iter.count() as u64, + let rows = match db.iter(&tx, *table_id) { + Ok(iter) => normalize_rows(iter.map(|row| row.to_product_value()).collect()), Err(err) => { let _ = db.release_tx(tx); return Err(err.into()); } }; + let count = rows.len() as u64; row_counts.push(TableRowCount { table, count }); + table_rows.push(TableRows { table, rows }); let schema = match db.schema_for_table(&tx, *table_id) { Ok(schema) => schema, @@ -178,6 +183,17 @@ impl EngineTarget { .collect::>(); indexes.sort(); + let mut unique_constraints = schema + .constraints + .iter() + .filter_map(|constraint| { + constraint.data.unique_columns().map(|columns| UniqueConstraintState { + columns: columns.iter().map(|col| col.0 as usize).collect(), + }) + }) + .collect::>(); + unique_constraints.sort(); + let mut sequences = schema .sequences .iter() @@ -202,6 +218,7 @@ impl EngineTarget { }) .collect(), indexes, + unique_constraints, sequences, }); } @@ -209,6 +226,7 @@ impl EngineTarget { let _ = db.release_tx(tx); Ok(CountState { row_counts, + table_rows, schema: SchemaState { tables: schema_tables }, }) } @@ -385,9 +403,15 @@ impl EngineTarget { fn expected_step_from_auto_step(step: &AutoMigrateStep<'_>) -> anyhow::Result { match step { + AutoMigrateStep::AddTable(_) => Ok(ExpectedStep::AddTable), + AutoMigrateStep::RemoveTable(_) => Ok(ExpectedStep::RemoveTable), AutoMigrateStep::AddColumns(_) => Ok(ExpectedStep::AddColumns), AutoMigrateStep::AddIndex(_) => Ok(ExpectedStep::AddIndex), AutoMigrateStep::RemoveIndex(_) => Ok(ExpectedStep::RemoveIndex), + AutoMigrateStep::AddSequence(_) => Ok(ExpectedStep::AddSequence), + AutoMigrateStep::RemoveSequence(_) => Ok(ExpectedStep::RemoveSequence), + AutoMigrateStep::AddConstraint(_) => Ok(ExpectedStep::AddConstraint), + AutoMigrateStep::RemoveConstraint(_) => Ok(ExpectedStep::RemoveConstraint), AutoMigrateStep::ChangeAccess(_) => Ok(ExpectedStep::ChangeAccess), AutoMigrateStep::ChangePrimaryKey(_) => Ok(ExpectedStep::ChangePrimaryKey), AutoMigrateStep::ChangeColumns(_) => Ok(ExpectedStep::ChangeColumns), @@ -453,6 +477,7 @@ impl TestSuite for EngineTest { #[cfg(test)] mod tests { use spacetimedb_lib::AlgebraicValue; + use spacetimedb_runtime::sim::Runtime as SimRuntime; use spacetimedb_sats::product; use super::migrations::{Migration, MigrationOp}; diff --git a/crates/dst/src/engine/generation.rs b/crates/dst/src/engine/generation.rs new file mode 100644 index 00000000000..9698e13b1df --- /dev/null +++ b/crates/dst/src/engine/generation.rs @@ -0,0 +1,263 @@ +use spacetimedb_lib::{AlgebraicValue, ProductValue}; +use spacetimedb_runtime::sim::Rng; +use spacetimedb_sats::ArrayValue; + +use super::migrations::Migration; +use super::model::{ColumnDomain, Model}; +use super::workload::Row; +use crate::schema::{SchemaDecisions, Type}; + +pub(crate) struct GenCtx<'a> { + rng: &'a Rng, + model: &'a Model, +} + +impl<'a> GenCtx<'a> { + pub(crate) fn new(rng: &'a Rng, model: &'a Model) -> Self { + Self { rng, model } + } + + pub(crate) fn gen_insert_row(&self, table: usize) -> Row { + ValueGen::new(self.rng, self.model).gen_insert_row(table) + } + + pub(crate) fn gen_migration(&self) -> Option { + MigrationGen::new(self.rng, self.model).choose() + } +} + +pub(crate) fn pick_weighted(rng: &Rng, weights: &[u64]) -> usize { + let total: u64 = weights.iter().sum(); + + assert!(total > 0, "at least one interaction weight must be non-zero"); + + let mut selected = rng.next_u64() % total; + + for (idx, weight) in weights.iter().copied().enumerate() { + if selected < weight { + return idx; + } + + selected -= weight; + } + + unreachable!("selected value is always inside total weight") +} + +struct ValueGen<'a> { + rng: &'a Rng, + model: &'a Model, +} + +impl<'a> ValueGen<'a> { + fn new(rng: &'a Rng, model: &'a Model) -> Self { + Self { rng, model } + } + + fn gen_insert_row(&self, table: usize) -> Row { + self.model.schema().tables[table] + .columns + .iter() + .enumerate() + .map(|(column, _)| { + let domain = self.model.column_domain(table, column); + self.gen_insert_value(&domain) + }) + .collect::() + } + + fn gen_insert_value(&self, domain: &ColumnDomain) -> AlgebraicValue { + if domain.sequenced { + return sequence_placeholder(domain.ty); + } + + if domain.unique { + return self.gen_fresh_value(domain); + } + + self.gen_value(domain) + } + + fn gen_value(&self, domain: &ColumnDomain) -> AlgebraicValue { + match pick_weighted(self.rng, &[45, 15, 15, 10, 10, 5]) { + 0 => self.gen_random_value(domain.ty), + 1 => self.gen_small_value(domain.ty), + 2 => self.gen_edge_value(domain.ty), + 3 => self + .near_existing_value(domain) + .unwrap_or_else(|| self.gen_random_value(domain.ty)), + 4 => self + .existing_value(domain) + .unwrap_or_else(|| self.gen_random_value(domain.ty)), + 5 => self.gen_weird_value(domain.ty), + _ => unreachable!(), + } + } + + fn gen_fresh_value(&self, domain: &ColumnDomain) -> AlgebraicValue { + for _ in 0..32 { + let value = self.gen_fresh_candidate(domain.ty); + if !domain.values.contains(&value) { + return value; + } + } + + self.gen_counter_value(domain.ty, self.rng.next_u64()) + } + + fn gen_fresh_candidate(&self, ty: Type) -> AlgebraicValue { + match pick_weighted(self.rng, &[50, 20, 20, 10]) { + 0 => self.gen_random_value(ty), + 1 => self.gen_small_value(ty), + 2 => self.gen_edge_value(ty), + 3 => self.gen_weird_value(ty), + _ => unreachable!(), + } + } + + fn existing_value(&self, domain: &ColumnDomain) -> Option { + (!domain.values.is_empty()).then(|| domain.values[self.rng.index(domain.values.len())].clone()) + } + + fn near_existing_value(&self, domain: &ColumnDomain) -> Option { + let existing = self.existing_value(domain)?; + Some(match (domain.ty, existing) { + (Type::Bool, AlgebraicValue::Bool(value)) => AlgebraicValue::Bool(!value), + (Type::I64, AlgebraicValue::I64(value)) => AlgebraicValue::I64(value.saturating_add(1)), + (Type::U64, AlgebraicValue::U64(value)) => AlgebraicValue::U64(value.saturating_add(1)), + (Type::String, AlgebraicValue::String(value)) => { + let mut value = value.to_string(); + value.push('a'); + AlgebraicValue::String(value.into()) + } + (Type::Bytes, AlgebraicValue::Array(ArrayValue::U8(value))) => { + let mut value = value.to_vec(); + value.push(0); + AlgebraicValue::Array(ArrayValue::U8(value.into())) + } + (Type::Sum { .. }, value) => value, + (_, value) => value, + }) + } + + fn gen_random_value(&self, ty: Type) -> AlgebraicValue { + match ty { + Type::Bool => AlgebraicValue::Bool(self.rng.next_u64().is_multiple_of(2)), + Type::I64 => AlgebraicValue::I64(self.rng.next_u64() as i64), + Type::U64 => AlgebraicValue::U64(self.rng.next_u64()), + Type::String => AlgebraicValue::String(format!("v_{}", self.rng.next_u64()).into()), + Type::Bytes => { + let len = (self.rng.next_u64() % 16) as usize; + let bytes: Vec = (0..len).map(|_| self.rng.next_u64() as u8).collect(); + AlgebraicValue::Array(ArrayValue::U8(bytes.into())) + } + Type::Sum { variants } => { + let tag = self.rng.index(variants as usize) as u8; + AlgebraicValue::sum(tag, AlgebraicValue::U8(self.rng.next_u64() as u8)) + } + } + } + + fn gen_small_value(&self, ty: Type) -> AlgebraicValue { + match ty { + Type::Bool => AlgebraicValue::Bool(false), + Type::I64 => { + const VALUES: &[i64] = &[-1, 0, 1, 2, 3]; + AlgebraicValue::I64(VALUES[self.rng.index(VALUES.len())]) + } + Type::U64 => { + const VALUES: &[u64] = &[0, 1, 2, 3]; + AlgebraicValue::U64(VALUES[self.rng.index(VALUES.len())]) + } + Type::String => { + const VALUES: &[&str] = &["", "a", "aa", "ab", "b", "v_0"]; + AlgebraicValue::String(VALUES[self.rng.index(VALUES.len())].into()) + } + Type::Bytes => { + const VALUES: &[&[u8]] = &[&[], &[0], &[1], &[0, 255]]; + AlgebraicValue::Array(ArrayValue::U8(VALUES[self.rng.index(VALUES.len())].to_vec().into())) + } + Type::Sum { variants } => { + let tag = if variants <= 1 { + 0 + } else { + self.rng.index(variants as usize) as u8 + }; + AlgebraicValue::sum(tag, AlgebraicValue::U8(0)) + } + } + } + + fn gen_edge_value(&self, ty: Type) -> AlgebraicValue { + match ty { + Type::Bool => AlgebraicValue::Bool(true), + Type::I64 => { + const VALUES: &[i64] = &[i64::MIN, i64::MIN + 1, i64::MAX - 1, i64::MAX]; + AlgebraicValue::I64(VALUES[self.rng.index(VALUES.len())]) + } + Type::U64 => { + const VALUES: &[u64] = &[0, 1, u64::MAX - 1, u64::MAX]; + AlgebraicValue::U64(VALUES[self.rng.index(VALUES.len())]) + } + Type::String => AlgebraicValue::String("x".repeat(128).into()), + Type::Bytes => AlgebraicValue::Array(ArrayValue::U8(vec![255; 32].into())), + Type::Sum { variants } => AlgebraicValue::sum(variants.saturating_sub(1), AlgebraicValue::U8(u8::MAX)), + } + } + + fn gen_weird_value(&self, ty: Type) -> AlgebraicValue { + match ty { + Type::String => { + const VALUES: &[&str] = &["quote'", "double\"quote", "back\\slash", "line\nbreak", "\0"]; + AlgebraicValue::String(VALUES[self.rng.index(VALUES.len())].into()) + } + Type::Bytes => { + const VALUES: &[&[u8]] = &[&[0], &[0, 0, 0], &[255], &[0, 255, 0, 255]]; + AlgebraicValue::Array(ArrayValue::U8(VALUES[self.rng.index(VALUES.len())].to_vec().into())) + } + _ => self.gen_edge_value(ty), + } + } + + fn gen_counter_value(&self, ty: Type, counter: u64) -> AlgebraicValue { + match ty { + Type::Bool => AlgebraicValue::Bool(counter.is_multiple_of(2)), + Type::I64 => AlgebraicValue::I64(counter as i64), + Type::U64 => AlgebraicValue::U64(counter), + Type::String => AlgebraicValue::String(format!("fresh_{counter}").into()), + Type::Bytes => AlgebraicValue::Array(ArrayValue::U8(counter.to_le_bytes().to_vec().into())), + Type::Sum { variants } => AlgebraicValue::sum( + if variants == 0 { + 0 + } else { + (counter % variants as u64) as u8 + }, + AlgebraicValue::U8(counter as u8), + ), + } + } +} + +fn sequence_placeholder(ty: Type) -> AlgebraicValue { + match ty { + Type::I64 => AlgebraicValue::I64(0), + Type::U64 => AlgebraicValue::U64(0), + _ => unreachable!("sequence columns are integral"), + } +} + +struct MigrationGen<'a> { + rng: &'a Rng, + model: &'a Model, +} + +impl<'a> MigrationGen<'a> { + fn new(rng: &'a Rng, model: &'a Model) -> Self { + Self { rng, model } + } + + fn choose(&self) -> Option { + let candidates = Migration::candidates(self.model); + SchemaDecisions::choose_index(self.rng, candidates.len()).map(|idx| candidates[idx].clone()) + } +} diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs index 90015e2e557..b21e9f79645 100644 --- a/crates/dst/src/engine/migrations.rs +++ b/crates/dst/src/engine/migrations.rs @@ -1,11 +1,13 @@ +use super::model::{ColumnDomain, Model}; +use crate::schema::{ + ColumnPlan, IndexAlgorithm, IndexPlan, SchemaNames, SchemaPlan, SequencePlan, TablePlan, Type, UniqueConstraintPlan, +}; use spacetimedb_lib::AlgebraicValue; -use spacetimedb_runtime::sim::Rng; - -use crate::schema::{ColumnPlan, IndexAlgorithm, SchemaDecisions, SchemaNames, SchemaPlan, TablePlan, Type}; const MAX_SUM_VARIANTS: u8 = 32; const MAX_EVENT_COLUMNS: usize = 32; const MAX_TABLE_COLUMNS: usize = 32; +const MAX_TABLES: usize = 128; #[derive(Debug, Clone, PartialEq, Eq)] pub struct Migration { @@ -13,21 +15,58 @@ pub struct Migration { pub ops: Vec, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub enum MigrationOp { + AddTable { + is_event: bool, + }, + RemoveTable, ChangeAccess, - AddColumn { ty: Type }, - ChangePrimaryKey { column: Option }, - ChangeIndex { index: usize }, - ChangeColumnType { column: usize }, + AddColumn { + ty: Type, + }, + AddIndex { + columns: Vec, + algorithm: IndexAlgorithm, + }, + RemoveIndex { + index: usize, + }, + AddSequence { + sequence: SequencePlan, + }, + RemoveSequence { + sequence: usize, + }, + AddUniqueConstraint { + columns: Vec, + }, + RemoveUniqueConstraint { + constraint: usize, + }, + ChangePrimaryKey { + column: Option, + }, + ChangeIndex { + index: usize, + }, + ChangeColumnType { + column: usize, + }, ReschemaEventTable, } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum ExpectedStep { + AddTable, + RemoveTable, AddColumns, AddIndex, RemoveIndex, + AddSequence, + RemoveSequence, + AddConstraint, + RemoveConstraint, ChangeAccess, ChangePrimaryKey, ChangeColumns, @@ -36,17 +75,32 @@ pub enum ExpectedStep { } impl Migration { - pub fn choose(schema: &SchemaPlan, rng: &Rng, table_row_count: impl Fn(usize) -> usize) -> Option { - let candidates = Self::candidates(schema, table_row_count); - SchemaDecisions::choose_index(rng, candidates.len()).map(|idx| candidates[idx].clone()) - } - - pub fn candidates(schema: &SchemaPlan, table_row_count: impl Fn(usize) -> usize) -> Vec { + pub(crate) fn candidates(model: &Model) -> Vec { + let schema = model.schema(); let mut candidates = Vec::new(); + let non_event_tables = schema.tables.iter().filter(|table| !table.is_event).count(); + + if schema.tables.len() < MAX_TABLES { + candidates.push(Self::new( + schema.tables.len(), + [MigrationOp::AddTable { is_event: false }], + )); + candidates.push(Self::new( + schema.tables.len(), + [MigrationOp::AddTable { is_event: true }], + )); + } for (table, table_plan) in schema.tables.iter().enumerate() { + let row_count = model.row_count(table); + let pristine = row_count == 0 && !model.ever_inserted(table); + + if (table_plan.is_event && row_count == 0) || (!table_plan.is_event && pristine && non_event_tables > 1) { + candidates.push(Self::new(table, [MigrationOp::RemoveTable])); + } + if table_plan.is_event { - if table_row_count(table) == 0 && table_plan.columns.len() < MAX_EVENT_COLUMNS { + if row_count == 0 && table_plan.columns.len() < MAX_EVENT_COLUMNS { candidates.push(Self::new( table, [MigrationOp::ChangeAccess, MigrationOp::ReschemaEventTable], @@ -64,6 +118,52 @@ impl Migration { ); } + if pristine { + if let Some(sequence) = first_addable_sequence(table_plan) { + candidates.push(Self::new(table, [MigrationOp::AddSequence { sequence }])); + } + + if let Some(columns) = first_addable_unique_constraint_columns(table_plan) { + let mut ops = Vec::new(); + if !has_index(table_plan, &columns) { + ops.push(MigrationOp::AddIndex { + columns: columns.clone(), + algorithm: IndexAlgorithm::BTree, + }); + } + ops.push(MigrationOp::AddUniqueConstraint { columns }); + candidates.push(Self::new(table, ops)); + } + } + + if row_count > 0 { + if let Some(sequence) = + first_addable_sequence_boundary_probe(table_plan, |column| model.column_domain(table, column)) + { + candidates.push(Self::new(table, [MigrationOp::AddSequence { sequence }])); + } + } + + if let Some(sequence) = first_removable_sequence(table_plan) { + candidates.push(Self::new(table, [MigrationOp::RemoveSequence { sequence }])); + } + + if let Some(constraint) = first_removable_unique_constraint(table_plan) { + candidates.push(Self::new(table, [MigrationOp::RemoveUniqueConstraint { constraint }])); + } + + if let Some((columns, algorithm)) = first_addable_index(table_plan) { + candidates.push(Self::new(table, [MigrationOp::AddIndex { columns, algorithm }])); + } + + if let Some(index) = first_changeable_index(table_plan) { + candidates.push(Self::new( + table, + [MigrationOp::ChangeAccess, MigrationOp::ChangeIndex { index }], + )); + candidates.push(Self::new(table, [MigrationOp::RemoveIndex { index }])); + } + if let Some(column) = first_widenable_sum_column(table_plan) { candidates.push(Self::new( table, @@ -80,13 +180,6 @@ impl Migration { )); } } - - if let Some(index) = first_changeable_index(table_plan) { - candidates.push(Self::new( - table, - [MigrationOp::ChangeAccess, MigrationOp::ChangeIndex { index }], - )); - } } candidates @@ -94,18 +187,41 @@ impl Migration { pub fn apply_to(&self, schema: &SchemaPlan) -> anyhow::Result { let mut next = schema.clone(); + if let Some(is_event) = self.adds_table() { + next.tables.push(new_table(&next.tables, is_event)); + return Ok(next); + } + if self.removes_table() { + if self.table >= next.tables.len() { + anyhow::bail!("remove-table migration references missing table {}", self.table); + } + next.tables.remove(self.table); + return Ok(next); + } + let table = next .tables .get_mut(self.table) .ok_or_else(|| anyhow::anyhow!("migration references missing table {}", self.table))?; for op in &self.ops { - apply_op(table, *op)?; + apply_op(table, op.clone())?; } Ok(next) } + pub fn adds_table(&self) -> Option { + match self.ops.as_slice() { + [MigrationOp::AddTable { is_event }] => Some(*is_event), + _ => None, + } + } + + pub fn removes_table(&self) -> bool { + matches!(self.ops.as_slice(), [MigrationOp::RemoveTable]) + } + pub fn added_column_defaults(&self) -> Vec { self.ops .iter() @@ -121,11 +237,22 @@ impl Migration { for op in &self.ops { match op { + MigrationOp::AddTable { .. } => expected.push(ExpectedStep::AddTable), + MigrationOp::RemoveTable => { + expected.push(ExpectedStep::RemoveTable); + expected.push(ExpectedStep::DisconnectAllUsers); + } MigrationOp::ChangeAccess => expected.push(ExpectedStep::ChangeAccess), MigrationOp::AddColumn { .. } => { expected.push(ExpectedStep::AddColumns); expected.push(ExpectedStep::DisconnectAllUsers); } + MigrationOp::AddIndex { .. } => expected.push(ExpectedStep::AddIndex), + MigrationOp::RemoveIndex { .. } => expected.push(ExpectedStep::RemoveIndex), + MigrationOp::AddSequence { .. } => expected.push(ExpectedStep::AddSequence), + MigrationOp::RemoveSequence { .. } => expected.push(ExpectedStep::RemoveSequence), + MigrationOp::AddUniqueConstraint { .. } => expected.push(ExpectedStep::AddConstraint), + MigrationOp::RemoveUniqueConstraint { .. } => expected.push(ExpectedStep::RemoveConstraint), MigrationOp::ChangePrimaryKey { .. } => expected.push(ExpectedStep::ChangePrimaryKey), MigrationOp::ChangeIndex { .. } => { expected.push(ExpectedStep::RemoveIndex); @@ -155,6 +282,12 @@ impl Migration { fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { match op { + MigrationOp::AddTable { .. } => { + anyhow::bail!("add-table migration should be applied by SchemaPlan::push before table ops"); + } + MigrationOp::RemoveTable => { + anyhow::bail!("remove-table migration should be applied by SchemaPlan::remove before table ops"); + } MigrationOp::ChangeAccess => { table.is_public = !table.is_public; } @@ -169,6 +302,73 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { ty, }); } + MigrationOp::AddIndex { columns, algorithm } => { + ensure_columns_exist(table, &columns)?; + anyhow::ensure!( + !has_index(table, &columns), + "add-index migration selected existing index columns" + ); + table.indexes.push(IndexPlan { columns, algorithm }); + } + MigrationOp::RemoveIndex { index } => { + let Some(index_plan) = table.indexes.get(index) else { + anyhow::bail!("remove-index migration references missing index {index}"); + }; + anyhow::ensure!( + !is_required_index(table, &index_plan.columns), + "remove-index migration selected a required index" + ); + table.indexes.remove(index); + } + MigrationOp::AddSequence { sequence } => { + let column = sequence.column; + let Some(column_plan) = table.columns.get(column) else { + anyhow::bail!("add-sequence migration references missing column {column}"); + }; + anyhow::ensure!( + column_plan.ty.is_integral(), + "add-sequence migration selected non-integral column" + ); + anyhow::ensure!( + table.sequences.iter().all(|existing| existing.column != column), + "add-sequence migration selected an already sequenced column" + ); + table.sequences.push(sequence); + } + MigrationOp::RemoveSequence { sequence } => { + anyhow::ensure!( + sequence < table.sequences.len(), + "remove-sequence migration references missing sequence" + ); + table.sequences.remove(sequence); + } + MigrationOp::AddUniqueConstraint { columns } => { + ensure_columns_exist(table, &columns)?; + anyhow::ensure!( + !table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == columns), + "add-constraint migration selected existing constraint columns" + ); + anyhow::ensure!( + has_index(table, &columns), + "add-constraint migration requires a matching index" + ); + table.unique_constraints.push(UniqueConstraintPlan { columns }); + } + MigrationOp::RemoveUniqueConstraint { constraint } => { + let Some(constraint_plan) = table.unique_constraints.get(constraint) else { + anyhow::bail!("remove-constraint migration references missing constraint {constraint}"); + }; + anyhow::ensure!( + !table + .primary_key + .is_some_and(|primary_key| constraint_plan.columns == [primary_key]), + "remove-constraint migration selected primary-key constraint" + ); + table.unique_constraints.remove(constraint); + } MigrationOp::ChangePrimaryKey { column } => { if let Some(column) = column { anyhow::ensure!( @@ -206,6 +406,57 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { Ok(()) } +fn new_table(tables: &[TablePlan], is_event: bool) -> TablePlan { + if is_event { + return TablePlan { + name: SchemaNames::fresh_table_name(tables, "added_events"), + columns: vec![ColumnPlan { + name: "payload".into(), + ty: Type::U64, + }], + primary_key: None, + indexes: vec![], + unique_constraints: vec![], + sequences: vec![], + is_public: true, + is_event: true, + }; + } + + TablePlan { + name: SchemaNames::fresh_table_name(tables, "added_table"), + columns: vec![ + ColumnPlan { + name: "id".into(), + ty: Type::U64, + }, + ColumnPlan { + name: "value".into(), + ty: Type::U64, + }, + ColumnPlan { + name: "kind".into(), + ty: Type::U64, + }, + ], + primary_key: Some(0), + indexes: vec![ + IndexPlan { + columns: vec![0], + algorithm: IndexAlgorithm::BTree, + }, + IndexPlan { + columns: vec![1], + algorithm: IndexAlgorithm::Hash, + }, + ], + unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], + sequences: vec![SequencePlan::new(0, Type::U64).expect("u64 is integral")], + is_public: true, + is_event: false, + } +} + fn first_widenable_sum_column(table: &TablePlan) -> Option { table.columns.iter().position(|column| match column.ty { Type::Sum { variants } => variants < MAX_SUM_VARIANTS, @@ -213,6 +464,13 @@ fn first_widenable_sum_column(table: &TablePlan) -> Option { }) } +fn first_addable_index(table: &TablePlan) -> Option<(Vec, IndexAlgorithm)> { + (0..table.columns.len()).find_map(|column| { + let columns = vec![column]; + (!has_index(table, &columns)).then_some((columns, IndexAlgorithm::BTree)) + }) +} + fn first_changeable_index(table: &TablePlan) -> Option { table .indexes @@ -220,6 +478,59 @@ fn first_changeable_index(table: &TablePlan) -> Option { .position(|index| !is_required_index(table, &index.columns)) } +fn first_addable_sequence(table: &TablePlan) -> Option { + table.columns.iter().enumerate().find_map(|(column, column_plan)| { + (column_plan.ty.is_integral() && table.sequences.iter().all(|sequence| sequence.column != column)) + .then(|| SequencePlan::new(column, column_plan.ty).expect("column type checked above")) + }) +} + +fn first_addable_sequence_boundary_probe( + table: &TablePlan, + column_domain: impl Fn(usize) -> ColumnDomain, +) -> Option { + table.columns.iter().enumerate().find_map(|(column, column_plan)| { + let domain = column_domain(column); + if column_plan.ty != Type::U64 + || domain.sequenced + || !domain.single_column_indexed + || !domain.single_column_unique + { + return None; + } + + let max_value = domain.positive_i128_value_above(2)?; + SequencePlan::with_existing_value_as_max(column, column_plan.ty, max_value) + }) +} + +fn first_removable_sequence(table: &TablePlan) -> Option { + (!table.sequences.is_empty()).then_some(0) +} + +fn first_addable_unique_constraint_columns(table: &TablePlan) -> Option> { + (0..table.columns.len()).find_map(|column| { + let columns = vec![column]; + (!table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == columns)) + .then_some(columns) + }) +} + +fn first_removable_unique_constraint(table: &TablePlan) -> Option { + table.unique_constraints.iter().position(|constraint| { + !table + .primary_key + .is_some_and(|primary_key| constraint.columns == [primary_key]) + }) +} + +fn has_index(table: &TablePlan, columns: &[usize]) -> bool { + table.indexes.iter().any(|index| index.columns == columns) +} + fn is_required_index(table: &TablePlan, columns: &[usize]) -> bool { table.primary_key.is_some_and(|primary_key| columns == [primary_key]) || table @@ -228,6 +539,15 @@ fn is_required_index(table: &TablePlan, columns: &[usize]) -> bool { .any(|constraint| constraint.columns == columns) } +fn ensure_columns_exist(table: &TablePlan, columns: &[usize]) -> anyhow::Result<()> { + anyhow::ensure!(!columns.is_empty(), "migration selected empty column list"); + anyhow::ensure!( + columns.iter().all(|&column| column < table.columns.len()), + "migration references missing column" + ); + Ok(()) +} + fn widen_sum_column(table: &mut TablePlan, column: Option) -> anyhow::Result<()> { let column = column.ok_or_else(|| anyhow::anyhow!("sum-widening migration missing column"))?; let Some(column_plan) = table.columns.get_mut(column) else { diff --git a/crates/dst/src/engine/model.rs b/crates/dst/src/engine/model.rs index 881329baf2d..748170da36f 100644 --- a/crates/dst/src/engine/model.rs +++ b/crates/dst/src/engine/model.rs @@ -1,8 +1,10 @@ +use spacetimedb_lib::AlgebraicValue; + use super::workload::{ normalize_rows, schema_state_for_plan, CommitDelta, CountState, InsertOutcome, Interaction, Observation, Row, - TableDelta, TableRowCount, + TableDelta, TableRowCount, TableRows, }; -use crate::schema::SchemaPlan; +use crate::schema::{SchemaPlan, Type}; #[derive(Debug)] pub struct Model { @@ -14,6 +16,7 @@ pub struct Model { #[derive(Debug)] struct TableState { rows: Vec, + ever_inserted: bool, } #[derive(Debug)] @@ -21,6 +24,30 @@ struct PendingTx { tables: Vec, } +#[derive(Debug, Clone)] +pub(crate) struct ColumnDomain { + pub(crate) ty: Type, + pub(crate) values: Vec, + pub(crate) unique: bool, + pub(crate) single_column_unique: bool, + pub(crate) single_column_indexed: bool, + pub(crate) sequenced: bool, +} + +impl ColumnDomain { + pub(crate) fn integral_values(&self) -> impl Iterator + '_ { + self.values.iter().filter_map(|value| match value { + AlgebraicValue::U64(value) => Some(*value as i128), + AlgebraicValue::I64(value) => Some(*value as i128), + _ => None, + }) + } + + pub(crate) fn positive_i128_value_above(&self, min_exclusive: i128) -> Option { + self.integral_values().find(|value| *value > min_exclusive) + } +} + // Keep mutable transactions as an overlay: committed rows stay shared, while // pending tables record only new rows and delete markers. #[derive(Debug, Default)] @@ -45,7 +72,14 @@ impl PendingTx { impl Model { pub fn new(schema: SchemaPlan) -> Self { - let committed_tables = schema.tables.iter().map(|_| TableState { rows: vec![] }).collect(); + let committed_tables = schema + .tables + .iter() + .map(|_| TableState { + rows: vec![], + ever_inserted: false, + }) + .collect(); Self { schema, committed_tables, @@ -115,6 +149,7 @@ impl Model { } Interaction::Insert { table, row } => { debug_assert!(self.pending_tx.is_some()); + self.committed_tables[*table].ever_inserted = true; // Properties feed the target-returned row here, so sequence-generated // values become part of the oracle before commit/replay checks run. if self.any_visible_row(*table, |visible_row| visible_row == row) { @@ -154,11 +189,20 @@ impl Model { } Interaction::Migrate(migration) => { debug_assert!(self.pending_tx.is_none()); + let adds_table = migration.adds_table().is_some(); + let removes_table = migration.removes_table(); let added_defaults = migration.added_column_defaults(); self.schema = migration .apply_to(&self.schema) .expect("generated migrations must be valid for the model schema"); - if !added_defaults.is_empty() { + if adds_table { + self.committed_tables.push(TableState { + rows: vec![], + ever_inserted: false, + }); + } else if removes_table { + self.committed_tables.remove(migration.table); + } else if !added_defaults.is_empty() { for row in &mut self.committed_tables[migration.table].rows { let mut elements = row.elements.to_vec(); elements.extend(added_defaults.iter().cloned()); @@ -233,6 +277,31 @@ impl Model { self.visible_count(table) as usize } + pub fn ever_inserted(&self, table: usize) -> bool { + self.committed_tables[table].ever_inserted + } + + pub(crate) fn column_domain(&self, table: usize, column: usize) -> ColumnDomain { + let table_plan = &self.schema.tables[table]; + ColumnDomain { + ty: table_plan.columns[column].ty, + values: self + .visible_rows(table) + .map(|row| row.elements[column].clone()) + .collect(), + unique: table_plan + .unique_constraints + .iter() + .any(|constraint| constraint.columns.contains(&column)), + single_column_unique: table_plan + .unique_constraints + .iter() + .any(|constraint| constraint.columns == [column]), + single_column_indexed: table_plan.indexes.iter().any(|index| index.columns == [column]), + sequenced: table_plan.sequences.iter().any(|sequence| sequence.column == column), + } + } + pub fn row(&self, table: usize, row: usize) -> Option<&Row> { self.visible_rows(table).nth(row) } @@ -249,8 +318,16 @@ impl Model { count: self.visible_count(table), }) .collect(); + let table_rows = (0..self.schema.tables.len()) + .map(|table| TableRows { + table, + rows: normalize_rows(self.visible_rows(table).cloned().collect()), + }) + .collect(); + CountState { row_counts, + table_rows, schema: schema_state_for_plan(&self.schema), } } diff --git a/crates/dst/src/engine/workload.rs b/crates/dst/src/engine/workload.rs index 9076985a9f0..04e6448d591 100644 --- a/crates/dst/src/engine/workload.rs +++ b/crates/dst/src/engine/workload.rs @@ -1,13 +1,12 @@ use std::fmt::{Debug, Error, Formatter}; -use spacetimedb_lib::bsatn::to_vec; -use spacetimedb_lib::{AlgebraicType, AlgebraicValue, ProductValue}; -use spacetimedb_runtime::sim::Rng; -use spacetimedb_sats::ArrayValue; - +use super::generation::{pick_weighted, GenCtx}; use super::migrations::Migration; use super::model::Model; -use crate::schema::{IndexAlgorithm, SchemaPlan, TablePlan, Type}; +use crate::schema::{IndexAlgorithm, SchemaPlan}; +use spacetimedb_lib::bsatn::to_vec; +use spacetimedb_lib::{AlgebraicType, ProductValue}; +use spacetimedb_runtime::sim::Rng; pub type Row = ProductValue; @@ -118,55 +117,11 @@ impl WorkloadGen { self.model.schema() } - fn gen_value(&self, ty: Type) -> AlgebraicValue { - match ty { - Type::Bool => AlgebraicValue::Bool(self.rng.next_u64().is_multiple_of(2)), - Type::I64 => AlgebraicValue::I64(self.rng.next_u64() as i64), - Type::U64 => AlgebraicValue::U64(self.rng.next_u64()), - Type::String => AlgebraicValue::String(format!("v_{}", self.rng.next_u64()).into()), - Type::Bytes => { - let len = (self.rng.next_u64() % 16) as usize; - let bytes: Vec = (0..len).map(|_| self.rng.next_u64() as u8).collect(); - AlgebraicValue::Array(ArrayValue::U8(bytes.into())) - } - Type::Sum { variants } => { - let tag = self.rng.index(variants as usize) as u8; - AlgebraicValue::sum(tag, AlgebraicValue::U8(self.rng.next_u64() as u8)) - } - } - } - - fn gen_row(&self, table: &TablePlan) -> Row { - table - .columns - .iter() - .map(|c| self.gen_value(c.ty)) - .collect::() - } - - fn gen_insert_row(&self, table_idx: usize) -> Row { - let table = &self.schema().tables[table_idx]; - let mut row = self.gen_row(table); - - if let Some(sequence) = table.sequences.first() { - row.elements[sequence.column] = match table.columns[sequence.column].ty { - Type::I64 => AlgebraicValue::I64(0), - Type::U64 => AlgebraicValue::U64(0), - _ => unreachable!("sequence columns are integral"), - }; - } - - row - } - - fn non_auto_inc_table_idx(&self) -> Option { - let auto_inc_table = self - .schema() - .auto_inc_table_and_column() - .map(|(table_idx, _)| table_idx); - - (0..self.schema().tables.len()) - .find(|&table_idx| Some(table_idx) != auto_inc_table && !self.schema().tables[table_idx].is_event) + fn non_sequenced_table_idx(&self) -> Option { + (0..self.schema().tables.len()).find(|&table_idx| { + let table = &self.schema().tables[table_idx]; + !table.is_event && table.sequences.is_empty() + }) } pub fn next_interaction(&mut self) -> Interaction { @@ -206,7 +161,7 @@ impl WorkloadGen { Interaction::Insert { table, - row: self.gen_insert_row(table), + row: GenCtx::new(&self.rng, &self.model).gen_insert_row(table), } } @@ -234,13 +189,16 @@ impl WorkloadGen { fn pick_interaction_choice(&mut self) -> InteractionChoice { let weights = self.weights; - match self.pick_weighted(&[ - weights.insert, - weights.delete, - weights.commit_tx, - weights.migrate, - weights.replay, - ]) { + match pick_weighted( + &self.rng, + &[ + weights.insert, + weights.delete, + weights.commit_tx, + weights.migrate, + weights.replay, + ], + ) { 0 => InteractionChoice::Insert, 1 => InteractionChoice::Delete, 2 => InteractionChoice::CommitTx, @@ -250,38 +208,26 @@ impl WorkloadGen { } } - fn pick_weighted(&mut self, weights: &[u64]) -> usize { - let total: u64 = weights.iter().sum(); - - assert!(total > 0, "at least one interaction weight must be non-zero"); - - let mut selected = self.rng.next_u64() % total; - - for (idx, weight) in weights.iter().copied().enumerate() { - if selected < weight { - return idx; - } - - selected -= weight; - } - - unreachable!("selected value is always inside total weight") - } - fn insert_table_idx(&self) -> usize { - let auto_inc_table_idx = self - .schema() - .auto_inc_table_and_column() - .map(|(table_idx, _)| table_idx) - .filter(|&table_idx| !self.schema().tables[table_idx].is_event); + let sequenced_tables = self.sequenced_table_indices(); let data_tables = self.data_table_indices(); - match auto_inc_table_idx { - Some(table_idx) if !self.rng.next_u64().is_multiple_of(3) => table_idx, - _ => data_tables[self.rng.index(data_tables.len())], + if !sequenced_tables.is_empty() && !self.rng.next_u64().is_multiple_of(3) { + sequenced_tables[self.rng.index(sequenced_tables.len())] + } else { + data_tables[self.rng.index(data_tables.len())] } } + fn sequenced_table_indices(&self) -> Vec { + self.schema() + .tables + .iter() + .enumerate() + .filter_map(|(table_idx, table)| (!table.is_event && !table.sequences.is_empty()).then_some(table_idx)) + .collect() + } + fn data_table_indices(&self) -> Vec { let data_tables: Vec<_> = self .schema() @@ -298,11 +244,11 @@ impl WorkloadGen { } fn gen_migration(&self) -> Option { - Migration::choose(self.schema(), &self.rng, |table| self.model.row_count(table)) + GenCtx::new(&self.rng, &self.model).gen_migration() } fn deletable_table_idx(&self) -> Option { - self.non_auto_inc_table_idx() + self.non_sequenced_table_idx() .filter(|&table_idx| self.model.row_count(table_idx) > 0) } } @@ -333,6 +279,7 @@ pub fn normalize_rows(mut rows: Vec) -> Vec { #[derive(Debug, Clone, PartialEq, Eq)] pub struct CountState { pub row_counts: Vec, + pub table_rows: Vec, pub schema: SchemaState, } @@ -342,6 +289,12 @@ pub struct TableRowCount { pub count: u64, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableRows { + pub table: usize, + pub rows: Vec, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct SchemaState { pub tables: Vec, @@ -356,6 +309,7 @@ pub struct TableSchemaState { pub primary_key: Option, pub columns: Vec, pub indexes: Vec, + pub unique_constraints: Vec, pub sequences: Vec, } @@ -379,6 +333,11 @@ pub enum IndexAlgorithmState { Unknown, } +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct UniqueConstraintState { + pub columns: Vec, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub struct SequenceState { pub column: usize, @@ -404,6 +363,15 @@ pub fn schema_state_for_plan(schema: &SchemaPlan) -> SchemaState { .collect::>(); indexes.sort(); + let mut unique_constraints = table_plan + .unique_constraints + .iter() + .map(|constraint| UniqueConstraintState { + columns: constraint.columns.clone(), + }) + .collect::>(); + unique_constraints.sort(); + let mut sequences = table_plan .sequences .iter() @@ -428,6 +396,7 @@ pub fn schema_state_for_plan(schema: &SchemaPlan) -> SchemaState { }) .collect(), indexes, + unique_constraints, sequences, } }) diff --git a/crates/dst/src/schema.rs b/crates/dst/src/schema.rs index 830f9102bcf..7f30e713973 100644 --- a/crates/dst/src/schema.rs +++ b/crates/dst/src/schema.rs @@ -166,6 +166,19 @@ impl SchemaNames { .join("_") ) } + + pub fn constraint_name(table: &TablePlan, constraint: &UniqueConstraintPlan) -> String { + format!( + "{}_{}_key", + table.name, + constraint + .columns + .iter() + .map(|&c| table.columns[c].name.as_str()) + .collect::>() + .join("_") + ) + } } // Schema plan — the canonical source of truth. @@ -179,6 +192,9 @@ impl SchemaPlan { pub fn ensure_engine_migration_coverage(&mut self) { self.ensure_widenable_sum_column(); self.ensure_event_table(); + self.ensure_sequence_mutation_column(); + self.ensure_unique_constraint_mutation_column(); + self.ensure_standalone_index(); } pub fn auto_inc_table_and_column(&self) -> Option<(usize, usize)> { @@ -260,6 +276,112 @@ impl SchemaPlan { is_event: true, }); } + + fn ensure_sequence_mutation_column(&mut self) { + if self.tables.iter().any(Self::has_sequence_mutation_column) { + return; + } + + let table = self + .tables + .iter_mut() + .find(|table| !table.is_event) + .expect("schema must contain at least one non-event table"); + let column = table.columns.len(); + table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "dst_seq"), + ty: Type::U64, + }); + table.indexes.push(IndexPlan { + columns: vec![column], + algorithm: IndexAlgorithm::BTree, + }); + table + .unique_constraints + .push(UniqueConstraintPlan { columns: vec![column] }); + } + + fn ensure_unique_constraint_mutation_column(&mut self) { + if self.tables.iter().any(|table| { + !table.is_event + && table.columns.iter().enumerate().any(|(column, _)| { + !table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == [column]) + }) + }) { + return; + } + + let table = self + .tables + .iter_mut() + .find(|table| !table.is_event) + .expect("schema must contain at least one non-event table"); + table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "dst_unique"), + ty: Type::U64, + }); + } + + fn has_sequence_mutation_column(table: &TablePlan) -> bool { + !table.is_event + && table.columns.iter().enumerate().any(|(column, plan)| { + plan.ty == Type::U64 + && table.sequences.iter().all(|seq| seq.column != column) + && table.indexes.iter().any(|index| index.columns == [column]) + && table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == [column]) + }) + } + + fn ensure_standalone_index(&mut self) { + if self.tables.iter().any(|table| { + !table.is_event + && table.indexes.iter().any(|index| { + !table.primary_key.is_some_and(|pk| index.columns == [pk]) + && !table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == index.columns) + }) + }) { + return; + } + + let table = self + .tables + .iter_mut() + .find(|table| !table.is_event) + .expect("schema must contain at least one non-event table"); + let column = if table.columns.len() < 2 { + table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "dst_indexed"), + ty: Type::U64, + }); + table.columns.len() - 1 + } else { + (0..table.columns.len()) + .find(|&column| { + !table.primary_key.is_some_and(|pk| pk == column) + && !table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == [column]) + }) + .unwrap_or(0) + }; + + if !table.indexes.iter().any(|index| index.columns == [column]) { + table.indexes.push(IndexPlan { + columns: vec![column], + algorithm: IndexAlgorithm::BTree, + }); + } + } } #[derive(Debug, Clone)] @@ -300,10 +422,14 @@ pub struct UniqueConstraintPlan { } /// A sequence on a specific integral column. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct SequencePlan { /// Index into `TablePlan.columns`. pub column: usize, + pub start: Option, + pub min_value: Option, + pub max_value: Option, + pub increment: i128, } impl SequencePlan { @@ -312,7 +438,44 @@ impl SequencePlan { if !ty.is_integral() { return None; } - Some(Self { column }) + Some(Self { + column, + start: None, + min_value: None, + max_value: None, + increment: 1, + }) + } + + pub fn with_bounds( + column: usize, + ty: Type, + start: i128, + min_value: i128, + max_value: i128, + increment: i128, + ) -> Option { + if !ty.is_integral() || increment == 0 || min_value >= max_value || start < min_value || start > max_value { + return None; + } + Some(Self { + column, + start: Some(start), + min_value: Some(min_value), + max_value: Some(max_value), + increment, + }) + } + + pub fn with_existing_value_as_max(column: usize, ty: Type, existing_value: i128) -> Option { + const DOMAIN_SIZE: i128 = 3; + + if existing_value < DOMAIN_SIZE { + return None; + } + + let min_value = existing_value - (DOMAIN_SIZE - 1); + Self::with_bounds(column, ty, min_value, min_value, existing_value, 1) } } @@ -325,7 +488,20 @@ pub fn to_raw_def(schema: &SchemaPlan) -> RawModuleDefV10 { to_raw_def_table(&mut builder, table); } - builder.finish() + let mut raw = builder.finish(); + apply_sequence_bounds(schema, &mut raw); + raw +} + +fn apply_sequence_bounds(schema: &SchemaPlan, raw: &mut RawModuleDefV10) { + for (table_plan, raw_table) in schema.tables.iter().zip(raw.tables_mut_for_tests().iter_mut()) { + for (sequence_plan, raw_sequence) in table_plan.sequences.iter().zip(raw_table.sequences.iter_mut()) { + raw_sequence.start = sequence_plan.start; + raw_sequence.min_value = sequence_plan.min_value; + raw_sequence.max_value = sequence_plan.max_value; + raw_sequence.increment = sequence_plan.increment; + } + } } fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { @@ -456,10 +632,10 @@ impl SchemaGenerator { } } // Ensure PK has a matching unique constraint. - if let Some(pk) = pk - && !seen.iter().any(|cols| cols.len() == 1 && cols[0] == *pk) - { - result.push(UniqueConstraintPlan { columns: vec![*pk] }); + if let Some(pk) = pk { + if !seen.iter().any(|cols| cols.len() == 1 && cols[0] == *pk) { + result.push(UniqueConstraintPlan { columns: vec![*pk] }); + } } result } diff --git a/crates/dst/src/traits.rs b/crates/dst/src/traits.rs index 2185e0ec918..5abfc770092 100644 --- a/crates/dst/src/traits.rs +++ b/crates/dst/src/traits.rs @@ -1,4 +1,7 @@ +use std::panic::{resume_unwind, AssertUnwindSafe}; + use anyhow::Error; +use futures::FutureExt; use spacetimedb_runtime::sim::Rng; /// This should be implemented by System under test. @@ -39,19 +42,24 @@ pub trait TestSuite { async move { let (mut interactions, mut target, mut properties) = self.build(rng).await?; - let result = async { + let result = AssertUnwindSafe(async { for interaction in interactions.by_ref().take(max_interactions) { let observation = target.execute(&interaction).await?; properties.observe(&interaction, &observation)?; } Ok(()) - } + }) + .catch_unwind() .await; + eprintln!("final interaction counts: {interactions:?}"); tracing::info!(interaction_counts = ?interactions, "final interaction counts"); - result + match result { + Ok(result) => result, + Err(payload) => resume_unwind(payload), + } } } } From def3e4ef6e3b2d101159fae6fba820d9cb4240d4 Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Mon, 13 Jul 2026 17:34:55 +0530 Subject: [PATCH 03/10] proptest like gen --- crates/dst/src/engine.rs | 9 +- crates/dst/src/engine/generation.rs | 244 ++++++++++++++++++++++------ crates/dst/src/engine/model.rs | 6 +- crates/dst/src/engine/state.rs | 145 +++++++++++++++++ crates/dst/src/engine/workload.rs | 146 +---------------- 5 files changed, 347 insertions(+), 203 deletions(-) create mode 100644 crates/dst/src/engine/state.rs diff --git a/crates/dst/src/engine.rs b/crates/dst/src/engine.rs index ac1b02b1087..137da2169ac 100644 --- a/crates/dst/src/engine.rs +++ b/crates/dst/src/engine.rs @@ -22,14 +22,15 @@ mod generation; mod migrations; mod model; mod properties; +mod state; mod workload; use self::migrations::{ExpectedStep, Migration}; -use self::workload::{ - normalize_rows, row_to_bytes, ColumnState, CommitDelta, CountState, IndexAlgorithmState, IndexState, InsertOutcome, - Interaction, Observation, SchemaState, SequenceState, TableDelta, TableRowCount, TableRows, TableSchemaState, - UniqueConstraintState, +use self::state::{ + ColumnState, CommitDelta, CountState, IndexAlgorithmState, IndexState, SchemaState, SequenceState, TableDelta, + TableRowCount, TableRows, TableSchemaState, UniqueConstraintState, }; +use self::workload::{normalize_rows, row_to_bytes, InsertOutcome, Interaction, Observation}; use crate::engine::model::Model; use crate::engine::properties::EngineProperties; diff --git a/crates/dst/src/engine/generation.rs b/crates/dst/src/engine/generation.rs index 9698e13b1df..0f7521bb2af 100644 --- a/crates/dst/src/engine/generation.rs +++ b/crates/dst/src/engine/generation.rs @@ -26,6 +26,34 @@ impl<'a> GenCtx<'a> { } } +#[derive(Clone, Copy)] +pub(crate) struct Choice { + weight: u64, + value: T, +} + +pub(crate) const fn choice(weight: u64, value: T) -> Choice { + Choice { weight, value } +} + +pub(crate) fn frequency(rng: &Rng, choices: &[Choice]) -> T { + let total: u64 = choices.iter().map(|choice| choice.weight).sum(); + + assert!(total > 0, "at least one choice weight must be non-zero"); + + let mut selected = rng.next_u64() % total; + + for choice in choices.iter().copied() { + if selected < choice.weight { + return choice.value; + } + + selected -= choice.weight; + } + + unreachable!("selected value is always inside total weight") +} + pub(crate) fn pick_weighted(rng: &Rng, weights: &[u64]) -> usize { let total: u64 = weights.iter().sum(); @@ -44,6 +72,59 @@ pub(crate) fn pick_weighted(rng: &Rng, weights: &[u64]) -> usize { unreachable!("selected value is always inside total weight") } +#[derive(Clone, Copy)] +enum ValueCase { + Random, + Small, + Edge, + NearExisting, + Existing, + Weird, +} + +#[derive(Clone, Copy)] +enum FreshValueCase { + Random, + Small, + Edge, + Weird, +} + +#[derive(Clone, Copy)] +enum I64Case { + Random, + Small, + Edge, +} + +#[derive(Clone, Copy)] +enum U64Case { + Random, + Small, + Edge, +} + +#[derive(Clone, Copy)] +enum StringCase { + RandomTagged, + Empty, + SmallAscii, + OrderedPrefix, + SqlEscaped, + NullByte, + Long, +} + +#[derive(Clone, Copy)] +enum BytesCase { + Random, + Empty, + Small, + RepeatedZero, + RepeatedMax, + Alternating, +} + struct ValueGen<'a> { rng: &'a Rng, model: &'a Model, @@ -79,18 +160,27 @@ impl<'a> ValueGen<'a> { } fn gen_value(&self, domain: &ColumnDomain) -> AlgebraicValue { - match pick_weighted(self.rng, &[45, 15, 15, 10, 10, 5]) { - 0 => self.gen_random_value(domain.ty), - 1 => self.gen_small_value(domain.ty), - 2 => self.gen_edge_value(domain.ty), - 3 => self + match frequency( + self.rng, + &[ + choice(45, ValueCase::Random), + choice(15, ValueCase::Small), + choice(15, ValueCase::Edge), + choice(10, ValueCase::NearExisting), + choice(10, ValueCase::Existing), + choice(5, ValueCase::Weird), + ], + ) { + ValueCase::Random => self.gen_random_value(domain.ty), + ValueCase::Small => self.gen_small_value(domain.ty), + ValueCase::Edge => self.gen_edge_value(domain.ty), + ValueCase::NearExisting => self .near_existing_value(domain) .unwrap_or_else(|| self.gen_random_value(domain.ty)), - 4 => self + ValueCase::Existing => self .existing_value(domain) .unwrap_or_else(|| self.gen_random_value(domain.ty)), - 5 => self.gen_weird_value(domain.ty), - _ => unreachable!(), + ValueCase::Weird => self.gen_weird_value(domain.ty), } } @@ -106,12 +196,19 @@ impl<'a> ValueGen<'a> { } fn gen_fresh_candidate(&self, ty: Type) -> AlgebraicValue { - match pick_weighted(self.rng, &[50, 20, 20, 10]) { - 0 => self.gen_random_value(ty), - 1 => self.gen_small_value(ty), - 2 => self.gen_edge_value(ty), - 3 => self.gen_weird_value(ty), - _ => unreachable!(), + match frequency( + self.rng, + &[ + choice(50, FreshValueCase::Random), + choice(20, FreshValueCase::Small), + choice(20, FreshValueCase::Edge), + choice(10, FreshValueCase::Weird), + ], + ) { + FreshValueCase::Random => self.gen_random_value(ty), + FreshValueCase::Small => self.gen_small_value(ty), + FreshValueCase::Edge => self.gen_edge_value(ty), + FreshValueCase::Weird => self.gen_weird_value(ty), } } @@ -143,14 +240,10 @@ impl<'a> ValueGen<'a> { fn gen_random_value(&self, ty: Type) -> AlgebraicValue { match ty { Type::Bool => AlgebraicValue::Bool(self.rng.next_u64().is_multiple_of(2)), - Type::I64 => AlgebraicValue::I64(self.rng.next_u64() as i64), - Type::U64 => AlgebraicValue::U64(self.rng.next_u64()), - Type::String => AlgebraicValue::String(format!("v_{}", self.rng.next_u64()).into()), - Type::Bytes => { - let len = (self.rng.next_u64() % 16) as usize; - let bytes: Vec = (0..len).map(|_| self.rng.next_u64() as u8).collect(); - AlgebraicValue::Array(ArrayValue::U8(bytes.into())) - } + Type::I64 => AlgebraicValue::I64(self.gen_i64_value(I64Case::Random)), + Type::U64 => AlgebraicValue::U64(self.gen_u64_value(U64Case::Random)), + Type::String => AlgebraicValue::String(self.gen_string_value(StringCase::RandomTagged).into()), + Type::Bytes => AlgebraicValue::Array(ArrayValue::U8(self.gen_bytes_value(BytesCase::Random).into())), Type::Sum { variants } => { let tag = self.rng.index(variants as usize) as u8; AlgebraicValue::sum(tag, AlgebraicValue::U8(self.rng.next_u64() as u8)) @@ -161,22 +254,10 @@ impl<'a> ValueGen<'a> { fn gen_small_value(&self, ty: Type) -> AlgebraicValue { match ty { Type::Bool => AlgebraicValue::Bool(false), - Type::I64 => { - const VALUES: &[i64] = &[-1, 0, 1, 2, 3]; - AlgebraicValue::I64(VALUES[self.rng.index(VALUES.len())]) - } - Type::U64 => { - const VALUES: &[u64] = &[0, 1, 2, 3]; - AlgebraicValue::U64(VALUES[self.rng.index(VALUES.len())]) - } - Type::String => { - const VALUES: &[&str] = &["", "a", "aa", "ab", "b", "v_0"]; - AlgebraicValue::String(VALUES[self.rng.index(VALUES.len())].into()) - } - Type::Bytes => { - const VALUES: &[&[u8]] = &[&[], &[0], &[1], &[0, 255]]; - AlgebraicValue::Array(ArrayValue::U8(VALUES[self.rng.index(VALUES.len())].to_vec().into())) - } + Type::I64 => AlgebraicValue::I64(self.gen_i64_value(I64Case::Small)), + Type::U64 => AlgebraicValue::U64(self.gen_u64_value(U64Case::Small)), + Type::String => AlgebraicValue::String(self.gen_string_value(StringCase::SmallAscii).into()), + Type::Bytes => AlgebraicValue::Array(ArrayValue::U8(self.gen_bytes_value(BytesCase::Small).into())), Type::Sum { variants } => { let tag = if variants <= 1 { 0 @@ -191,16 +272,10 @@ impl<'a> ValueGen<'a> { fn gen_edge_value(&self, ty: Type) -> AlgebraicValue { match ty { Type::Bool => AlgebraicValue::Bool(true), - Type::I64 => { - const VALUES: &[i64] = &[i64::MIN, i64::MIN + 1, i64::MAX - 1, i64::MAX]; - AlgebraicValue::I64(VALUES[self.rng.index(VALUES.len())]) - } - Type::U64 => { - const VALUES: &[u64] = &[0, 1, u64::MAX - 1, u64::MAX]; - AlgebraicValue::U64(VALUES[self.rng.index(VALUES.len())]) - } - Type::String => AlgebraicValue::String("x".repeat(128).into()), - Type::Bytes => AlgebraicValue::Array(ArrayValue::U8(vec![255; 32].into())), + Type::I64 => AlgebraicValue::I64(self.gen_i64_value(I64Case::Edge)), + Type::U64 => AlgebraicValue::U64(self.gen_u64_value(U64Case::Edge)), + Type::String => AlgebraicValue::String(self.gen_string_value(StringCase::Long).into()), + Type::Bytes => AlgebraicValue::Array(ArrayValue::U8(self.gen_bytes_value(BytesCase::RepeatedMax).into())), Type::Sum { variants } => AlgebraicValue::sum(variants.saturating_sub(1), AlgebraicValue::U8(u8::MAX)), } } @@ -208,17 +283,78 @@ impl<'a> ValueGen<'a> { fn gen_weird_value(&self, ty: Type) -> AlgebraicValue { match ty { Type::String => { - const VALUES: &[&str] = &["quote'", "double\"quote", "back\\slash", "line\nbreak", "\0"]; - AlgebraicValue::String(VALUES[self.rng.index(VALUES.len())].into()) + let case = frequency( + self.rng, + &[ + choice(35, StringCase::SqlEscaped), + choice(25, StringCase::NullByte), + choice(25, StringCase::OrderedPrefix), + choice(15, StringCase::Empty), + ], + ); + AlgebraicValue::String(self.gen_string_value(case).into()) } Type::Bytes => { - const VALUES: &[&[u8]] = &[&[0], &[0, 0, 0], &[255], &[0, 255, 0, 255]]; - AlgebraicValue::Array(ArrayValue::U8(VALUES[self.rng.index(VALUES.len())].to_vec().into())) + let case = frequency( + self.rng, + &[ + choice(25, BytesCase::Empty), + choice(20, BytesCase::RepeatedZero), + choice(20, BytesCase::RepeatedMax), + choice(20, BytesCase::Alternating), + choice(15, BytesCase::Small), + ], + ); + AlgebraicValue::Array(ArrayValue::U8(self.gen_bytes_value(case).into())) } _ => self.gen_edge_value(ty), } } + fn gen_i64_value(&self, case: I64Case) -> i64 { + match case { + I64Case::Random => self.rng.next_u64() as i64, + I64Case::Small => self.sample(&[-3, -2, -1, 0, 1, 2, 3]), + I64Case::Edge => self.sample(&[i64::MIN, i64::MIN + 1, -1, 0, 1, i64::MAX - 1, i64::MAX]), + } + } + + fn gen_u64_value(&self, case: U64Case) -> u64 { + match case { + U64Case::Random => self.rng.next_u64(), + U64Case::Small => self.sample(&[0, 1, 2, 3, 4, 5]), + U64Case::Edge => self.sample(&[0, 1, 2, u64::MAX - 1, u64::MAX]), + } + } + + fn gen_string_value(&self, case: StringCase) -> String { + match case { + StringCase::RandomTagged => format!("v_{}", self.rng.next_u64()), + StringCase::Empty => String::new(), + StringCase::SmallAscii => self.sample(&["a", "aa", "ab", "b", "z", "v_0", "v_1"]).to_owned(), + StringCase::OrderedPrefix => self.sample(&["a", "aa", "aaa", "ab", "aba", "abb", "b"]).to_owned(), + StringCase::SqlEscaped => self + .sample(&["quote'", "double\"quote", "back\\slash", "line\nbreak"]) + .to_owned(), + StringCase::NullByte => "nul\0byte".to_owned(), + StringCase::Long => "x".repeat(128), + } + } + + fn gen_bytes_value(&self, case: BytesCase) -> Vec { + match case { + BytesCase::Random => { + let len = (self.rng.next_u64() % 16) as usize; + (0..len).map(|_| self.rng.next_u64() as u8).collect() + } + BytesCase::Empty => Vec::new(), + BytesCase::Small => self.sample(&[&[][..], &[0][..], &[1][..], &[0, 255][..]]).to_vec(), + BytesCase::RepeatedZero => vec![0; 32], + BytesCase::RepeatedMax => vec![255; 32], + BytesCase::Alternating => vec![0, 255, 0, 255, 0, 255], + } + } + fn gen_counter_value(&self, ty: Type, counter: u64) -> AlgebraicValue { match ty { Type::Bool => AlgebraicValue::Bool(counter.is_multiple_of(2)), @@ -236,6 +372,10 @@ impl<'a> ValueGen<'a> { ), } } + + fn sample(&self, values: &[T]) -> T { + values[self.rng.index(values.len())] + } } fn sequence_placeholder(ty: Type) -> AlgebraicValue { diff --git a/crates/dst/src/engine/model.rs b/crates/dst/src/engine/model.rs index 748170da36f..834deca3889 100644 --- a/crates/dst/src/engine/model.rs +++ b/crates/dst/src/engine/model.rs @@ -1,9 +1,7 @@ use spacetimedb_lib::AlgebraicValue; -use super::workload::{ - normalize_rows, schema_state_for_plan, CommitDelta, CountState, InsertOutcome, Interaction, Observation, Row, - TableDelta, TableRowCount, TableRows, -}; +use super::state::{schema_state_for_plan, CommitDelta, CountState, TableDelta, TableRowCount, TableRows}; +use super::workload::{normalize_rows, InsertOutcome, Interaction, Observation, Row}; use crate::schema::{SchemaPlan, Type}; #[derive(Debug)] diff --git a/crates/dst/src/engine/state.rs b/crates/dst/src/engine/state.rs new file mode 100644 index 00000000000..35512196024 --- /dev/null +++ b/crates/dst/src/engine/state.rs @@ -0,0 +1,145 @@ +use spacetimedb_lib::AlgebraicType; + +use super::workload::Row; +use crate::schema::{IndexAlgorithm, SchemaPlan}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CountState { + pub row_counts: Vec, + pub table_rows: Vec, + pub schema: SchemaState, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TableRowCount { + pub table: usize, + pub count: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableRows { + pub table: usize, + pub rows: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SchemaState { + pub tables: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableSchemaState { + pub table: usize, + pub name: String, + pub is_public: bool, + pub is_event: bool, + pub primary_key: Option, + pub columns: Vec, + pub indexes: Vec, + pub unique_constraints: Vec, + pub sequences: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ColumnState { + pub name: String, + pub ty: AlgebraicType, +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct IndexState { + pub columns: Vec, + pub algorithm: IndexAlgorithmState, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum IndexAlgorithmState { + BTree, + Hash, + Direct, + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct UniqueConstraintState { + pub columns: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub struct SequenceState { + pub column: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CommitDelta { + pub tables: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableDelta { + pub table: usize, + pub inserts: Vec, + pub deletes: Vec, + pub truncated: bool, +} + +pub fn schema_state_for_plan(schema: &SchemaPlan) -> SchemaState { + SchemaState { + tables: schema + .tables + .iter() + .enumerate() + .map(|(table, table_plan)| { + let mut indexes = table_plan + .indexes + .iter() + .map(|index| IndexState { + columns: index.columns.clone(), + algorithm: match index.algorithm { + IndexAlgorithm::BTree => IndexAlgorithmState::BTree, + IndexAlgorithm::Hash => IndexAlgorithmState::Hash, + }, + }) + .collect::>(); + indexes.sort(); + + let mut unique_constraints = table_plan + .unique_constraints + .iter() + .map(|constraint| UniqueConstraintState { + columns: constraint.columns.clone(), + }) + .collect::>(); + unique_constraints.sort(); + + let mut sequences = table_plan + .sequences + .iter() + .map(|sequence| SequenceState { + column: sequence.column, + }) + .collect::>(); + sequences.sort(); + + TableSchemaState { + table, + name: table_plan.name.clone(), + is_public: table_plan.is_public, + is_event: table_plan.is_event, + primary_key: table_plan.primary_key, + columns: table_plan + .columns + .iter() + .map(|column| ColumnState { + name: column.name.clone(), + ty: column.ty.to_algebraic(), + }) + .collect(), + indexes, + unique_constraints, + sequences, + } + }) + .collect(), + } +} diff --git a/crates/dst/src/engine/workload.rs b/crates/dst/src/engine/workload.rs index 04e6448d591..20bc2678f3d 100644 --- a/crates/dst/src/engine/workload.rs +++ b/crates/dst/src/engine/workload.rs @@ -3,9 +3,10 @@ use std::fmt::{Debug, Error, Formatter}; use super::generation::{pick_weighted, GenCtx}; use super::migrations::Migration; use super::model::Model; -use crate::schema::{IndexAlgorithm, SchemaPlan}; +use super::state::{CommitDelta, CountState}; +use crate::schema::SchemaPlan; use spacetimedb_lib::bsatn::to_vec; -use spacetimedb_lib::{AlgebraicType, ProductValue}; +use spacetimedb_lib::ProductValue; use spacetimedb_runtime::sim::Rng; pub type Row = ProductValue; @@ -275,144 +276,3 @@ pub fn normalize_rows(mut rows: Vec) -> Vec { rows.sort_by_key(row_to_bytes); rows } - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CountState { - pub row_counts: Vec, - pub table_rows: Vec, - pub schema: SchemaState, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct TableRowCount { - pub table: usize, - pub count: u64, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct TableRows { - pub table: usize, - pub rows: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct SchemaState { - pub tables: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct TableSchemaState { - pub table: usize, - pub name: String, - pub is_public: bool, - pub is_event: bool, - pub primary_key: Option, - pub columns: Vec, - pub indexes: Vec, - pub unique_constraints: Vec, - pub sequences: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ColumnState { - pub name: String, - pub ty: AlgebraicType, -} - -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] -pub struct IndexState { - pub columns: Vec, - pub algorithm: IndexAlgorithmState, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -pub enum IndexAlgorithmState { - BTree, - Hash, - Direct, - Unknown, -} - -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] -pub struct UniqueConstraintState { - pub columns: Vec, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -pub struct SequenceState { - pub column: usize, -} - -pub fn schema_state_for_plan(schema: &SchemaPlan) -> SchemaState { - SchemaState { - tables: schema - .tables - .iter() - .enumerate() - .map(|(table, table_plan)| { - let mut indexes = table_plan - .indexes - .iter() - .map(|index| IndexState { - columns: index.columns.clone(), - algorithm: match index.algorithm { - IndexAlgorithm::BTree => IndexAlgorithmState::BTree, - IndexAlgorithm::Hash => IndexAlgorithmState::Hash, - }, - }) - .collect::>(); - indexes.sort(); - - let mut unique_constraints = table_plan - .unique_constraints - .iter() - .map(|constraint| UniqueConstraintState { - columns: constraint.columns.clone(), - }) - .collect::>(); - unique_constraints.sort(); - - let mut sequences = table_plan - .sequences - .iter() - .map(|sequence| SequenceState { - column: sequence.column, - }) - .collect::>(); - sequences.sort(); - - TableSchemaState { - table, - name: table_plan.name.clone(), - is_public: table_plan.is_public, - is_event: table_plan.is_event, - primary_key: table_plan.primary_key, - columns: table_plan - .columns - .iter() - .map(|column| ColumnState { - name: column.name.clone(), - ty: column.ty.to_algebraic(), - }) - .collect(), - indexes, - unique_constraints, - sequences, - } - }) - .collect(), - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CommitDelta { - pub tables: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct TableDelta { - pub table: usize, - pub inserts: Vec, - pub deletes: Vec, - pub truncated: bool, -} From c1b2251ff782ea4ddac48dcc410bacb013646f2b Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Mon, 13 Jul 2026 18:07:54 +0530 Subject: [PATCH 04/10] more cleanup --- crates/dst/src/engine.rs | 95 +++-------- crates/dst/src/engine/generation.rs | 53 +----- crates/dst/src/engine/migrations.rs | 242 ++++++++++++++-------------- crates/dst/src/engine/model.rs | 36 +++-- crates/dst/src/engine/properties.rs | 3 +- crates/dst/src/engine/row.rs | 13 ++ crates/dst/src/engine/state.rs | 196 ++++++++++++++++------ crates/dst/src/engine/workload.rs | 17 +- crates/dst/src/lib.rs | 1 + crates/dst/src/rng.rs | 58 +++++++ crates/dst/src/schema.rs | 10 +- 11 files changed, 394 insertions(+), 330 deletions(-) create mode 100644 crates/dst/src/engine/row.rs create mode 100644 crates/dst/src/rng.rs diff --git a/crates/dst/src/engine.rs b/crates/dst/src/engine.rs index 137da2169ac..87a4452c6b6 100644 --- a/crates/dst/src/engine.rs +++ b/crates/dst/src/engine.rs @@ -7,14 +7,13 @@ use spacetimedb_engine::error::{DBError, DatastoreError, IndexError}; use spacetimedb_engine::persistence::{DiskSizeFn, Durability as EngineDurability, Persistence}; use spacetimedb_engine::relational_db::{MutTx, RelationalDB}; use spacetimedb_engine::update::{update_database, UpdateLogger}; -use spacetimedb_lib::db::auth::StAccess; use spacetimedb_lib::identity::AuthCtx; use spacetimedb_lib::{Identity, RawModuleDef}; use spacetimedb_primitives::TableId; use spacetimedb_runtime::sim::{Rng, Runtime as SimRuntime}; use spacetimedb_runtime::Handle; use spacetimedb_schema::auto_migrate::{ponder_migrate, AutoMigrateStep, MigratePlan}; -use spacetimedb_schema::def::{IndexAlgorithm as SchemaIndexAlgorithm, ModuleDef}; +use spacetimedb_schema::def::ModuleDef; use spacetimedb_schema::schema::{Schema, TableSchema}; use spacetimedb_table::page_pool::PagePool; @@ -22,15 +21,16 @@ mod generation; mod migrations; mod model; mod properties; +mod row; mod state; mod workload; use self::migrations::{ExpectedStep, Migration}; +use self::row::{normalize_rows, row_to_bytes}; use self::state::{ - ColumnState, CommitDelta, CountState, IndexAlgorithmState, IndexState, SchemaState, SequenceState, TableDelta, - TableRowCount, TableRows, TableSchemaState, UniqueConstraintState, + table_schema_state_for_schema, CommitDelta, CountState, SchemaState, TableDelta, TableRowCount, TableRows, }; -use self::workload::{normalize_rows, row_to_bytes, InsertOutcome, Interaction, Observation}; +use self::workload::{InsertOutcome, Interaction, Observation}; use crate::engine::model::Model; use crate::engine::properties::EngineProperties; @@ -174,54 +174,7 @@ impl EngineTarget { return Err(err.into()); } }; - let mut indexes = schema - .indexes - .iter() - .map(|index| IndexState { - columns: index_columns(&index.index_algorithm), - algorithm: index_algorithm_state(&index.index_algorithm), - }) - .collect::>(); - indexes.sort(); - - let mut unique_constraints = schema - .constraints - .iter() - .filter_map(|constraint| { - constraint.data.unique_columns().map(|columns| UniqueConstraintState { - columns: columns.iter().map(|col| col.0 as usize).collect(), - }) - }) - .collect::>(); - unique_constraints.sort(); - - let mut sequences = schema - .sequences - .iter() - .map(|sequence| SequenceState { - column: sequence.col_pos.0 as usize, - }) - .collect::>(); - sequences.sort(); - - schema_tables.push(TableSchemaState { - table, - name: schema.table_name.to_string(), - is_public: schema.table_access == StAccess::Public, - is_event: schema.is_event, - primary_key: schema.primary_key.map(|col| col.0 as usize), - columns: schema - .columns - .iter() - .map(|column| ColumnState { - name: column.col_name.to_string(), - ty: column.col_type.clone(), - }) - .collect(), - indexes, - unique_constraints, - sequences, - }); + schema_tables.push(table_schema_state_for_schema(table, &schema)); } let _ = db.release_tx(tx); @@ -422,19 +375,6 @@ fn expected_step_from_auto_step(step: &AutoMigrateStep<'_>) -> anyhow::Result Vec { - algorithm.columns().iter().map(|col| col.0 as usize).collect() -} - -fn index_algorithm_state(algorithm: &SchemaIndexAlgorithm) -> IndexAlgorithmState { - match algorithm { - SchemaIndexAlgorithm::BTree(_) => IndexAlgorithmState::BTree, - SchemaIndexAlgorithm::Hash(_) => IndexAlgorithmState::Hash, - SchemaIndexAlgorithm::Direct(_) => IndexAlgorithmState::Direct, - _ => IndexAlgorithmState::Unknown, - } -} - struct DstUpdateLogger; impl UpdateLogger for DstUpdateLogger { @@ -481,7 +421,7 @@ mod tests { use spacetimedb_runtime::sim::Runtime as SimRuntime; use spacetimedb_sats::product; - use super::migrations::{Migration, MigrationOp}; + use super::migrations::{Migration, TableMigrationOp}; use super::*; use crate::schema::{ColumnPlan, IndexAlgorithm, IndexPlan, TablePlan, Type, UniqueConstraintPlan}; @@ -596,9 +536,12 @@ mod tests { let mut target = EngineTarget::init(add_column_replay_schema(), 0)?; insert_u64_rows(&mut target)?; - target.execute(&Interaction::Migrate(Migration { + target.execute(&Interaction::Migrate(Migration::AlterTable { table: 0, - ops: vec![MigrationOp::ChangeAccess, MigrationOp::AddColumn { ty: Type::U64 }], + ops: vec![ + TableMigrationOp::ChangeAccess, + TableMigrationOp::AddColumn { ty: Type::U64 }, + ], }))?; target.execute(&Interaction::Replay)?; @@ -610,9 +553,12 @@ mod tests { let mut target = EngineTarget::init(change_index_replay_schema(), 0)?; insert_u64_rows(&mut target)?; - target.execute(&Interaction::Migrate(Migration { + target.execute(&Interaction::Migrate(Migration::AlterTable { table: 0, - ops: vec![MigrationOp::ChangeAccess, MigrationOp::ChangeIndex { index: 1 }], + ops: vec![ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeIndex { index: 1 }, + ], }))?; target.execute(&Interaction::Replay)?; @@ -632,9 +578,12 @@ mod tests { } target.execute(&Interaction::CommitTx)?; - target.execute(&Interaction::Migrate(Migration { + target.execute(&Interaction::Migrate(Migration::AlterTable { table: 0, - ops: vec![MigrationOp::ChangeAccess, MigrationOp::ChangeColumnType { column: 1 }], + ops: vec![ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeColumnType { column: 1 }, + ], }))?; target.execute(&Interaction::Replay)?; diff --git a/crates/dst/src/engine/generation.rs b/crates/dst/src/engine/generation.rs index 0f7521bb2af..d81f7237628 100644 --- a/crates/dst/src/engine/generation.rs +++ b/crates/dst/src/engine/generation.rs @@ -4,8 +4,9 @@ use spacetimedb_sats::ArrayValue; use super::migrations::Migration; use super::model::{ColumnDomain, Model}; -use super::workload::Row; -use crate::schema::{SchemaDecisions, Type}; +use super::row::Row; +use crate::rng::{choice, choose_index, frequency}; +use crate::schema::Type; pub(crate) struct GenCtx<'a> { rng: &'a Rng, @@ -26,52 +27,6 @@ impl<'a> GenCtx<'a> { } } -#[derive(Clone, Copy)] -pub(crate) struct Choice { - weight: u64, - value: T, -} - -pub(crate) const fn choice(weight: u64, value: T) -> Choice { - Choice { weight, value } -} - -pub(crate) fn frequency(rng: &Rng, choices: &[Choice]) -> T { - let total: u64 = choices.iter().map(|choice| choice.weight).sum(); - - assert!(total > 0, "at least one choice weight must be non-zero"); - - let mut selected = rng.next_u64() % total; - - for choice in choices.iter().copied() { - if selected < choice.weight { - return choice.value; - } - - selected -= choice.weight; - } - - unreachable!("selected value is always inside total weight") -} - -pub(crate) fn pick_weighted(rng: &Rng, weights: &[u64]) -> usize { - let total: u64 = weights.iter().sum(); - - assert!(total > 0, "at least one interaction weight must be non-zero"); - - let mut selected = rng.next_u64() % total; - - for (idx, weight) in weights.iter().copied().enumerate() { - if selected < weight { - return idx; - } - - selected -= weight; - } - - unreachable!("selected value is always inside total weight") -} - #[derive(Clone, Copy)] enum ValueCase { Random, @@ -398,6 +353,6 @@ impl<'a> MigrationGen<'a> { fn choose(&self) -> Option { let candidates = Migration::candidates(self.model); - SchemaDecisions::choose_index(self.rng, candidates.len()).map(|idx| candidates[idx].clone()) + choose_index(self.rng, candidates.len()).map(|idx| candidates[idx].clone()) } } diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs index b21e9f79645..59835aebab4 100644 --- a/crates/dst/src/engine/migrations.rs +++ b/crates/dst/src/engine/migrations.rs @@ -10,17 +10,14 @@ const MAX_TABLE_COLUMNS: usize = 32; const MAX_TABLES: usize = 128; #[derive(Debug, Clone, PartialEq, Eq)] -pub struct Migration { - pub table: usize, - pub ops: Vec, +pub enum Migration { + AddTable { is_event: bool }, + RemoveTable { table: usize }, + AlterTable { table: usize, ops: Vec }, } #[derive(Debug, Clone, PartialEq, Eq)] -pub enum MigrationOp { - AddTable { - is_event: bool, - }, - RemoveTable, +pub enum TableMigrationOp { ChangeAccess, AddColumn { ty: Type, @@ -81,14 +78,8 @@ impl Migration { let non_event_tables = schema.tables.iter().filter(|table| !table.is_event).count(); if schema.tables.len() < MAX_TABLES { - candidates.push(Self::new( - schema.tables.len(), - [MigrationOp::AddTable { is_event: false }], - )); - candidates.push(Self::new( - schema.tables.len(), - [MigrationOp::AddTable { is_event: true }], - )); + candidates.push(Self::AddTable { is_event: false }); + candidates.push(Self::AddTable { is_event: true }); } for (table, table_plan) in schema.tables.iter().enumerate() { @@ -96,43 +87,43 @@ impl Migration { let pristine = row_count == 0 && !model.ever_inserted(table); if (table_plan.is_event && row_count == 0) || (!table_plan.is_event && pristine && non_event_tables > 1) { - candidates.push(Self::new(table, [MigrationOp::RemoveTable])); + candidates.push(Self::RemoveTable { table }); } if table_plan.is_event { if row_count == 0 && table_plan.columns.len() < MAX_EVENT_COLUMNS { - candidates.push(Self::new( + candidates.push(Self::alter_table( table, - [MigrationOp::ChangeAccess, MigrationOp::ReschemaEventTable], + [TableMigrationOp::ChangeAccess, TableMigrationOp::ReschemaEventTable], )); } continue; } if table_plan.columns.len() < MAX_TABLE_COLUMNS { - candidates.extend( - Type::ALL - .iter() - .copied() - .map(|ty| Self::new(table, [MigrationOp::ChangeAccess, MigrationOp::AddColumn { ty }])), - ); + candidates.extend(Type::ALL.iter().copied().map(|ty| { + Self::alter_table( + table, + [TableMigrationOp::ChangeAccess, TableMigrationOp::AddColumn { ty }], + ) + })); } if pristine { if let Some(sequence) = first_addable_sequence(table_plan) { - candidates.push(Self::new(table, [MigrationOp::AddSequence { sequence }])); + candidates.push(Self::alter_table(table, [TableMigrationOp::AddSequence { sequence }])); } if let Some(columns) = first_addable_unique_constraint_columns(table_plan) { let mut ops = Vec::new(); if !has_index(table_plan, &columns) { - ops.push(MigrationOp::AddIndex { + ops.push(TableMigrationOp::AddIndex { columns: columns.clone(), algorithm: IndexAlgorithm::BTree, }); } - ops.push(MigrationOp::AddUniqueConstraint { columns }); - candidates.push(Self::new(table, ops)); + ops.push(TableMigrationOp::AddUniqueConstraint { columns }); + candidates.push(Self::alter_table(table, ops)); } } @@ -140,42 +131,54 @@ impl Migration { if let Some(sequence) = first_addable_sequence_boundary_probe(table_plan, |column| model.column_domain(table, column)) { - candidates.push(Self::new(table, [MigrationOp::AddSequence { sequence }])); + candidates.push(Self::alter_table(table, [TableMigrationOp::AddSequence { sequence }])); } } if let Some(sequence) = first_removable_sequence(table_plan) { - candidates.push(Self::new(table, [MigrationOp::RemoveSequence { sequence }])); + candidates.push(Self::alter_table( + table, + [TableMigrationOp::RemoveSequence { sequence }], + )); } if let Some(constraint) = first_removable_unique_constraint(table_plan) { - candidates.push(Self::new(table, [MigrationOp::RemoveUniqueConstraint { constraint }])); + candidates.push(Self::alter_table( + table, + [TableMigrationOp::RemoveUniqueConstraint { constraint }], + )); } if let Some((columns, algorithm)) = first_addable_index(table_plan) { - candidates.push(Self::new(table, [MigrationOp::AddIndex { columns, algorithm }])); + candidates.push(Self::alter_table( + table, + [TableMigrationOp::AddIndex { columns, algorithm }], + )); } if let Some(index) = first_changeable_index(table_plan) { - candidates.push(Self::new( + candidates.push(Self::alter_table( table, - [MigrationOp::ChangeAccess, MigrationOp::ChangeIndex { index }], + [TableMigrationOp::ChangeAccess, TableMigrationOp::ChangeIndex { index }], )); - candidates.push(Self::new(table, [MigrationOp::RemoveIndex { index }])); + candidates.push(Self::alter_table(table, [TableMigrationOp::RemoveIndex { index }])); } if let Some(column) = first_widenable_sum_column(table_plan) { - candidates.push(Self::new( + candidates.push(Self::alter_table( table, - [MigrationOp::ChangeAccess, MigrationOp::ChangeColumnType { column }], + [ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeColumnType { column }, + ], )); if table_plan.primary_key.is_some() && table_plan.sequences.is_empty() { - candidates.push(Self::new( + candidates.push(Self::alter_table( table, [ - MigrationOp::ChangePrimaryKey { column: None }, - MigrationOp::ChangeColumnType { column }, + TableMigrationOp::ChangePrimaryKey { column: None }, + TableMigrationOp::ChangeColumnType { column }, ], )); } @@ -187,81 +190,60 @@ impl Migration { pub fn apply_to(&self, schema: &SchemaPlan) -> anyhow::Result { let mut next = schema.clone(); - if let Some(is_event) = self.adds_table() { - next.tables.push(new_table(&next.tables, is_event)); - return Ok(next); - } - if self.removes_table() { - if self.table >= next.tables.len() { - anyhow::bail!("remove-table migration references missing table {}", self.table); - } - next.tables.remove(self.table); - return Ok(next); - } - let table = next - .tables - .get_mut(self.table) - .ok_or_else(|| anyhow::anyhow!("migration references missing table {}", self.table))?; - - for op in &self.ops { - apply_op(table, op.clone())?; + match self { + Self::AddTable { is_event } => { + next.tables.push(new_table(&next.tables, *is_event)); + } + Self::RemoveTable { table } => { + anyhow::ensure!( + *table < next.tables.len(), + "remove-table migration references missing table {table}" + ); + next.tables.remove(*table); + } + Self::AlterTable { table, ops } => { + let table_plan = next + .tables + .get_mut(*table) + .ok_or_else(|| anyhow::anyhow!("migration references missing table {table}"))?; + + for op in ops { + apply_table_op(table_plan, op.clone())?; + } + } } Ok(next) } - pub fn adds_table(&self) -> Option { - match self.ops.as_slice() { - [MigrationOp::AddTable { is_event }] => Some(*is_event), - _ => None, - } - } - - pub fn removes_table(&self) -> bool { - matches!(self.ops.as_slice(), [MigrationOp::RemoveTable]) - } - - pub fn added_column_defaults(&self) -> Vec { - self.ops + pub fn added_column_defaults(&self) -> Option<(usize, Vec)> { + let Self::AlterTable { table, ops } = self else { + return None; + }; + let defaults = ops .iter() .filter_map(|op| match op { - MigrationOp::AddColumn { ty } => Some(ty.default_value()), + TableMigrationOp::AddColumn { ty } => Some(ty.default_value()), _ => None, }) - .collect() + .collect::>(); + + (!defaults.is_empty()).then_some((*table, defaults)) } pub fn expected_steps(&self) -> Vec { let mut expected = Vec::new(); - for op in &self.ops { - match op { - MigrationOp::AddTable { .. } => expected.push(ExpectedStep::AddTable), - MigrationOp::RemoveTable => { - expected.push(ExpectedStep::RemoveTable); - expected.push(ExpectedStep::DisconnectAllUsers); - } - MigrationOp::ChangeAccess => expected.push(ExpectedStep::ChangeAccess), - MigrationOp::AddColumn { .. } => { - expected.push(ExpectedStep::AddColumns); - expected.push(ExpectedStep::DisconnectAllUsers); - } - MigrationOp::AddIndex { .. } => expected.push(ExpectedStep::AddIndex), - MigrationOp::RemoveIndex { .. } => expected.push(ExpectedStep::RemoveIndex), - MigrationOp::AddSequence { .. } => expected.push(ExpectedStep::AddSequence), - MigrationOp::RemoveSequence { .. } => expected.push(ExpectedStep::RemoveSequence), - MigrationOp::AddUniqueConstraint { .. } => expected.push(ExpectedStep::AddConstraint), - MigrationOp::RemoveUniqueConstraint { .. } => expected.push(ExpectedStep::RemoveConstraint), - MigrationOp::ChangePrimaryKey { .. } => expected.push(ExpectedStep::ChangePrimaryKey), - MigrationOp::ChangeIndex { .. } => { - expected.push(ExpectedStep::RemoveIndex); - expected.push(ExpectedStep::AddIndex); - } - MigrationOp::ChangeColumnType { .. } => expected.push(ExpectedStep::ChangeColumns), - MigrationOp::ReschemaEventTable => { - expected.push(ExpectedStep::ReschemaEventTable); - expected.push(ExpectedStep::DisconnectAllUsers); + match self { + Self::AddTable { .. } => expected.push(ExpectedStep::AddTable), + Self::RemoveTable { .. } => { + expected.push(ExpectedStep::RemoveTable); + expected.push(ExpectedStep::DisconnectAllUsers); + } + Self::AlterTable { ops, .. } => { + for op in ops { + expected_steps_for_table_op(op, &mut expected); } } } @@ -275,23 +257,43 @@ impl Migration { expected } - fn new(table: usize, ops: impl Into>) -> Self { - Self { table, ops: ops.into() } + fn alter_table(table: usize, ops: impl Into>) -> Self { + Self::AlterTable { table, ops: ops.into() } } } -fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { +fn expected_steps_for_table_op(op: &TableMigrationOp, expected: &mut Vec) { match op { - MigrationOp::AddTable { .. } => { - anyhow::bail!("add-table migration should be applied by SchemaPlan::push before table ops"); + TableMigrationOp::ChangeAccess => expected.push(ExpectedStep::ChangeAccess), + TableMigrationOp::AddColumn { .. } => { + expected.push(ExpectedStep::AddColumns); + expected.push(ExpectedStep::DisconnectAllUsers); + } + TableMigrationOp::AddIndex { .. } => expected.push(ExpectedStep::AddIndex), + TableMigrationOp::RemoveIndex { .. } => expected.push(ExpectedStep::RemoveIndex), + TableMigrationOp::AddSequence { .. } => expected.push(ExpectedStep::AddSequence), + TableMigrationOp::RemoveSequence { .. } => expected.push(ExpectedStep::RemoveSequence), + TableMigrationOp::AddUniqueConstraint { .. } => expected.push(ExpectedStep::AddConstraint), + TableMigrationOp::RemoveUniqueConstraint { .. } => expected.push(ExpectedStep::RemoveConstraint), + TableMigrationOp::ChangePrimaryKey { .. } => expected.push(ExpectedStep::ChangePrimaryKey), + TableMigrationOp::ChangeIndex { .. } => { + expected.push(ExpectedStep::RemoveIndex); + expected.push(ExpectedStep::AddIndex); } - MigrationOp::RemoveTable => { - anyhow::bail!("remove-table migration should be applied by SchemaPlan::remove before table ops"); + TableMigrationOp::ChangeColumnType { .. } => expected.push(ExpectedStep::ChangeColumns), + TableMigrationOp::ReschemaEventTable => { + expected.push(ExpectedStep::ReschemaEventTable); + expected.push(ExpectedStep::DisconnectAllUsers); } - MigrationOp::ChangeAccess => { + } +} + +fn apply_table_op(table: &mut TablePlan, op: TableMigrationOp) -> anyhow::Result<()> { + match op { + TableMigrationOp::ChangeAccess => { table.is_public = !table.is_public; } - MigrationOp::AddColumn { ty } => { + TableMigrationOp::AddColumn { ty } => { anyhow::ensure!(!table.is_event, "add-column migration selected event table"); anyhow::ensure!( table.columns.len() < MAX_TABLE_COLUMNS, @@ -302,7 +304,7 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { ty, }); } - MigrationOp::AddIndex { columns, algorithm } => { + TableMigrationOp::AddIndex { columns, algorithm } => { ensure_columns_exist(table, &columns)?; anyhow::ensure!( !has_index(table, &columns), @@ -310,7 +312,7 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { ); table.indexes.push(IndexPlan { columns, algorithm }); } - MigrationOp::RemoveIndex { index } => { + TableMigrationOp::RemoveIndex { index } => { let Some(index_plan) = table.indexes.get(index) else { anyhow::bail!("remove-index migration references missing index {index}"); }; @@ -320,7 +322,7 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { ); table.indexes.remove(index); } - MigrationOp::AddSequence { sequence } => { + TableMigrationOp::AddSequence { sequence } => { let column = sequence.column; let Some(column_plan) = table.columns.get(column) else { anyhow::bail!("add-sequence migration references missing column {column}"); @@ -335,14 +337,14 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { ); table.sequences.push(sequence); } - MigrationOp::RemoveSequence { sequence } => { + TableMigrationOp::RemoveSequence { sequence } => { anyhow::ensure!( sequence < table.sequences.len(), "remove-sequence migration references missing sequence" ); table.sequences.remove(sequence); } - MigrationOp::AddUniqueConstraint { columns } => { + TableMigrationOp::AddUniqueConstraint { columns } => { ensure_columns_exist(table, &columns)?; anyhow::ensure!( !table @@ -357,7 +359,7 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { ); table.unique_constraints.push(UniqueConstraintPlan { columns }); } - MigrationOp::RemoveUniqueConstraint { constraint } => { + TableMigrationOp::RemoveUniqueConstraint { constraint } => { let Some(constraint_plan) = table.unique_constraints.get(constraint) else { anyhow::bail!("remove-constraint migration references missing constraint {constraint}"); }; @@ -369,7 +371,7 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { ); table.unique_constraints.remove(constraint); } - MigrationOp::ChangePrimaryKey { column } => { + TableMigrationOp::ChangePrimaryKey { column } => { if let Some(column) = column { anyhow::ensure!( column < table.columns.len(), @@ -378,7 +380,7 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { } table.primary_key = column; } - MigrationOp::ChangeIndex { index } => { + TableMigrationOp::ChangeIndex { index } => { let Some(index_plan) = table.indexes.get_mut(index) else { anyhow::bail!("index migration references missing index {index}"); }; @@ -387,10 +389,10 @@ fn apply_op(table: &mut TablePlan, op: MigrationOp) -> anyhow::Result<()> { IndexAlgorithm::Hash => IndexAlgorithm::BTree, }; } - MigrationOp::ChangeColumnType { column } => { + TableMigrationOp::ChangeColumnType { column } => { widen_sum_column(table, Some(column))?; } - MigrationOp::ReschemaEventTable => { + TableMigrationOp::ReschemaEventTable => { anyhow::ensure!(table.is_event, "event-table migration selected non-event table"); anyhow::ensure!( table.columns.len() < MAX_EVENT_COLUMNS, diff --git a/crates/dst/src/engine/model.rs b/crates/dst/src/engine/model.rs index 834deca3889..5dfcd068225 100644 --- a/crates/dst/src/engine/model.rs +++ b/crates/dst/src/engine/model.rs @@ -1,7 +1,9 @@ use spacetimedb_lib::AlgebraicValue; +use super::migrations::Migration; +use super::row::{normalize_rows, Row}; use super::state::{schema_state_for_plan, CommitDelta, CountState, TableDelta, TableRowCount, TableRows}; -use super::workload::{normalize_rows, InsertOutcome, Interaction, Observation, Row}; +use super::workload::{InsertOutcome, Interaction, Observation}; use crate::schema::{SchemaPlan, Type}; #[derive(Debug)] @@ -187,24 +189,28 @@ impl Model { } Interaction::Migrate(migration) => { debug_assert!(self.pending_tx.is_none()); - let adds_table = migration.adds_table().is_some(); - let removes_table = migration.removes_table(); let added_defaults = migration.added_column_defaults(); self.schema = migration .apply_to(&self.schema) .expect("generated migrations must be valid for the model schema"); - if adds_table { - self.committed_tables.push(TableState { - rows: vec![], - ever_inserted: false, - }); - } else if removes_table { - self.committed_tables.remove(migration.table); - } else if !added_defaults.is_empty() { - for row in &mut self.committed_tables[migration.table].rows { - let mut elements = row.elements.to_vec(); - elements.extend(added_defaults.iter().cloned()); - row.elements = elements.into_boxed_slice(); + match migration { + Migration::AddTable { .. } => { + self.committed_tables.push(TableState { + rows: vec![], + ever_inserted: false, + }); + } + Migration::RemoveTable { table } => { + self.committed_tables.remove(*table); + } + Migration::AlterTable { .. } => { + if let Some((table, defaults)) = added_defaults { + for row in &mut self.committed_tables[table].rows { + let mut elements = row.elements.to_vec(); + elements.extend(defaults.iter().cloned()); + row.elements = elements.into_boxed_slice(); + } + } } } Observation::Migrated diff --git a/crates/dst/src/engine/properties.rs b/crates/dst/src/engine/properties.rs index 667eec09510..9511c37de48 100644 --- a/crates/dst/src/engine/properties.rs +++ b/crates/dst/src/engine/properties.rs @@ -1,5 +1,6 @@ use super::model::Model; -use super::workload::{InsertOutcome, Interaction, Observation, Row}; +use super::row::Row; +use super::workload::{InsertOutcome, Interaction, Observation}; use crate::schema::SchemaPlan; use crate::traits::Properties; diff --git a/crates/dst/src/engine/row.rs b/crates/dst/src/engine/row.rs new file mode 100644 index 00000000000..a3008d48d42 --- /dev/null +++ b/crates/dst/src/engine/row.rs @@ -0,0 +1,13 @@ +use spacetimedb_lib::bsatn::to_vec; +use spacetimedb_lib::ProductValue; + +pub type Row = ProductValue; + +pub fn row_to_bytes(row: &Row) -> Vec { + to_vec(row).expect("row serialization must not fail") +} + +pub fn normalize_rows(mut rows: Vec) -> Vec { + rows.sort_by_key(row_to_bytes); + rows +} diff --git a/crates/dst/src/engine/state.rs b/crates/dst/src/engine/state.rs index 35512196024..db579034ebb 100644 --- a/crates/dst/src/engine/state.rs +++ b/crates/dst/src/engine/state.rs @@ -1,7 +1,13 @@ +use spacetimedb_lib::db::auth::StAccess; use spacetimedb_lib::AlgebraicType; +use spacetimedb_primitives::ColId; +use spacetimedb_schema::def::IndexAlgorithm as SchemaIndexAlgorithm; +use spacetimedb_schema::schema::TableSchema; -use super::workload::Row; -use crate::schema::{IndexAlgorithm, SchemaPlan}; +use super::row::Row; +use crate::schema::{ + IndexAlgorithm as PlanIndexAlgorithm, IndexPlan, SchemaPlan, SequencePlan, TablePlan, UniqueConstraintPlan, +}; #[derive(Debug, Clone, PartialEq, Eq)] pub struct CountState { @@ -83,63 +89,149 @@ pub struct TableDelta { pub truncated: bool, } +impl From for IndexAlgorithmState { + fn from(algorithm: PlanIndexAlgorithm) -> Self { + match algorithm { + PlanIndexAlgorithm::BTree => Self::BTree, + PlanIndexAlgorithm::Hash => Self::Hash, + } + } +} + +impl From<&SchemaIndexAlgorithm> for IndexAlgorithmState { + fn from(algorithm: &SchemaIndexAlgorithm) -> Self { + match algorithm { + SchemaIndexAlgorithm::BTree(_) => Self::BTree, + SchemaIndexAlgorithm::Hash(_) => Self::Hash, + SchemaIndexAlgorithm::Direct(_) => Self::Direct, + _ => Self::Unknown, + } + } +} + +impl IndexState { + fn from_plan(index: &IndexPlan) -> Self { + Self { + columns: index.columns.clone(), + algorithm: index.algorithm.into(), + } + } + + fn from_schema(algorithm: &SchemaIndexAlgorithm) -> Self { + Self { + columns: schema_index_columns(algorithm), + algorithm: algorithm.into(), + } + } +} + +impl UniqueConstraintState { + fn from_plan(constraint: &UniqueConstraintPlan) -> Self { + Self { + columns: constraint.columns.clone(), + } + } + + fn from_schema_columns(columns: impl IntoIterator) -> Self { + Self { + columns: columns.into_iter().map(|col| col.0 as usize).collect(), + } + } +} + +impl SequenceState { + fn from_plan(sequence: &SequencePlan) -> Self { + Self { + column: sequence.column, + } + } + + fn from_schema_column(column: ColId) -> Self { + Self { + column: column.0 as usize, + } + } +} + pub fn schema_state_for_plan(schema: &SchemaPlan) -> SchemaState { SchemaState { tables: schema .tables .iter() .enumerate() - .map(|(table, table_plan)| { - let mut indexes = table_plan - .indexes - .iter() - .map(|index| IndexState { - columns: index.columns.clone(), - algorithm: match index.algorithm { - IndexAlgorithm::BTree => IndexAlgorithmState::BTree, - IndexAlgorithm::Hash => IndexAlgorithmState::Hash, - }, - }) - .collect::>(); - indexes.sort(); - - let mut unique_constraints = table_plan - .unique_constraints - .iter() - .map(|constraint| UniqueConstraintState { - columns: constraint.columns.clone(), - }) - .collect::>(); - unique_constraints.sort(); - - let mut sequences = table_plan - .sequences - .iter() - .map(|sequence| SequenceState { - column: sequence.column, - }) - .collect::>(); - sequences.sort(); - - TableSchemaState { - table, - name: table_plan.name.clone(), - is_public: table_plan.is_public, - is_event: table_plan.is_event, - primary_key: table_plan.primary_key, - columns: table_plan - .columns - .iter() - .map(|column| ColumnState { - name: column.name.clone(), - ty: column.ty.to_algebraic(), - }) - .collect(), - indexes, - unique_constraints, - sequences, - } + .map(|(table, table_plan)| table_schema_state_for_plan(table, table_plan)) + .collect(), + } +} + +pub fn table_schema_state_for_schema(table: usize, schema: &TableSchema) -> TableSchemaState { + TableSchemaState { + table, + name: schema.table_name.to_string(), + is_public: schema.table_access == StAccess::Public, + is_event: schema.is_event, + primary_key: schema.primary_key.map(|col| col.0 as usize), + columns: schema + .columns + .iter() + .map(|column| ColumnState { + name: column.col_name.to_string(), + ty: column.col_type.clone(), + }) + .collect(), + indexes: sorted( + schema + .indexes + .iter() + .map(|index| IndexState::from_schema(&index.index_algorithm)), + ), + unique_constraints: sorted(schema.constraints.iter().filter_map(|constraint| { + constraint + .data + .unique_columns() + .map(|columns| UniqueConstraintState::from_schema_columns(columns.iter())) + })), + sequences: sorted( + schema + .sequences + .iter() + .map(|sequence| SequenceState::from_schema_column(sequence.col_pos)), + ), + } +} + +fn table_schema_state_for_plan(table: usize, table_plan: &TablePlan) -> TableSchemaState { + TableSchemaState { + table, + name: table_plan.name.clone(), + is_public: table_plan.is_public, + is_event: table_plan.is_event, + primary_key: table_plan.primary_key, + columns: table_plan + .columns + .iter() + .map(|column| ColumnState { + name: column.name.clone(), + ty: column.ty.to_algebraic(), }) .collect(), + indexes: sorted(table_plan.indexes.iter().map(IndexState::from_plan)), + unique_constraints: sorted( + table_plan + .unique_constraints + .iter() + .map(UniqueConstraintState::from_plan), + ), + sequences: sorted(table_plan.sequences.iter().map(SequenceState::from_plan)), } } + +fn schema_index_columns(algorithm: &SchemaIndexAlgorithm) -> Vec { + algorithm.columns().iter().map(|col| col.0 as usize).collect() +} + +fn sorted(values: impl IntoIterator) -> Vec { + let mut values = values.into_iter().collect::>(); + values.sort(); + values +} diff --git a/crates/dst/src/engine/workload.rs b/crates/dst/src/engine/workload.rs index 20bc2678f3d..f4fa39f6579 100644 --- a/crates/dst/src/engine/workload.rs +++ b/crates/dst/src/engine/workload.rs @@ -1,16 +1,14 @@ use std::fmt::{Debug, Error, Formatter}; -use super::generation::{pick_weighted, GenCtx}; +use super::generation::GenCtx; use super::migrations::Migration; use super::model::Model; +use super::row::Row; use super::state::{CommitDelta, CountState}; +use crate::rng::pick_weighted; use crate::schema::SchemaPlan; -use spacetimedb_lib::bsatn::to_vec; -use spacetimedb_lib::ProductValue; use spacetimedb_runtime::sim::Rng; -pub type Row = ProductValue; - #[derive(Debug, Clone)] pub enum Interaction { BeginMutTx, @@ -267,12 +265,3 @@ impl Iterator for WorkloadGen { Some(self.next_interaction()) } } - -pub fn row_to_bytes(row: &Row) -> Vec { - to_vec(row).expect("row serialization must not fail") -} - -pub fn normalize_rows(mut rows: Vec) -> Vec { - rows.sort_by_key(row_to_bytes); - rows -} diff --git a/crates/dst/src/lib.rs b/crates/dst/src/lib.rs index 8d12c575e4c..ef8e754f6de 100644 --- a/crates/dst/src/lib.rs +++ b/crates/dst/src/lib.rs @@ -1,4 +1,5 @@ pub mod engine; +mod rng; pub mod schema; pub mod sim; pub mod traits; diff --git a/crates/dst/src/rng.rs b/crates/dst/src/rng.rs new file mode 100644 index 00000000000..5f471105321 --- /dev/null +++ b/crates/dst/src/rng.rs @@ -0,0 +1,58 @@ +use spacetimedb_runtime::sim::Rng; + +#[derive(Clone, Copy)] +pub(crate) struct Choice { + weight: u64, + value: T, +} + +pub(crate) const fn choice(weight: u64, value: T) -> Choice { + Choice { weight, value } +} + +pub(crate) fn frequency(rng: &Rng, choices: &[Choice]) -> T { + let total: u64 = choices.iter().map(|choice| choice.weight).sum(); + + assert!(total > 0, "at least one choice weight must be non-zero"); + + let mut selected = rng.next_u64() % total; + + for choice in choices.iter().copied() { + if selected < choice.weight { + return choice.value; + } + + selected -= choice.weight; + } + + unreachable!("selected value is always inside total weight") +} + +pub(crate) fn pick_weighted(rng: &Rng, weights: &[u64]) -> usize { + let total: u64 = weights.iter().sum(); + + assert!(total > 0, "at least one weight must be non-zero"); + + let mut selected = rng.next_u64() % total; + + for (idx, weight) in weights.iter().copied().enumerate() { + if selected < weight { + return idx; + } + + selected -= weight; + } + + unreachable!("selected value is always inside total weight") +} + +pub(crate) fn choose_index(rng: &Rng, len: usize) -> Option { + (len > 0).then(|| rng.index(len)) +} + +pub(crate) fn range_inclusive(rng: &Rng, lo: usize, hi: usize) -> usize { + if lo >= hi { + return lo; + } + lo + (rng.next_u64() as usize % (hi - lo + 1)) +} diff --git a/crates/dst/src/schema.rs b/crates/dst/src/schema.rs index 7f30e713973..cdfc7e166aa 100644 --- a/crates/dst/src/schema.rs +++ b/crates/dst/src/schema.rs @@ -1,3 +1,4 @@ +use crate::rng; use spacetimedb_lib::db::raw_def::v10::*; use spacetimedb_lib::db::raw_def::v9::{RawIndexAlgorithm, TableAccess, TableType}; use spacetimedb_primitives::{ColId, ColList}; @@ -68,18 +69,15 @@ pub struct SchemaDecisions; impl SchemaDecisions { pub fn range(rng: &Rng, (lo, hi): (usize, usize)) -> usize { - if lo >= hi { - return lo; - } - lo + (rng.next_u64() as usize % (hi - lo + 1)) + rng::range_inclusive(rng, lo, hi) } pub fn index(rng: &Rng, len: usize) -> usize { - rng.index(len) + rng::choose_index(rng, len).expect("len must be non-zero") } pub fn choose_index(rng: &Rng, len: usize) -> Option { - (len > 0).then(|| Self::index(rng, len)) + rng::choose_index(rng, len) } pub fn sample_probability(rng: &Rng, probability: f64) -> bool { From b14c6f00de8dea07fc6e76278f75f71b966db765 Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Tue, 14 Jul 2026 14:46:17 +0530 Subject: [PATCH 05/10] cleanups in migrations --- crates/dst/src/engine.rs | 113 +++----- crates/dst/src/engine/generation.rs | 110 ++++--- crates/dst/src/engine/migrations.rs | 428 ++++++++++++++-------------- crates/dst/src/engine/model.rs | 118 ++++++-- crates/dst/src/rng.rs | 8 + crates/dst/src/schema.rs | 10 +- 6 files changed, 430 insertions(+), 357 deletions(-) diff --git a/crates/dst/src/engine.rs b/crates/dst/src/engine.rs index 87a4452c6b6..fd0822ff1aa 100644 --- a/crates/dst/src/engine.rs +++ b/crates/dst/src/engine.rs @@ -12,7 +12,7 @@ use spacetimedb_lib::{Identity, RawModuleDef}; use spacetimedb_primitives::TableId; use spacetimedb_runtime::sim::{Rng, Runtime as SimRuntime}; use spacetimedb_runtime::Handle; -use spacetimedb_schema::auto_migrate::{ponder_migrate, AutoMigrateStep, MigratePlan}; +use spacetimedb_schema::auto_migrate::{ponder_migrate, MigratePlan}; use spacetimedb_schema::def::ModuleDef; use spacetimedb_schema::schema::{Schema, TableSchema}; use spacetimedb_table::page_pool::PagePool; @@ -25,7 +25,7 @@ mod row; mod state; mod workload; -use self::migrations::{ExpectedStep, Migration}; +use self::migrations::Migration; use self::row::{normalize_rows, row_to_bytes}; use self::state::{ table_schema_state_for_schema, CommitDelta, CountState, SchemaState, TableDelta, TableRowCount, TableRows, @@ -223,12 +223,15 @@ impl EngineTarget { self.active_mut_tx.is_none(), "migration while mutable transaction is active" ); + anyhow::ensure!( + migration.schema() != &self.schema, + "engine DST generated a no-op migration" + ); - let new_schema = migration.apply_to(&self.schema)?; let old_module_def = Self::module_def(&self.schema)?; - let new_module_def = Self::module_def(&new_schema)?; + let new_module_def = Self::module_def(migration.schema())?; let plan = ponder_migrate(&old_module_def, &new_module_def)?; - self.ensure_expected_plan(migration, &plan)?; + ensure_auto_plan(&plan)?; let db = self .db @@ -240,33 +243,11 @@ impl EngineTarget { anyhow::bail!("migration commit produced no transaction data"); }; - self.schema = new_schema; + self.schema = migration.schema().clone(); self.table_ids = Self::load_table_ids(db, &self.schema)?; Ok(()) } - fn ensure_expected_plan(&self, migration: &Migration, plan: &MigratePlan<'_>) -> anyhow::Result<()> { - let MigratePlan::Auto(plan) = plan else { - anyhow::bail!("engine DST generated a manual migration plan"); - }; - - let mut actual = plan - .steps - .iter() - .map(expected_step_from_auto_step) - .collect::>>()?; - actual.sort(); - - let mut expected = migration.expected_steps(); - expected.sort(); - - anyhow::ensure!( - actual == expected, - "engine DST generated unexpected migration steps: actual={actual:?}, expected={expected:?}" - ); - Ok(()) - } - pub fn execute(&mut self, interaction: &Interaction) -> anyhow::Result { tracing::debug!(?interaction, "executing interaction"); @@ -355,24 +336,11 @@ impl EngineTarget { } } -fn expected_step_from_auto_step(step: &AutoMigrateStep<'_>) -> anyhow::Result { - match step { - AutoMigrateStep::AddTable(_) => Ok(ExpectedStep::AddTable), - AutoMigrateStep::RemoveTable(_) => Ok(ExpectedStep::RemoveTable), - AutoMigrateStep::AddColumns(_) => Ok(ExpectedStep::AddColumns), - AutoMigrateStep::AddIndex(_) => Ok(ExpectedStep::AddIndex), - AutoMigrateStep::RemoveIndex(_) => Ok(ExpectedStep::RemoveIndex), - AutoMigrateStep::AddSequence(_) => Ok(ExpectedStep::AddSequence), - AutoMigrateStep::RemoveSequence(_) => Ok(ExpectedStep::RemoveSequence), - AutoMigrateStep::AddConstraint(_) => Ok(ExpectedStep::AddConstraint), - AutoMigrateStep::RemoveConstraint(_) => Ok(ExpectedStep::RemoveConstraint), - AutoMigrateStep::ChangeAccess(_) => Ok(ExpectedStep::ChangeAccess), - AutoMigrateStep::ChangePrimaryKey(_) => Ok(ExpectedStep::ChangePrimaryKey), - AutoMigrateStep::ChangeColumns(_) => Ok(ExpectedStep::ChangeColumns), - AutoMigrateStep::ReschemaEventTable(_) => Ok(ExpectedStep::ReschemaEventTable), - AutoMigrateStep::DisconnectAllUsers => Ok(ExpectedStep::DisconnectAllUsers), - step => anyhow::bail!("engine DST generated unsupported migration step: {step:?}"), - } +fn ensure_auto_plan(plan: &MigratePlan<'_>) -> anyhow::Result<()> { + let MigratePlan::Auto(_) = plan else { + anyhow::bail!("engine DST generated a manual migration plan"); + }; + Ok(()) } struct DstUpdateLogger; @@ -421,7 +389,7 @@ mod tests { use spacetimedb_runtime::sim::Runtime as SimRuntime; use spacetimedb_sats::product; - use super::migrations::{Migration, TableMigrationOp}; + use super::migrations::{Migration, SchemaRewrite, TableMigrationOp}; use super::*; use crate::schema::{ColumnPlan, IndexAlgorithm, IndexPlan, TablePlan, Type, UniqueConstraintPlan}; @@ -536,13 +504,16 @@ mod tests { let mut target = EngineTarget::init(add_column_replay_schema(), 0)?; insert_u64_rows(&mut target)?; - target.execute(&Interaction::Migrate(Migration::AlterTable { - table: 0, - ops: vec![ - TableMigrationOp::ChangeAccess, - TableMigrationOp::AddColumn { ty: Type::U64 }, - ], - }))?; + target.execute(&Interaction::Migrate(Migration::from_rewrites( + &target.schema, + vec![SchemaRewrite::AlterTable { + table: "items".into(), + ops: vec![ + TableMigrationOp::ChangeAccess, + TableMigrationOp::AddColumn { ty: Type::U64 }, + ], + }], + )?))?; target.execute(&Interaction::Replay)?; Ok(()) @@ -553,13 +524,16 @@ mod tests { let mut target = EngineTarget::init(change_index_replay_schema(), 0)?; insert_u64_rows(&mut target)?; - target.execute(&Interaction::Migrate(Migration::AlterTable { - table: 0, - ops: vec![ - TableMigrationOp::ChangeAccess, - TableMigrationOp::ChangeIndex { index: 1 }, - ], - }))?; + target.execute(&Interaction::Migrate(Migration::from_rewrites( + &target.schema, + vec![SchemaRewrite::AlterTable { + table: "items".into(), + ops: vec![ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeIndex { columns: vec![1] }, + ], + }], + )?))?; target.execute(&Interaction::Replay)?; Ok(()) @@ -578,13 +552,16 @@ mod tests { } target.execute(&Interaction::CommitTx)?; - target.execute(&Interaction::Migrate(Migration::AlterTable { - table: 0, - ops: vec![ - TableMigrationOp::ChangeAccess, - TableMigrationOp::ChangeColumnType { column: 1 }, - ], - }))?; + target.execute(&Interaction::Migrate(Migration::from_rewrites( + &target.schema, + vec![SchemaRewrite::AlterTable { + table: "items".into(), + ops: vec![ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeColumnType { column: 1 }, + ], + }], + )?))?; target.execute(&Interaction::Replay)?; Ok(()) diff --git a/crates/dst/src/engine/generation.rs b/crates/dst/src/engine/generation.rs index d81f7237628..d46b3386d4e 100644 --- a/crates/dst/src/engine/generation.rs +++ b/crates/dst/src/engine/generation.rs @@ -5,7 +5,7 @@ use spacetimedb_sats::ArrayValue; use super::migrations::Migration; use super::model::{ColumnDomain, Model}; use super::row::Row; -use crate::rng::{choice, choose_index, frequency}; +use crate::rng::{choice, choose_index, Choice, WeightedChoice}; use crate::schema::Type; pub(crate) struct GenCtx<'a> { @@ -37,6 +37,17 @@ enum ValueCase { Weird, } +impl WeightedChoice for ValueCase { + const CHOICES: &'static [Choice] = &[ + choice(45, Self::Random), + choice(15, Self::Small), + choice(15, Self::Edge), + choice(10, Self::NearExisting), + choice(10, Self::Existing), + choice(5, Self::Weird), + ]; +} + #[derive(Clone, Copy)] enum FreshValueCase { Random, @@ -45,6 +56,15 @@ enum FreshValueCase { Weird, } +impl WeightedChoice for FreshValueCase { + const CHOICES: &'static [Choice] = &[ + choice(50, Self::Random), + choice(20, Self::Small), + choice(20, Self::Edge), + choice(10, Self::Weird), + ]; +} + #[derive(Clone, Copy)] enum I64Case { Random, @@ -70,6 +90,19 @@ enum StringCase { Long, } +impl StringCase { + const WEIRD_CHOICES: &'static [Choice] = &[ + choice(35, Self::SqlEscaped), + choice(25, Self::NullByte), + choice(25, Self::OrderedPrefix), + choice(15, Self::Empty), + ]; + + fn pick_weird(rng: &Rng) -> Self { + crate::rng::frequency(rng, Self::WEIRD_CHOICES) + } +} + #[derive(Clone, Copy)] enum BytesCase { Random, @@ -80,6 +113,20 @@ enum BytesCase { Alternating, } +impl BytesCase { + const WEIRD_CHOICES: &'static [Choice] = &[ + choice(25, Self::Empty), + choice(20, Self::RepeatedZero), + choice(20, Self::RepeatedMax), + choice(20, Self::Alternating), + choice(15, Self::Small), + ]; + + fn pick_weird(rng: &Rng) -> Self { + crate::rng::frequency(rng, Self::WEIRD_CHOICES) + } +} + struct ValueGen<'a> { rng: &'a Rng, model: &'a Model, @@ -115,17 +162,7 @@ impl<'a> ValueGen<'a> { } fn gen_value(&self, domain: &ColumnDomain) -> AlgebraicValue { - match frequency( - self.rng, - &[ - choice(45, ValueCase::Random), - choice(15, ValueCase::Small), - choice(15, ValueCase::Edge), - choice(10, ValueCase::NearExisting), - choice(10, ValueCase::Existing), - choice(5, ValueCase::Weird), - ], - ) { + match ValueCase::pick(self.rng) { ValueCase::Random => self.gen_random_value(domain.ty), ValueCase::Small => self.gen_small_value(domain.ty), ValueCase::Edge => self.gen_edge_value(domain.ty), @@ -151,15 +188,7 @@ impl<'a> ValueGen<'a> { } fn gen_fresh_candidate(&self, ty: Type) -> AlgebraicValue { - match frequency( - self.rng, - &[ - choice(50, FreshValueCase::Random), - choice(20, FreshValueCase::Small), - choice(20, FreshValueCase::Edge), - choice(10, FreshValueCase::Weird), - ], - ) { + match FreshValueCase::pick(self.rng) { FreshValueCase::Random => self.gen_random_value(ty), FreshValueCase::Small => self.gen_small_value(ty), FreshValueCase::Edge => self.gen_edge_value(ty), @@ -238,28 +267,11 @@ impl<'a> ValueGen<'a> { fn gen_weird_value(&self, ty: Type) -> AlgebraicValue { match ty { Type::String => { - let case = frequency( - self.rng, - &[ - choice(35, StringCase::SqlEscaped), - choice(25, StringCase::NullByte), - choice(25, StringCase::OrderedPrefix), - choice(15, StringCase::Empty), - ], - ); + let case = StringCase::pick_weird(self.rng); AlgebraicValue::String(self.gen_string_value(case).into()) } Type::Bytes => { - let case = frequency( - self.rng, - &[ - choice(25, BytesCase::Empty), - choice(20, BytesCase::RepeatedZero), - choice(20, BytesCase::RepeatedMax), - choice(20, BytesCase::Alternating), - choice(15, BytesCase::Small), - ], - ); + let case = BytesCase::pick_weird(self.rng); AlgebraicValue::Array(ArrayValue::U8(self.gen_bytes_value(case).into())) } _ => self.gen_edge_value(ty), @@ -352,7 +364,21 @@ impl<'a> MigrationGen<'a> { } fn choose(&self) -> Option { - let candidates = Migration::candidates(self.model); - choose_index(self.rng, candidates.len()).map(|idx| candidates[idx].clone()) + let original = self.model.schema(); + let mut schema = original.clone(); + let steps = 1 + self.rng.index(10); + + for _ in 0..steps { + let candidates = Migration::candidates(&schema, self.model); + let Some(idx) = choose_index(self.rng, candidates.len()) else { + break; + }; + let rewrite = candidates[idx].clone(); + rewrite + .apply_to(&mut schema) + .expect("generated rewrite must be valid for the draft schema"); + } + + (schema != *original).then(|| Migration::from_schema(schema)) } } diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs index 59835aebab4..73263545858 100644 --- a/crates/dst/src/engine/migrations.rs +++ b/crates/dst/src/engine/migrations.rs @@ -2,7 +2,6 @@ use super::model::{ColumnDomain, Model}; use crate::schema::{ ColumnPlan, IndexAlgorithm, IndexPlan, SchemaNames, SchemaPlan, SequencePlan, TablePlan, Type, UniqueConstraintPlan, }; -use spacetimedb_lib::AlgebraicValue; const MAX_SUM_VARIANTS: u8 = 32; const MAX_EVENT_COLUMNS: usize = 32; @@ -10,10 +9,15 @@ const MAX_TABLE_COLUMNS: usize = 32; const MAX_TABLES: usize = 128; #[derive(Debug, Clone, PartialEq, Eq)] -pub enum Migration { +pub struct Migration { + schema: SchemaPlan, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SchemaRewrite { AddTable { is_event: bool }, - RemoveTable { table: usize }, - AlterTable { table: usize, ops: Vec }, + RemoveTable { table: String }, + AlterTable { table: String, ops: Vec }, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -27,25 +31,25 @@ pub enum TableMigrationOp { algorithm: IndexAlgorithm, }, RemoveIndex { - index: usize, + columns: Vec, }, AddSequence { sequence: SequencePlan, }, RemoveSequence { - sequence: usize, + column: usize, }, AddUniqueConstraint { columns: Vec, }, RemoveUniqueConstraint { - constraint: usize, + columns: Vec, }, ChangePrimaryKey { column: Option, }, ChangeIndex { - index: usize, + columns: Vec, }, ChangeColumnType { column: usize, @@ -53,47 +57,47 @@ pub enum TableMigrationOp { ReschemaEventTable, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -pub enum ExpectedStep { - AddTable, - RemoveTable, - AddColumns, - AddIndex, - RemoveIndex, - AddSequence, - RemoveSequence, - AddConstraint, - RemoveConstraint, - ChangeAccess, - ChangePrimaryKey, - ChangeColumns, - ReschemaEventTable, - DisconnectAllUsers, -} - impl Migration { - pub(crate) fn candidates(model: &Model) -> Vec { - let schema = model.schema(); + pub(crate) fn from_schema(schema: SchemaPlan) -> Self { + Self { schema } + } + + #[cfg(test)] + pub(crate) fn from_rewrites(base: &SchemaPlan, rewrites: Vec) -> anyhow::Result { + let mut schema = base.clone(); + for rewrite in &rewrites { + rewrite.apply_to(&mut schema)?; + } + Ok(Self::from_schema(schema)) + } + + pub(crate) fn schema(&self) -> &SchemaPlan { + &self.schema + } + + pub(crate) fn candidates(schema: &SchemaPlan, model: &Model) -> Vec { let mut candidates = Vec::new(); let non_event_tables = schema.tables.iter().filter(|table| !table.is_event).count(); if schema.tables.len() < MAX_TABLES { - candidates.push(Self::AddTable { is_event: false }); - candidates.push(Self::AddTable { is_event: true }); + candidates.push(SchemaRewrite::AddTable { is_event: false }); + candidates.push(SchemaRewrite::AddTable { is_event: true }); } - for (table, table_plan) in schema.tables.iter().enumerate() { - let row_count = model.row_count(table); - let pristine = row_count == 0 && !model.ever_inserted(table); + for table_plan in &schema.tables { + let row_count = model.row_count_by_table_name(&table_plan.name); + let pristine = row_count == 0 && !model.ever_inserted_by_table_name(&table_plan.name); if (table_plan.is_event && row_count == 0) || (!table_plan.is_event && pristine && non_event_tables > 1) { - candidates.push(Self::RemoveTable { table }); + candidates.push(SchemaRewrite::RemoveTable { + table: table_plan.name.clone(), + }); } if table_plan.is_event { if row_count == 0 && table_plan.columns.len() < MAX_EVENT_COLUMNS { - candidates.push(Self::alter_table( - table, + candidates.push(SchemaRewrite::alter_table( + table_plan, [TableMigrationOp::ChangeAccess, TableMigrationOp::ReschemaEventTable], )); } @@ -102,19 +106,22 @@ impl Migration { if table_plan.columns.len() < MAX_TABLE_COLUMNS { candidates.extend(Type::ALL.iter().copied().map(|ty| { - Self::alter_table( - table, + SchemaRewrite::alter_table( + table_plan, [TableMigrationOp::ChangeAccess, TableMigrationOp::AddColumn { ty }], ) })); } if pristine { - if let Some(sequence) = first_addable_sequence(table_plan) { - candidates.push(Self::alter_table(table, [TableMigrationOp::AddSequence { sequence }])); + for sequence in addable_sequences(table_plan) { + candidates.push(SchemaRewrite::alter_table( + table_plan, + [TableMigrationOp::AddSequence { sequence }], + )); } - if let Some(columns) = first_addable_unique_constraint_columns(table_plan) { + for columns in addable_unique_constraint_columns(table_plan) { let mut ops = Vec::new(); if !has_index(table_plan, &columns) { ops.push(TableMigrationOp::AddIndex { @@ -123,50 +130,61 @@ impl Migration { }); } ops.push(TableMigrationOp::AddUniqueConstraint { columns }); - candidates.push(Self::alter_table(table, ops)); + candidates.push(SchemaRewrite::alter_table(table_plan, ops)); } } if row_count > 0 { - if let Some(sequence) = - first_addable_sequence_boundary_probe(table_plan, |column| model.column_domain(table, column)) + for sequence in + addable_sequence_boundary_probes(table_plan, |column| column_domain(model, table_plan, column)) { - candidates.push(Self::alter_table(table, [TableMigrationOp::AddSequence { sequence }])); + candidates.push(SchemaRewrite::alter_table( + table_plan, + [TableMigrationOp::AddSequence { sequence }], + )); } } - if let Some(sequence) = first_removable_sequence(table_plan) { - candidates.push(Self::alter_table( - table, - [TableMigrationOp::RemoveSequence { sequence }], + for column in removable_sequence_columns(table_plan) { + candidates.push(SchemaRewrite::alter_table( + table_plan, + [TableMigrationOp::RemoveSequence { column }], )); } - if let Some(constraint) = first_removable_unique_constraint(table_plan) { - candidates.push(Self::alter_table( - table, - [TableMigrationOp::RemoveUniqueConstraint { constraint }], + for columns in removable_unique_constraint_columns(table_plan) { + candidates.push(SchemaRewrite::alter_table( + table_plan, + [TableMigrationOp::RemoveUniqueConstraint { columns }], )); } - if let Some((columns, algorithm)) = first_addable_index(table_plan) { - candidates.push(Self::alter_table( - table, + for (columns, algorithm) in addable_indexes(table_plan) { + candidates.push(SchemaRewrite::alter_table( + table_plan, [TableMigrationOp::AddIndex { columns, algorithm }], )); } - if let Some(index) = first_changeable_index(table_plan) { - candidates.push(Self::alter_table( - table, - [TableMigrationOp::ChangeAccess, TableMigrationOp::ChangeIndex { index }], + for columns in changeable_index_columns(table_plan) { + candidates.push(SchemaRewrite::alter_table( + table_plan, + [ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeIndex { + columns: columns.clone(), + }, + ], + )); + candidates.push(SchemaRewrite::alter_table( + table_plan, + [TableMigrationOp::RemoveIndex { columns }], )); - candidates.push(Self::alter_table(table, [TableMigrationOp::RemoveIndex { index }])); } - if let Some(column) = first_widenable_sum_column(table_plan) { - candidates.push(Self::alter_table( - table, + for column in widenable_sum_columns(table_plan) { + candidates.push(SchemaRewrite::alter_table( + table_plan, [ TableMigrationOp::ChangeAccess, TableMigrationOp::ChangeColumnType { column }, @@ -174,8 +192,8 @@ impl Migration { )); if table_plan.primary_key.is_some() && table_plan.sequences.is_empty() { - candidates.push(Self::alter_table( - table, + candidates.push(SchemaRewrite::alter_table( + table_plan, [ TableMigrationOp::ChangePrimaryKey { column: None }, TableMigrationOp::ChangeColumnType { column }, @@ -187,105 +205,43 @@ impl Migration { candidates } +} - pub fn apply_to(&self, schema: &SchemaPlan) -> anyhow::Result { - let mut next = schema.clone(); +impl SchemaRewrite { + fn alter_table(table: &TablePlan, ops: impl Into>) -> Self { + Self::AlterTable { + table: table.name.clone(), + ops: ops.into(), + } + } + pub(crate) fn apply_to(&self, schema: &mut SchemaPlan) -> anyhow::Result<()> { match self { Self::AddTable { is_event } => { - next.tables.push(new_table(&next.tables, *is_event)); + schema.tables.push(new_table(&schema.tables, *is_event)); } Self::RemoveTable { table } => { - anyhow::ensure!( - *table < next.tables.len(), - "remove-table migration references missing table {table}" - ); - next.tables.remove(*table); + let table = table_position(schema, table)?; + schema.tables.remove(table); } Self::AlterTable { table, ops } => { - let table_plan = next - .tables - .get_mut(*table) - .ok_or_else(|| anyhow::anyhow!("migration references missing table {table}"))?; - + let table = table_position(schema, table)?; + let table_plan = &mut schema.tables[table]; for op in ops { apply_table_op(table_plan, op.clone())?; } } } - - Ok(next) - } - - pub fn added_column_defaults(&self) -> Option<(usize, Vec)> { - let Self::AlterTable { table, ops } = self else { - return None; - }; - let defaults = ops - .iter() - .filter_map(|op| match op { - TableMigrationOp::AddColumn { ty } => Some(ty.default_value()), - _ => None, - }) - .collect::>(); - - (!defaults.is_empty()).then_some((*table, defaults)) - } - - pub fn expected_steps(&self) -> Vec { - let mut expected = Vec::new(); - - match self { - Self::AddTable { .. } => expected.push(ExpectedStep::AddTable), - Self::RemoveTable { .. } => { - expected.push(ExpectedStep::RemoveTable); - expected.push(ExpectedStep::DisconnectAllUsers); - } - Self::AlterTable { ops, .. } => { - for op in ops { - expected_steps_for_table_op(op, &mut expected); - } - } - } - - if expected.contains(&ExpectedStep::AddColumns) { - expected.retain(|step| *step != ExpectedStep::ChangeColumns); - } - - expected.sort(); - expected.dedup(); - expected - } - - fn alter_table(table: usize, ops: impl Into>) -> Self { - Self::AlterTable { table, ops: ops.into() } + Ok(()) } } -fn expected_steps_for_table_op(op: &TableMigrationOp, expected: &mut Vec) { - match op { - TableMigrationOp::ChangeAccess => expected.push(ExpectedStep::ChangeAccess), - TableMigrationOp::AddColumn { .. } => { - expected.push(ExpectedStep::AddColumns); - expected.push(ExpectedStep::DisconnectAllUsers); - } - TableMigrationOp::AddIndex { .. } => expected.push(ExpectedStep::AddIndex), - TableMigrationOp::RemoveIndex { .. } => expected.push(ExpectedStep::RemoveIndex), - TableMigrationOp::AddSequence { .. } => expected.push(ExpectedStep::AddSequence), - TableMigrationOp::RemoveSequence { .. } => expected.push(ExpectedStep::RemoveSequence), - TableMigrationOp::AddUniqueConstraint { .. } => expected.push(ExpectedStep::AddConstraint), - TableMigrationOp::RemoveUniqueConstraint { .. } => expected.push(ExpectedStep::RemoveConstraint), - TableMigrationOp::ChangePrimaryKey { .. } => expected.push(ExpectedStep::ChangePrimaryKey), - TableMigrationOp::ChangeIndex { .. } => { - expected.push(ExpectedStep::RemoveIndex); - expected.push(ExpectedStep::AddIndex); - } - TableMigrationOp::ChangeColumnType { .. } => expected.push(ExpectedStep::ChangeColumns), - TableMigrationOp::ReschemaEventTable => { - expected.push(ExpectedStep::ReschemaEventTable); - expected.push(ExpectedStep::DisconnectAllUsers); - } - } +fn table_position(schema: &SchemaPlan, table: &str) -> anyhow::Result { + schema + .tables + .iter() + .position(|table_plan| table_plan.name == table) + .ok_or_else(|| anyhow::anyhow!("migration references missing table {table}")) } fn apply_table_op(table: &mut TablePlan, op: TableMigrationOp) -> anyhow::Result<()> { @@ -312,12 +268,13 @@ fn apply_table_op(table: &mut TablePlan, op: TableMigrationOp) -> anyhow::Result ); table.indexes.push(IndexPlan { columns, algorithm }); } - TableMigrationOp::RemoveIndex { index } => { - let Some(index_plan) = table.indexes.get(index) else { - anyhow::bail!("remove-index migration references missing index {index}"); + TableMigrationOp::RemoveIndex { columns } => { + ensure_columns_exist(table, &columns)?; + let Some(index) = table.indexes.iter().position(|index| index.columns == columns) else { + anyhow::bail!("remove-index migration references missing index on columns {columns:?}"); }; anyhow::ensure!( - !is_required_index(table, &index_plan.columns), + !is_required_index(table, &columns), "remove-index migration selected a required index" ); table.indexes.remove(index); @@ -337,11 +294,14 @@ fn apply_table_op(table: &mut TablePlan, op: TableMigrationOp) -> anyhow::Result ); table.sequences.push(sequence); } - TableMigrationOp::RemoveSequence { sequence } => { + TableMigrationOp::RemoveSequence { column } => { anyhow::ensure!( - sequence < table.sequences.len(), - "remove-sequence migration references missing sequence" + column < table.columns.len(), + "remove-sequence migration references missing column" ); + let Some(sequence) = table.sequences.iter().position(|sequence| sequence.column == column) else { + anyhow::bail!("remove-sequence migration references missing sequence on column {column}"); + }; table.sequences.remove(sequence); } TableMigrationOp::AddUniqueConstraint { columns } => { @@ -359,14 +319,17 @@ fn apply_table_op(table: &mut TablePlan, op: TableMigrationOp) -> anyhow::Result ); table.unique_constraints.push(UniqueConstraintPlan { columns }); } - TableMigrationOp::RemoveUniqueConstraint { constraint } => { - let Some(constraint_plan) = table.unique_constraints.get(constraint) else { - anyhow::bail!("remove-constraint migration references missing constraint {constraint}"); + TableMigrationOp::RemoveUniqueConstraint { columns } => { + ensure_columns_exist(table, &columns)?; + let Some(constraint) = table + .unique_constraints + .iter() + .position(|constraint| constraint.columns == columns) + else { + anyhow::bail!("remove-constraint migration references missing constraint on columns {columns:?}"); }; anyhow::ensure!( - !table - .primary_key - .is_some_and(|primary_key| constraint_plan.columns == [primary_key]), + !table.primary_key.is_some_and(|primary_key| columns == [primary_key]), "remove-constraint migration selected primary-key constraint" ); table.unique_constraints.remove(constraint); @@ -380,9 +343,14 @@ fn apply_table_op(table: &mut TablePlan, op: TableMigrationOp) -> anyhow::Result } table.primary_key = column; } - TableMigrationOp::ChangeIndex { index } => { - let Some(index_plan) = table.indexes.get_mut(index) else { - anyhow::bail!("index migration references missing index {index}"); + TableMigrationOp::ChangeIndex { columns } => { + ensure_columns_exist(table, &columns)?; + anyhow::ensure!( + !is_required_index(table, &columns), + "index migration selected a required index" + ); + let Some(index_plan) = table.indexes.iter_mut().find(|index| index.columns == columns) else { + anyhow::bail!("index migration references missing index on columns {columns:?}"); }; index_plan.algorithm = match index_plan.algorithm { IndexAlgorithm::BTree => IndexAlgorithm::Hash, @@ -459,74 +427,104 @@ fn new_table(tables: &[TablePlan], is_event: bool) -> TablePlan { } } -fn first_widenable_sum_column(table: &TablePlan) -> Option { - table.columns.iter().position(|column| match column.ty { - Type::Sum { variants } => variants < MAX_SUM_VARIANTS, - _ => false, - }) +fn widenable_sum_columns(table: &TablePlan) -> Vec { + table + .columns + .iter() + .enumerate() + .filter_map(|(column, column_plan)| match column_plan.ty { + Type::Sum { variants } if variants < MAX_SUM_VARIANTS => Some(column), + _ => None, + }) + .collect() } -fn first_addable_index(table: &TablePlan) -> Option<(Vec, IndexAlgorithm)> { - (0..table.columns.len()).find_map(|column| { - let columns = vec![column]; - (!has_index(table, &columns)).then_some((columns, IndexAlgorithm::BTree)) - }) +fn addable_indexes(table: &TablePlan) -> Vec<(Vec, IndexAlgorithm)> { + (0..table.columns.len()) + .filter_map(|column| { + let columns = vec![column]; + (!has_index(table, &columns)).then_some((columns, IndexAlgorithm::BTree)) + }) + .collect() } -fn first_changeable_index(table: &TablePlan) -> Option { +fn changeable_index_columns(table: &TablePlan) -> Vec> { table .indexes .iter() - .position(|index| !is_required_index(table, &index.columns)) + .filter_map(|index| (!is_required_index(table, &index.columns)).then_some(index.columns.clone())) + .collect() } -fn first_addable_sequence(table: &TablePlan) -> Option { - table.columns.iter().enumerate().find_map(|(column, column_plan)| { - (column_plan.ty.is_integral() && table.sequences.iter().all(|sequence| sequence.column != column)) - .then(|| SequencePlan::new(column, column_plan.ty).expect("column type checked above")) - }) +fn addable_sequences(table: &TablePlan) -> Vec { + table + .columns + .iter() + .enumerate() + .filter_map(|(column, column_plan)| { + (column_plan.ty.is_integral() && table.sequences.iter().all(|sequence| sequence.column != column)) + .then(|| SequencePlan::new(column, column_plan.ty).expect("column type checked above")) + }) + .collect() } -fn first_addable_sequence_boundary_probe( +fn addable_sequence_boundary_probes( table: &TablePlan, - column_domain: impl Fn(usize) -> ColumnDomain, -) -> Option { - table.columns.iter().enumerate().find_map(|(column, column_plan)| { - let domain = column_domain(column); - if column_plan.ty != Type::U64 - || domain.sequenced - || !domain.single_column_indexed - || !domain.single_column_unique - { - return None; - } + column_domain: impl Fn(usize) -> Option, +) -> Vec { + table + .columns + .iter() + .enumerate() + .filter_map(|(column, column_plan)| { + let domain = column_domain(column)?; + if column_plan.ty != Type::U64 + || domain.sequenced + || !domain.single_column_indexed + || !domain.single_column_unique + { + return None; + } - let max_value = domain.positive_i128_value_above(2)?; - SequencePlan::with_existing_value_as_max(column, column_plan.ty, max_value) - }) + let max_value = domain.positive_i128_value_above(2)?; + SequencePlan::with_existing_value_as_max(column, column_plan.ty, max_value) + }) + .collect() } -fn first_removable_sequence(table: &TablePlan) -> Option { - (!table.sequences.is_empty()).then_some(0) +fn column_domain(model: &Model, table: &TablePlan, column: usize) -> Option { + let column_name = &table.columns.get(column)?.name; + model.column_domain_by_name(&table.name, column_name) } -fn first_addable_unique_constraint_columns(table: &TablePlan) -> Option> { - (0..table.columns.len()).find_map(|column| { - let columns = vec![column]; - (!table - .unique_constraints - .iter() - .any(|constraint| constraint.columns == columns)) - .then_some(columns) - }) +fn removable_sequence_columns(table: &TablePlan) -> Vec { + table.sequences.iter().map(|sequence| sequence.column).collect() } -fn first_removable_unique_constraint(table: &TablePlan) -> Option { - table.unique_constraints.iter().position(|constraint| { - !table - .primary_key - .is_some_and(|primary_key| constraint.columns == [primary_key]) - }) +fn addable_unique_constraint_columns(table: &TablePlan) -> Vec> { + (0..table.columns.len()) + .filter_map(|column| { + let columns = vec![column]; + (!table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == columns)) + .then_some(columns) + }) + .collect() +} + +fn removable_unique_constraint_columns(table: &TablePlan) -> Vec> { + table + .unique_constraints + .iter() + .filter_map(|constraint| { + (!table + .primary_key + .is_some_and(|primary_key| constraint.columns == [primary_key])) + .then_some(constraint.columns.clone()) + }) + .collect() } fn has_index(table: &TablePlan, columns: &[usize]) -> bool { @@ -561,7 +559,7 @@ fn widen_sum_column(table: &mut TablePlan, column: Option) -> anyhow::Res }; anyhow::ensure!( *variants < MAX_SUM_VARIANTS, - "sum column already has the configured maximum number of variants" + "sum-widening migration selected maxed-out sum column" ); *variants += 1; Ok(()) diff --git a/crates/dst/src/engine/model.rs b/crates/dst/src/engine/model.rs index 5dfcd068225..e6a194765b9 100644 --- a/crates/dst/src/engine/model.rs +++ b/crates/dst/src/engine/model.rs @@ -1,10 +1,9 @@ -use spacetimedb_lib::AlgebraicValue; +use spacetimedb_lib::{AlgebraicValue, ProductValue}; -use super::migrations::Migration; use super::row::{normalize_rows, Row}; use super::state::{schema_state_for_plan, CommitDelta, CountState, TableDelta, TableRowCount, TableRows}; use super::workload::{InsertOutcome, Interaction, Observation}; -use crate::schema::{SchemaPlan, Type}; +use crate::schema::{ColumnPlan, SchemaPlan, TablePlan, Type}; #[derive(Debug)] pub struct Model { @@ -189,30 +188,9 @@ impl Model { } Interaction::Migrate(migration) => { debug_assert!(self.pending_tx.is_none()); - let added_defaults = migration.added_column_defaults(); - self.schema = migration - .apply_to(&self.schema) - .expect("generated migrations must be valid for the model schema"); - match migration { - Migration::AddTable { .. } => { - self.committed_tables.push(TableState { - rows: vec![], - ever_inserted: false, - }); - } - Migration::RemoveTable { table } => { - self.committed_tables.remove(*table); - } - Migration::AlterTable { .. } => { - if let Some((table, defaults)) = added_defaults { - for row in &mut self.committed_tables[table].rows { - let mut elements = row.elements.to_vec(); - elements.extend(defaults.iter().cloned()); - row.elements = elements.into_boxed_slice(); - } - } - } - } + let old_schema = std::mem::replace(&mut self.schema, migration.schema().clone()); + let old_tables = std::mem::take(&mut self.committed_tables); + self.committed_tables = remap_table_states(&old_schema, &self.schema, old_tables); Observation::Migrated } Interaction::Replay => { @@ -285,6 +263,31 @@ impl Model { self.committed_tables[table].ever_inserted } + pub(crate) fn row_count_by_table_name(&self, table: &str) -> usize { + self.table_index(table).map_or(0, |table| self.row_count(table)) + } + + pub(crate) fn ever_inserted_by_table_name(&self, table: &str) -> bool { + self.table_index(table) + .is_some_and(|table| self.committed_tables[table].ever_inserted) + } + + pub(crate) fn column_domain_by_name(&self, table: &str, column: &str) -> Option { + let table = self.table_index(table)?; + let column = self.schema.tables[table] + .columns + .iter() + .position(|column_plan| column_plan.name == column)?; + Some(self.column_domain(table, column)) + } + + fn table_index(&self, table: &str) -> Option { + self.schema + .tables + .iter() + .position(|table_plan| table_plan.name == table) + } + pub(crate) fn column_domain(&self, table: usize, column: usize) -> ColumnDomain { let table_plan = &self.schema.tables[table]; ColumnDomain { @@ -337,6 +340,67 @@ impl Model { } } +fn remap_table_states( + old_schema: &SchemaPlan, + new_schema: &SchemaPlan, + old_tables: Vec, +) -> Vec { + let mut old_tables = old_tables.into_iter().map(Some).collect::>(); + new_schema + .tables + .iter() + .map(|new_table| { + let Some(old_table_idx) = old_schema + .tables + .iter() + .position(|old_table| old_table.name == new_table.name) + else { + return TableState { + rows: vec![], + ever_inserted: false, + }; + }; + + let old_table = &old_schema.tables[old_table_idx]; + let old_state = old_tables[old_table_idx] + .take() + .expect("old table state is consumed once"); + remap_table_state(old_table, new_table, old_state) + }) + .collect() +} + +fn remap_table_state(old_table: &TablePlan, new_table: &TablePlan, state: TableState) -> TableState { + TableState { + rows: state + .rows + .into_iter() + .map(|row| remap_row(old_table, new_table, row)) + .collect(), + ever_inserted: state.ever_inserted, + } +} + +fn remap_row(old_table: &TablePlan, new_table: &TablePlan, row: Row) -> Row { + let elements = new_table + .columns + .iter() + .map(|new_column| remap_value(old_table, new_column, &row)) + .collect::>(); + ProductValue { + elements: elements.into_boxed_slice(), + } +} + +fn remap_value(old_table: &TablePlan, new_column: &ColumnPlan, row: &Row) -> AlgebraicValue { + old_table + .columns + .iter() + .position(|old_column| old_column.name == new_column.name) + .map(|old_column| row.elements[old_column].clone()) + .unwrap_or_else(|| new_column.ty.default_value()) +} + #[cfg(test)] mod tests { use spacetimedb_lib::AlgebraicValue; diff --git a/crates/dst/src/rng.rs b/crates/dst/src/rng.rs index 5f471105321..bc1a5a465fa 100644 --- a/crates/dst/src/rng.rs +++ b/crates/dst/src/rng.rs @@ -28,6 +28,14 @@ pub(crate) fn frequency(rng: &Rng, choices: &[Choice]) -> T { unreachable!("selected value is always inside total weight") } +pub(crate) trait WeightedChoice: Copy + 'static { + const CHOICES: &'static [Choice]; + + fn pick(rng: &Rng) -> Self { + frequency(rng, Self::CHOICES) + } +} + pub(crate) fn pick_weighted(rng: &Rng, weights: &[u64]) -> usize { let total: u64 = weights.iter().sum(); diff --git a/crates/dst/src/schema.rs b/crates/dst/src/schema.rs index cdfc7e166aa..b8562559ac1 100644 --- a/crates/dst/src/schema.rs +++ b/crates/dst/src/schema.rs @@ -181,7 +181,7 @@ impl SchemaNames { // Schema plan — the canonical source of truth. // This Schema should be able to translate to valid `RawModuleDefV10`. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct SchemaPlan { pub tables: Vec, } @@ -382,7 +382,7 @@ impl SchemaPlan { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct TablePlan { pub name: String, pub columns: Vec, @@ -394,13 +394,13 @@ pub struct TablePlan { pub is_event: bool, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct ColumnPlan { pub name: String, pub ty: Type, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct IndexPlan { /// Indices into `TablePlan.columns`. pub columns: Vec, @@ -413,7 +413,7 @@ pub enum IndexAlgorithm { Hash, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct UniqueConstraintPlan { /// Indices into `TablePlan.columns`. Non-empty. pub columns: Vec, From 8f81a734511ec64a23ff4462303b27617d2172ae Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Tue, 14 Jul 2026 15:04:44 +0530 Subject: [PATCH 06/10] consistent choices --- crates/dst/src/engine/generation.rs | 10 +- crates/dst/src/engine/migrations.rs | 381 ++++++++++++++++++++-------- crates/dst/src/engine/workload.rs | 41 ++- crates/dst/src/rng.rs | 28 +- 4 files changed, 307 insertions(+), 153 deletions(-) diff --git a/crates/dst/src/engine/generation.rs b/crates/dst/src/engine/generation.rs index d46b3386d4e..70204fa91b1 100644 --- a/crates/dst/src/engine/generation.rs +++ b/crates/dst/src/engine/generation.rs @@ -5,7 +5,7 @@ use spacetimedb_sats::ArrayValue; use super::migrations::Migration; use super::model::{ColumnDomain, Model}; use super::row::Row; -use crate::rng::{choice, choose_index, Choice, WeightedChoice}; +use crate::rng::{choice, Choice, WeightedChoice}; use crate::schema::Type; pub(crate) struct GenCtx<'a> { @@ -99,7 +99,7 @@ impl StringCase { ]; fn pick_weird(rng: &Rng) -> Self { - crate::rng::frequency(rng, Self::WEIRD_CHOICES) + crate::rng::pick_choice(rng, Self::WEIRD_CHOICES) } } @@ -123,7 +123,7 @@ impl BytesCase { ]; fn pick_weird(rng: &Rng) -> Self { - crate::rng::frequency(rng, Self::WEIRD_CHOICES) + crate::rng::pick_choice(rng, Self::WEIRD_CHOICES) } } @@ -369,11 +369,9 @@ impl<'a> MigrationGen<'a> { let steps = 1 + self.rng.index(10); for _ in 0..steps { - let candidates = Migration::candidates(&schema, self.model); - let Some(idx) = choose_index(self.rng, candidates.len()) else { + let Some(rewrite) = Migration::choose_rewrite(self.rng, &schema, self.model) else { break; }; - let rewrite = candidates[idx].clone(); rewrite .apply_to(&mut schema) .expect("generated rewrite must be valid for the draft schema"); diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs index 73263545858..1fb656820df 100644 --- a/crates/dst/src/engine/migrations.rs +++ b/crates/dst/src/engine/migrations.rs @@ -1,7 +1,9 @@ use super::model::{ColumnDomain, Model}; +use crate::rng::{choice, choose_index, Choice, WeightedChoice}; use crate::schema::{ ColumnPlan, IndexAlgorithm, IndexPlan, SchemaNames, SchemaPlan, SequencePlan, TablePlan, Type, UniqueConstraintPlan, }; +use spacetimedb_runtime::sim::Rng; const MAX_SUM_VARIANTS: u8 = 32; const MAX_EVENT_COLUMNS: usize = 32; @@ -57,6 +59,41 @@ pub enum TableMigrationOp { ReschemaEventTable, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum MigrationChoice { + AddTable, + RemoveTable, + AddColumn, + AddIndex, + RemoveIndex, + ChangeIndex, + AddSequence, + RemoveSequence, + AddUniqueConstraint, + RemoveUniqueConstraint, + DropPrimaryKeyAndWidenSum, + WidenSumColumn, + ReschemaEventTable, +} + +impl WeightedChoice for MigrationChoice { + const CHOICES: &'static [Choice] = &[ + choice(2, Self::AddTable), + choice(1, Self::RemoveTable), + choice(14, Self::AddColumn), + choice(10, Self::AddIndex), + choice(8, Self::RemoveIndex), + choice(10, Self::ChangeIndex), + choice(18, Self::AddSequence), + choice(8, Self::RemoveSequence), + choice(12, Self::AddUniqueConstraint), + choice(8, Self::RemoveUniqueConstraint), + choice(6, Self::DropPrimaryKeyAndWidenSum), + choice(12, Self::WidenSumColumn), + choice(8, Self::ReschemaEventTable), + ]; +} + impl Migration { pub(crate) fn from_schema(schema: SchemaPlan) -> Self { Self { schema } @@ -75,136 +112,268 @@ impl Migration { &self.schema } + pub(crate) fn choose_rewrite(rng: &Rng, schema: &SchemaPlan, model: &Model) -> Option { + for _ in 0..16 { + let choice = MigrationChoice::pick(rng); + if let Some(rewrite) = pick_rewrite(rng, Self::candidates_for(schema, model, choice)) { + return Some(rewrite); + } + } + + pick_rewrite(rng, Self::candidates(schema, model)) + } + pub(crate) fn candidates(schema: &SchemaPlan, model: &Model) -> Vec { - let mut candidates = Vec::new(); - let non_event_tables = schema.tables.iter().filter(|table| !table.is_event).count(); + ::CHOICES + .iter() + .flat_map(|choice| Self::candidates_for(schema, model, choice.value())) + .collect() + } - if schema.tables.len() < MAX_TABLES { - candidates.push(SchemaRewrite::AddTable { is_event: false }); - candidates.push(SchemaRewrite::AddTable { is_event: true }); + fn candidates_for(schema: &SchemaPlan, model: &Model, choice: MigrationChoice) -> Vec { + match choice { + MigrationChoice::AddTable => add_table_rewrites(schema), + MigrationChoice::RemoveTable => remove_table_rewrites(schema, model), + MigrationChoice::AddColumn => add_column_rewrites(schema), + MigrationChoice::AddIndex => add_index_rewrites(schema), + MigrationChoice::RemoveIndex => remove_index_rewrites(schema), + MigrationChoice::ChangeIndex => change_index_rewrites(schema), + MigrationChoice::AddSequence => add_sequence_rewrites(schema, model), + MigrationChoice::RemoveSequence => remove_sequence_rewrites(schema), + MigrationChoice::AddUniqueConstraint => add_unique_constraint_rewrites(schema, model), + MigrationChoice::RemoveUniqueConstraint => remove_unique_constraint_rewrites(schema), + MigrationChoice::DropPrimaryKeyAndWidenSum => drop_primary_key_and_widen_sum_rewrites(schema), + MigrationChoice::WidenSumColumn => widen_sum_column_rewrites(schema), + MigrationChoice::ReschemaEventTable => reschema_event_table_rewrites(schema, model), } + } +} - for table_plan in &schema.tables { - let row_count = model.row_count_by_table_name(&table_plan.name); - let pristine = row_count == 0 && !model.ever_inserted_by_table_name(&table_plan.name); +fn pick_rewrite(rng: &Rng, mut candidates: Vec) -> Option { + let idx = choose_index(rng, candidates.len())?; + Some(candidates.swap_remove(idx)) +} - if (table_plan.is_event && row_count == 0) || (!table_plan.is_event && pristine && non_event_tables > 1) { - candidates.push(SchemaRewrite::RemoveTable { - table: table_plan.name.clone(), - }); - } +fn add_table_rewrites(schema: &SchemaPlan) -> Vec { + if schema.tables.len() >= MAX_TABLES { + return Vec::new(); + } - if table_plan.is_event { - if row_count == 0 && table_plan.columns.len() < MAX_EVENT_COLUMNS { - candidates.push(SchemaRewrite::alter_table( - table_plan, - [TableMigrationOp::ChangeAccess, TableMigrationOp::ReschemaEventTable], - )); - } - continue; - } + vec![ + SchemaRewrite::AddTable { is_event: false }, + SchemaRewrite::AddTable { is_event: true }, + ] +} - if table_plan.columns.len() < MAX_TABLE_COLUMNS { - candidates.extend(Type::ALL.iter().copied().map(|ty| { - SchemaRewrite::alter_table( - table_plan, - [TableMigrationOp::ChangeAccess, TableMigrationOp::AddColumn { ty }], - ) - })); - } +fn remove_table_rewrites(schema: &SchemaPlan, model: &Model) -> Vec { + let non_event_tables = schema.tables.iter().filter(|table| !table.is_event).count(); - if pristine { - for sequence in addable_sequences(table_plan) { - candidates.push(SchemaRewrite::alter_table( - table_plan, - [TableMigrationOp::AddSequence { sequence }], - )); + schema + .tables + .iter() + .filter_map(|table| { + let row_count = model.row_count_by_table_name(&table.name); + let pristine = row_count == 0 && !model.ever_inserted_by_table_name(&table.name); + ((table.is_event && row_count == 0) || (!table.is_event && pristine && non_event_tables > 1)).then(|| { + SchemaRewrite::RemoveTable { + table: table.name.clone(), } + }) + }) + .collect() +} - for columns in addable_unique_constraint_columns(table_plan) { - let mut ops = Vec::new(); - if !has_index(table_plan, &columns) { - ops.push(TableMigrationOp::AddIndex { - columns: columns.clone(), - algorithm: IndexAlgorithm::BTree, - }); - } - ops.push(TableMigrationOp::AddUniqueConstraint { columns }); - candidates.push(SchemaRewrite::alter_table(table_plan, ops)); - } - } +fn add_column_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event && table.columns.len() < MAX_TABLE_COLUMNS) + .flat_map(|table| { + Type::ALL.iter().copied().map(|ty| { + SchemaRewrite::alter_table( + table, + [TableMigrationOp::ChangeAccess, TableMigrationOp::AddColumn { ty }], + ) + }) + }) + .collect() +} - if row_count > 0 { - for sequence in - addable_sequence_boundary_probes(table_plan, |column| column_domain(model, table_plan, column)) - { - candidates.push(SchemaRewrite::alter_table( - table_plan, - [TableMigrationOp::AddSequence { sequence }], - )); - } - } +fn add_index_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + addable_indexes(table).into_iter().map(|(columns, algorithm)| { + SchemaRewrite::alter_table(table, [TableMigrationOp::AddIndex { columns, algorithm }]) + }) + }) + .collect() +} - for column in removable_sequence_columns(table_plan) { - candidates.push(SchemaRewrite::alter_table( - table_plan, - [TableMigrationOp::RemoveSequence { column }], - )); - } +fn remove_index_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + changeable_index_columns(table) + .into_iter() + .map(|columns| SchemaRewrite::alter_table(table, [TableMigrationOp::RemoveIndex { columns }])) + }) + .collect() +} - for columns in removable_unique_constraint_columns(table_plan) { - candidates.push(SchemaRewrite::alter_table( - table_plan, - [TableMigrationOp::RemoveUniqueConstraint { columns }], - )); - } +fn change_index_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + changeable_index_columns(table).into_iter().map(|columns| { + SchemaRewrite::alter_table( + table, + [ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeIndex { columns }, + ], + ) + }) + }) + .collect() +} + +fn add_sequence_rewrites(schema: &SchemaPlan, model: &Model) -> Vec { + let mut rewrites = Vec::new(); + + for table in schema.tables.iter().filter(|table| !table.is_event) { + let row_count = model.row_count_by_table_name(&table.name); + let pristine = row_count == 0 && !model.ever_inserted_by_table_name(&table.name); + + if pristine { + rewrites.extend( + addable_sequences(table) + .into_iter() + .map(|sequence| SchemaRewrite::alter_table(table, [TableMigrationOp::AddSequence { sequence }])), + ); + } + + if row_count > 0 { + rewrites.extend( + addable_sequence_boundary_probes(table, |column| column_domain(model, table, column)) + .into_iter() + .map(|sequence| SchemaRewrite::alter_table(table, [TableMigrationOp::AddSequence { sequence }])), + ); + } + } + + rewrites +} + +fn remove_sequence_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + removable_sequence_columns(table) + .into_iter() + .map(|column| SchemaRewrite::alter_table(table, [TableMigrationOp::RemoveSequence { column }])) + }) + .collect() +} + +fn add_unique_constraint_rewrites(schema: &SchemaPlan, model: &Model) -> Vec { + let mut rewrites = Vec::new(); + + for table in schema.tables.iter().filter(|table| !table.is_event) { + let row_count = model.row_count_by_table_name(&table.name); + let pristine = row_count == 0 && !model.ever_inserted_by_table_name(&table.name); + if !pristine { + continue; + } - for (columns, algorithm) in addable_indexes(table_plan) { - candidates.push(SchemaRewrite::alter_table( - table_plan, - [TableMigrationOp::AddIndex { columns, algorithm }], - )); + for columns in addable_unique_constraint_columns(table) { + let mut ops = Vec::new(); + if !has_index(table, &columns) { + ops.push(TableMigrationOp::AddIndex { + columns: columns.clone(), + algorithm: IndexAlgorithm::BTree, + }); } + ops.push(TableMigrationOp::AddUniqueConstraint { columns }); + rewrites.push(SchemaRewrite::alter_table(table, ops)); + } + } - for columns in changeable_index_columns(table_plan) { - candidates.push(SchemaRewrite::alter_table( - table_plan, + rewrites +} + +fn remove_unique_constraint_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + removable_unique_constraint_columns(table).into_iter().map(|columns| { + SchemaRewrite::alter_table(table, [TableMigrationOp::RemoveUniqueConstraint { columns }]) + }) + }) + .collect() +} + +fn drop_primary_key_and_widen_sum_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event && table.primary_key.is_some() && table.sequences.is_empty()) + .flat_map(|table| { + widenable_sum_columns(table).into_iter().map(|column| { + SchemaRewrite::alter_table( + table, [ - TableMigrationOp::ChangeAccess, - TableMigrationOp::ChangeIndex { - columns: columns.clone(), - }, + TableMigrationOp::ChangePrimaryKey { column: None }, + TableMigrationOp::ChangeColumnType { column }, ], - )); - candidates.push(SchemaRewrite::alter_table( - table_plan, - [TableMigrationOp::RemoveIndex { columns }], - )); - } + ) + }) + }) + .collect() +} - for column in widenable_sum_columns(table_plan) { - candidates.push(SchemaRewrite::alter_table( - table_plan, +fn widen_sum_column_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + widenable_sum_columns(table).into_iter().map(|column| { + SchemaRewrite::alter_table( + table, [ TableMigrationOp::ChangeAccess, TableMigrationOp::ChangeColumnType { column }, ], - )); - - if table_plan.primary_key.is_some() && table_plan.sequences.is_empty() { - candidates.push(SchemaRewrite::alter_table( - table_plan, - [ - TableMigrationOp::ChangePrimaryKey { column: None }, - TableMigrationOp::ChangeColumnType { column }, - ], - )); - } - } - } + ) + }) + }) + .collect() +} - candidates - } +fn reschema_event_table_rewrites(schema: &SchemaPlan, model: &Model) -> Vec { + schema + .tables + .iter() + .filter(|table| table.is_event && model.row_count_by_table_name(&table.name) == 0) + .filter(|table| table.columns.len() < MAX_EVENT_COLUMNS) + .map(|table| { + SchemaRewrite::alter_table( + table, + [TableMigrationOp::ChangeAccess, TableMigrationOp::ReschemaEventTable], + ) + }) + .collect() } impl SchemaRewrite { diff --git a/crates/dst/src/engine/workload.rs b/crates/dst/src/engine/workload.rs index f4fa39f6579..f77933a3401 100644 --- a/crates/dst/src/engine/workload.rs +++ b/crates/dst/src/engine/workload.rs @@ -5,7 +5,7 @@ use super::migrations::Migration; use super::model::Model; use super::row::Row; use super::state::{CommitDelta, CountState}; -use crate::rng::pick_weighted; +use crate::rng::{choice, pick_choice, Choice}; use crate::schema::SchemaPlan; use spacetimedb_runtime::sim::Rng; @@ -91,6 +91,18 @@ enum InteractionChoice { Replay, } +impl InteractionWeights { + fn choices(self) -> [Choice; 5] { + [ + choice(self.insert, InteractionChoice::Insert), + choice(self.delete, InteractionChoice::Delete), + choice(self.commit_tx, InteractionChoice::CommitTx), + choice(self.migrate, InteractionChoice::Migrate), + choice(self.replay, InteractionChoice::Replay), + ] + } +} + pub struct WorkloadGen { rng: Rng, model: Model, @@ -100,11 +112,15 @@ pub struct WorkloadGen { impl WorkloadGen { pub fn new(rng: Rng, model: Model) -> Self { + Self::with_weights(rng, model, InteractionWeights::default()) + } + + pub fn with_weights(rng: Rng, model: Model, weights: InteractionWeights) -> Self { Self { rng, model, stats: InteractionCounts::default(), - weights: InteractionWeights::default(), + weights, } } @@ -186,25 +202,8 @@ impl WorkloadGen { } fn pick_interaction_choice(&mut self) -> InteractionChoice { - let weights = self.weights; - - match pick_weighted( - &self.rng, - &[ - weights.insert, - weights.delete, - weights.commit_tx, - weights.migrate, - weights.replay, - ], - ) { - 0 => InteractionChoice::Insert, - 1 => InteractionChoice::Delete, - 2 => InteractionChoice::CommitTx, - 3 => InteractionChoice::Migrate, - 4 => InteractionChoice::Replay, - _ => unreachable!(), - } + let choices = self.weights.choices(); + pick_choice(&self.rng, &choices) } fn insert_table_idx(&self) -> usize { diff --git a/crates/dst/src/rng.rs b/crates/dst/src/rng.rs index bc1a5a465fa..084ed909544 100644 --- a/crates/dst/src/rng.rs +++ b/crates/dst/src/rng.rs @@ -6,11 +6,17 @@ pub(crate) struct Choice { value: T, } +impl Choice { + pub(crate) const fn value(self) -> T { + self.value + } +} + pub(crate) const fn choice(weight: u64, value: T) -> Choice { Choice { weight, value } } -pub(crate) fn frequency(rng: &Rng, choices: &[Choice]) -> T { +pub(crate) fn pick_choice(rng: &Rng, choices: &[Choice]) -> T { let total: u64 = choices.iter().map(|choice| choice.weight).sum(); assert!(total > 0, "at least one choice weight must be non-zero"); @@ -32,26 +38,8 @@ pub(crate) trait WeightedChoice: Copy + 'static { const CHOICES: &'static [Choice]; fn pick(rng: &Rng) -> Self { - frequency(rng, Self::CHOICES) - } -} - -pub(crate) fn pick_weighted(rng: &Rng, weights: &[u64]) -> usize { - let total: u64 = weights.iter().sum(); - - assert!(total > 0, "at least one weight must be non-zero"); - - let mut selected = rng.next_u64() % total; - - for (idx, weight) in weights.iter().copied().enumerate() { - if selected < weight { - return idx; - } - - selected -= weight; + pick_choice(rng, Self::CHOICES) } - - unreachable!("selected value is always inside total weight") } pub(crate) fn choose_index(rng: &Rng, len: usize) -> Option { From afd13094fe2c470cc01eabd902bb3d77d8a14fca Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Tue, 14 Jul 2026 15:30:42 +0530 Subject: [PATCH 07/10] remove ensure_ helpers --- crates/dst/src/engine/migrations.rs | 19 +- crates/dst/src/schema.rs | 391 ++++++++++++---------------- 2 files changed, 185 insertions(+), 225 deletions(-) diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs index 1fb656820df..e8111a7a9ab 100644 --- a/crates/dst/src/engine/migrations.rs +++ b/crates/dst/src/engine/migrations.rs @@ -184,12 +184,13 @@ fn remove_table_rewrites(schema: &SchemaPlan, model: &Model) -> Vec Vec { + let types = addable_column_types(schema); schema .tables .iter() .filter(|table| !table.is_event && table.columns.len() < MAX_TABLE_COLUMNS) .flat_map(|table| { - Type::ALL.iter().copied().map(|ty| { + types.iter().copied().map(|ty| { SchemaRewrite::alter_table( table, [TableMigrationOp::ChangeAccess, TableMigrationOp::AddColumn { ty }], @@ -199,6 +200,22 @@ fn add_column_rewrites(schema: &SchemaPlan) -> Vec { .collect() } +fn addable_column_types(schema: &SchemaPlan) -> Vec { + Type::ALL + .iter() + .copied() + .filter(|ty| !matches!(ty, Type::Sum { .. }) || !schema_has_sum_column(schema)) + .collect() +} + +fn schema_has_sum_column(schema: &SchemaPlan) -> bool { + schema + .tables + .iter() + .flat_map(|table| table.columns.iter()) + .any(|column| matches!(column.ty, Type::Sum { .. })) +} + fn add_index_rewrites(schema: &SchemaPlan) -> Vec { schema .tables diff --git a/crates/dst/src/schema.rs b/crates/dst/src/schema.rs index b8562559ac1..ef98cf93ca2 100644 --- a/crates/dst/src/schema.rs +++ b/crates/dst/src/schema.rs @@ -8,11 +8,7 @@ use spacetimedb_sats::{ }; pub fn default_schema(rng: Rng) -> SchemaPlan { - let profile = SchemaProfile::default(); - let mut plan = SchemaGenerator::new(rng, profile).gen_schema(); - plan.ensure_auto_inc_table(); - plan.ensure_engine_migration_coverage(); - plan + SchemaGenerator::new(rng, SchemaProfile::engine_dst()).gen_schema() } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -26,7 +22,14 @@ pub enum Type { } impl Type { - pub const ALL: &'static [Type] = &[Type::Bool, Type::I64, Type::U64, Type::String, Type::Bytes]; + pub const ALL: &'static [Type] = &[ + Type::Bool, + Type::I64, + Type::U64, + Type::String, + Type::Bytes, + Type::Sum { variants: 1 }, + ]; pub fn to_algebraic(self) -> AlgebraicType { match self { @@ -84,10 +87,6 @@ impl SchemaDecisions { rng.sample_probability(probability) } - pub fn gen_type(rng: &Rng) -> Type { - Type::ALL[Self::index(rng, Type::ALL.len())] - } - pub fn gen_table_name(rng: &Rng, tables: &[TablePlan]) -> String { loop { let name = format!("tbl_{}", Self::gen_ident(rng)); @@ -186,202 +185,6 @@ pub struct SchemaPlan { pub tables: Vec, } -impl SchemaPlan { - pub fn ensure_engine_migration_coverage(&mut self) { - self.ensure_widenable_sum_column(); - self.ensure_event_table(); - self.ensure_sequence_mutation_column(); - self.ensure_unique_constraint_mutation_column(); - self.ensure_standalone_index(); - } - - pub fn auto_inc_table_and_column(&self) -> Option<(usize, usize)> { - self.tables - .iter() - .enumerate() - .find_map(|(table_idx, table)| table.sequences.first().map(|sequence| (table_idx, sequence.column))) - } - - pub fn ensure_auto_inc_table(&mut self) { - if self.auto_inc_table_and_column().is_some() { - return; - } - - let table = self.tables.first_mut().expect("schema must contain at least one table"); - if table.columns.is_empty() { - table.columns.push(ColumnPlan { - name: "id".into(), - ty: Type::U64, - }); - } else { - table.columns[0].ty = Type::U64; - } - - table.primary_key = Some(0); - if !table - .unique_constraints - .iter() - .any(|constraint| constraint.columns == [0]) - { - table.unique_constraints.push(UniqueConstraintPlan { columns: vec![0] }); - } - if !table.indexes.iter().any(|index| index.columns == [0]) { - table.indexes.push(IndexPlan { - columns: vec![0], - algorithm: IndexAlgorithm::BTree, - }); - } - table.sequences = vec![SequencePlan::new(0, Type::U64).expect("u64 is integral")]; - } - - fn ensure_widenable_sum_column(&mut self) { - if self - .tables - .iter() - .any(|table| !table.is_event && table.columns.iter().any(|column| matches!(column.ty, Type::Sum { .. }))) - { - return; - } - - let table = self - .tables - .iter_mut() - .find(|table| !table.is_event) - .expect("schema must contain at least one non-event table"); - table.columns.push(ColumnPlan { - name: SchemaNames::fresh_column_name(table, "dst_sum"), - ty: Type::Sum { variants: 1 }, - }); - } - - fn ensure_event_table(&mut self) { - if self.tables.iter().any(|table| table.is_event) { - return; - } - - let name = SchemaNames::fresh_table_name(&self.tables, "dst_events"); - self.tables.push(TablePlan { - name, - columns: vec![ColumnPlan { - name: "payload".into(), - ty: Type::U64, - }], - primary_key: None, - indexes: vec![], - unique_constraints: vec![], - sequences: vec![], - is_public: true, - is_event: true, - }); - } - - fn ensure_sequence_mutation_column(&mut self) { - if self.tables.iter().any(Self::has_sequence_mutation_column) { - return; - } - - let table = self - .tables - .iter_mut() - .find(|table| !table.is_event) - .expect("schema must contain at least one non-event table"); - let column = table.columns.len(); - table.columns.push(ColumnPlan { - name: SchemaNames::fresh_column_name(table, "dst_seq"), - ty: Type::U64, - }); - table.indexes.push(IndexPlan { - columns: vec![column], - algorithm: IndexAlgorithm::BTree, - }); - table - .unique_constraints - .push(UniqueConstraintPlan { columns: vec![column] }); - } - - fn ensure_unique_constraint_mutation_column(&mut self) { - if self.tables.iter().any(|table| { - !table.is_event - && table.columns.iter().enumerate().any(|(column, _)| { - !table - .unique_constraints - .iter() - .any(|constraint| constraint.columns == [column]) - }) - }) { - return; - } - - let table = self - .tables - .iter_mut() - .find(|table| !table.is_event) - .expect("schema must contain at least one non-event table"); - table.columns.push(ColumnPlan { - name: SchemaNames::fresh_column_name(table, "dst_unique"), - ty: Type::U64, - }); - } - - fn has_sequence_mutation_column(table: &TablePlan) -> bool { - !table.is_event - && table.columns.iter().enumerate().any(|(column, plan)| { - plan.ty == Type::U64 - && table.sequences.iter().all(|seq| seq.column != column) - && table.indexes.iter().any(|index| index.columns == [column]) - && table - .unique_constraints - .iter() - .any(|constraint| constraint.columns == [column]) - }) - } - - fn ensure_standalone_index(&mut self) { - if self.tables.iter().any(|table| { - !table.is_event - && table.indexes.iter().any(|index| { - !table.primary_key.is_some_and(|pk| index.columns == [pk]) - && !table - .unique_constraints - .iter() - .any(|constraint| constraint.columns == index.columns) - }) - }) { - return; - } - - let table = self - .tables - .iter_mut() - .find(|table| !table.is_event) - .expect("schema must contain at least one non-event table"); - let column = if table.columns.len() < 2 { - table.columns.push(ColumnPlan { - name: SchemaNames::fresh_column_name(table, "dst_indexed"), - ty: Type::U64, - }); - table.columns.len() - 1 - } else { - (0..table.columns.len()) - .find(|&column| { - !table.primary_key.is_some_and(|pk| pk == column) - && !table - .unique_constraints - .iter() - .any(|constraint| constraint.columns == [column]) - }) - .unwrap_or(0) - }; - - if !table.indexes.iter().any(|index| index.columns == [column]) { - table.indexes.push(IndexPlan { - columns: vec![column], - algorithm: IndexAlgorithm::BTree, - }); - } - } -} - #[derive(Debug, Clone, PartialEq, Eq)] pub struct TablePlan { pub name: String, @@ -565,6 +368,9 @@ fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { pub struct SchemaProfile { pub table_count: (usize, usize), pub columns: (usize, usize), + pub table_kind_weights: TableKindWeights, + pub type_weights: TypeWeights, + pub sum_variants: (usize, usize), pub pk_prob: f64, pub auto_inc_prob: f64, pub indexes: (usize, usize), @@ -573,21 +379,113 @@ pub struct SchemaProfile { pub private_prob: f64, } +impl SchemaProfile { + pub fn engine_dst() -> Self { + Self { + table_count: (3, 10), + columns: (1, 20), + table_kind_weights: TableKindWeights::default(), + type_weights: TypeWeights::default(), + sum_variants: (1, 4), + pk_prob: 0.65, + auto_inc_prob: 0.20, + indexes: (0, 5), + unique_constraints: (0, 3), + btree_prob: 0.65, + private_prob: 0.10, + } + } +} + impl Default for SchemaProfile { + fn default() -> Self { + Self::engine_dst() + } +} + +#[derive(Debug, Clone, Copy)] +pub struct TableKindWeights { + pub data: u64, + pub event: u64, +} + +impl Default for TableKindWeights { + fn default() -> Self { + Self { data: 9, event: 1 } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct TypeWeights { + pub bool_: u64, + pub i64_: u64, + pub u64_: u64, + pub string: u64, + pub bytes: u64, + pub sum: u64, +} + +impl Default for TypeWeights { fn default() -> Self { Self { - table_count: (1, 100), - columns: (1, 10), - pk_prob: 0.7, - auto_inc_prob: 0.3, - indexes: (0, 3), - unique_constraints: (0, 2), - btree_prob: 0.7, - private_prob: 0.1, + bool_: 12, + i64_: 24, + u64_: 28, + string: 16, + bytes: 12, + sum: 8, } } } +#[derive(Debug, Clone, Copy)] +enum TableKind { + Data, + Event, +} + +impl TableKindWeights { + fn choices(self) -> [rng::Choice; 2] { + [ + rng::choice(self.data, TableKind::Data), + rng::choice(self.event, TableKind::Event), + ] + } +} + +#[derive(Debug, Clone, Copy)] +enum TypeKind { + Bool, + I64, + U64, + String, + Bytes, + Sum, +} + +impl TypeWeights { + fn choices(self) -> [rng::Choice; 6] { + [ + rng::choice(self.bool_, TypeKind::Bool), + rng::choice(self.i64_, TypeKind::I64), + rng::choice(self.u64_, TypeKind::U64), + rng::choice(self.string, TypeKind::String), + rng::choice(self.bytes, TypeKind::Bytes), + rng::choice(self.sum, TypeKind::Sum), + ] + } + + fn non_sum_choices(self) -> [rng::Choice; 5] { + [ + rng::choice(self.bool_, TypeKind::Bool), + rng::choice(self.i64_, TypeKind::I64), + rng::choice(self.u64_, TypeKind::U64), + rng::choice(self.string, TypeKind::String), + rng::choice(self.bytes, TypeKind::Bytes), + ] + } +} + pub struct SchemaGenerator { rng: Rng, profile: SchemaProfile, @@ -598,7 +496,7 @@ impl SchemaGenerator { Self { rng, profile } } - fn gen_columns(&self) -> Vec { + fn gen_columns(&self, sum_available: &mut bool) -> Vec { let n = SchemaDecisions::range(&self.rng, self.profile.columns); let mut names = Vec::with_capacity(n); let mut seen = Vec::with_capacity(n); @@ -607,12 +505,36 @@ impl SchemaGenerator { seen.push(name.clone()); names.push(ColumnPlan { name, - ty: SchemaDecisions::gen_type(&self.rng), + ty: self.gen_type(sum_available), }); } names } + fn gen_type(&self, sum_available: &mut bool) -> Type { + let kind = if *sum_available { + let choices = self.profile.type_weights.choices(); + rng::pick_choice(&self.rng, &choices) + } else { + let choices = self.profile.type_weights.non_sum_choices(); + rng::pick_choice(&self.rng, &choices) + }; + + match kind { + TypeKind::Bool => Type::Bool, + TypeKind::I64 => Type::I64, + TypeKind::U64 => Type::U64, + TypeKind::String => Type::String, + TypeKind::Bytes => Type::Bytes, + TypeKind::Sum => { + *sum_available = false; + Type::Sum { + variants: SchemaDecisions::range(&self.rng, self.profile.sum_variants) as u8, + } + } + } + } + fn gen_unique_constraints(&self, columns: &[ColumnPlan], pk: &Option) -> Vec { let n = SchemaDecisions::range(&self.rng, self.profile.unique_constraints); let mut seen: Vec> = Vec::new(); @@ -696,8 +618,23 @@ impl SchemaGenerator { indexes } - fn gen_table(&self, existing_tables: &[TablePlan]) -> TablePlan { - let columns = self.gen_columns(); + fn gen_table(&self, existing_tables: &[TablePlan], is_event: bool, sum_available: &mut bool) -> TablePlan { + let columns = self.gen_columns(sum_available); + let name = SchemaDecisions::gen_table_name(&self.rng, existing_tables); + let is_public = !SchemaDecisions::sample_probability(&self.rng, self.profile.private_prob); + + if is_event { + return TablePlan { + name, + columns, + primary_key: None, + indexes: vec![], + unique_constraints: vec![], + sequences: vec![], + is_public, + is_event: true, + }; + } let primary_key = if SchemaDecisions::sample_probability(&self.rng, self.profile.pk_prob) && !columns.is_empty() { @@ -722,8 +659,6 @@ impl SchemaGenerator { let indexes = self.gen_indexes(&columns, &unique_constraints, &primary_key); - let name = SchemaDecisions::gen_table_name(&self.rng, existing_tables); - TablePlan { name, columns, @@ -731,16 +666,24 @@ impl SchemaGenerator { indexes, unique_constraints, sequences, - is_public: !SchemaDecisions::sample_probability(&self.rng, self.profile.private_prob), + is_public, is_event: false, } } + fn gen_table_kind(&self) -> TableKind { + let choices = self.profile.table_kind_weights.choices(); + rng::pick_choice(&self.rng, &choices) + } + pub fn gen_schema(&self) -> SchemaPlan { let table_count = SchemaDecisions::range(&self.rng, self.profile.table_count); - let mut tables = Vec::with_capacity(table_count); - for _ in 0..table_count { - tables.push(self.gen_table(&tables)); + let mut tables: Vec = Vec::with_capacity(table_count); + let mut sum_available = true; + for table_idx in 0..table_count { + let must_be_data = table_idx + 1 == table_count && !tables.iter().any(|table| !table.is_event); + let is_event = !must_be_data && matches!(self.gen_table_kind(), TableKind::Event); + tables.push(self.gen_table(&tables, is_event, &mut sum_available)); } SchemaPlan { tables } } From a5b2981bb47c00fb8c67314dba21ad91ac77f458 Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Tue, 14 Jul 2026 15:37:43 +0530 Subject: [PATCH 08/10] common new table --- crates/dst/src/engine/migrations.rs | 85 +++++++---------------------- crates/dst/src/schema.rs | 13 +++++ 2 files changed, 32 insertions(+), 66 deletions(-) diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs index e8111a7a9ab..91b10333026 100644 --- a/crates/dst/src/engine/migrations.rs +++ b/crates/dst/src/engine/migrations.rs @@ -1,7 +1,8 @@ use super::model::{ColumnDomain, Model}; use crate::rng::{choice, choose_index, Choice, WeightedChoice}; use crate::schema::{ - ColumnPlan, IndexAlgorithm, IndexPlan, SchemaNames, SchemaPlan, SequencePlan, TablePlan, Type, UniqueConstraintPlan, + ColumnPlan, IndexAlgorithm, IndexPlan, SchemaGenerator, SchemaNames, SchemaPlan, SchemaProfile, SequencePlan, + TablePlan, Type, UniqueConstraintPlan, }; use spacetimedb_runtime::sim::Rng; @@ -17,7 +18,7 @@ pub struct Migration { #[derive(Debug, Clone, PartialEq, Eq)] pub enum SchemaRewrite { - AddTable { is_event: bool }, + AddTable { table: TablePlan }, RemoveTable { table: String }, AlterTable { table: String, ops: Vec }, } @@ -115,24 +116,24 @@ impl Migration { pub(crate) fn choose_rewrite(rng: &Rng, schema: &SchemaPlan, model: &Model) -> Option { for _ in 0..16 { let choice = MigrationChoice::pick(rng); - if let Some(rewrite) = pick_rewrite(rng, Self::candidates_for(schema, model, choice)) { + if let Some(rewrite) = pick_rewrite(rng, Self::candidates_for(rng, schema, model, choice)) { return Some(rewrite); } } - pick_rewrite(rng, Self::candidates(schema, model)) + pick_rewrite(rng, Self::candidates(rng, schema, model)) } - pub(crate) fn candidates(schema: &SchemaPlan, model: &Model) -> Vec { + pub(crate) fn candidates(rng: &Rng, schema: &SchemaPlan, model: &Model) -> Vec { ::CHOICES .iter() - .flat_map(|choice| Self::candidates_for(schema, model, choice.value())) + .flat_map(|choice| Self::candidates_for(rng, schema, model, choice.value())) .collect() } - fn candidates_for(schema: &SchemaPlan, model: &Model, choice: MigrationChoice) -> Vec { + fn candidates_for(rng: &Rng, schema: &SchemaPlan, model: &Model, choice: MigrationChoice) -> Vec { match choice { - MigrationChoice::AddTable => add_table_rewrites(schema), + MigrationChoice::AddTable => add_table_rewrites(rng, schema), MigrationChoice::RemoveTable => remove_table_rewrites(schema, model), MigrationChoice::AddColumn => add_column_rewrites(schema), MigrationChoice::AddIndex => add_index_rewrites(schema), @@ -154,15 +155,18 @@ fn pick_rewrite(rng: &Rng, mut candidates: Vec) -> Option Vec { +fn add_table_rewrites(rng: &Rng, schema: &SchemaPlan) -> Vec { if schema.tables.len() >= MAX_TABLES { return Vec::new(); } - vec![ - SchemaRewrite::AddTable { is_event: false }, - SchemaRewrite::AddTable { is_event: true }, - ] + let generator = SchemaGenerator::new(rng.clone(), SchemaProfile::engine_dst()); + [false, true] + .into_iter() + .map(|is_event| SchemaRewrite::AddTable { + table: generator.gen_table_for_schema(schema, is_event), + }) + .collect() } fn remove_table_rewrites(schema: &SchemaPlan, model: &Model) -> Vec { @@ -403,8 +407,8 @@ impl SchemaRewrite { pub(crate) fn apply_to(&self, schema: &mut SchemaPlan) -> anyhow::Result<()> { match self { - Self::AddTable { is_event } => { - schema.tables.push(new_table(&schema.tables, *is_event)); + Self::AddTable { table } => { + schema.tables.push(table.clone()); } Self::RemoveTable { table } => { let table = table_position(schema, table)?; @@ -562,57 +566,6 @@ fn apply_table_op(table: &mut TablePlan, op: TableMigrationOp) -> anyhow::Result Ok(()) } -fn new_table(tables: &[TablePlan], is_event: bool) -> TablePlan { - if is_event { - return TablePlan { - name: SchemaNames::fresh_table_name(tables, "added_events"), - columns: vec![ColumnPlan { - name: "payload".into(), - ty: Type::U64, - }], - primary_key: None, - indexes: vec![], - unique_constraints: vec![], - sequences: vec![], - is_public: true, - is_event: true, - }; - } - - TablePlan { - name: SchemaNames::fresh_table_name(tables, "added_table"), - columns: vec![ - ColumnPlan { - name: "id".into(), - ty: Type::U64, - }, - ColumnPlan { - name: "value".into(), - ty: Type::U64, - }, - ColumnPlan { - name: "kind".into(), - ty: Type::U64, - }, - ], - primary_key: Some(0), - indexes: vec![ - IndexPlan { - columns: vec![0], - algorithm: IndexAlgorithm::BTree, - }, - IndexPlan { - columns: vec![1], - algorithm: IndexAlgorithm::Hash, - }, - ], - unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], - sequences: vec![SequencePlan::new(0, Type::U64).expect("u64 is integral")], - is_public: true, - is_event: false, - } -} - fn widenable_sum_columns(table: &TablePlan) -> Vec { table .columns diff --git a/crates/dst/src/schema.rs b/crates/dst/src/schema.rs index ef98cf93ca2..5ced13ed45a 100644 --- a/crates/dst/src/schema.rs +++ b/crates/dst/src/schema.rs @@ -618,6 +618,11 @@ impl SchemaGenerator { indexes } + pub fn gen_table_for_schema(&self, schema: &SchemaPlan, is_event: bool) -> TablePlan { + let mut sum_available = !schema_has_sum_column(schema); + self.gen_table(&schema.tables, is_event, &mut sum_available) + } + fn gen_table(&self, existing_tables: &[TablePlan], is_event: bool, sum_available: &mut bool) -> TablePlan { let columns = self.gen_columns(sum_available); let name = SchemaDecisions::gen_table_name(&self.rng, existing_tables); @@ -689,6 +694,14 @@ impl SchemaGenerator { } } +fn schema_has_sum_column(schema: &SchemaPlan) -> bool { + schema + .tables + .iter() + .flat_map(|table| table.columns.iter()) + .any(|column| matches!(column.ty, Type::Sum { .. })) +} + #[cfg(test)] mod tests { use super::*; From 1072154744c3485ad45e45800ce2a6ed6621f3be Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Tue, 14 Jul 2026 15:49:17 +0530 Subject: [PATCH 09/10] code redability --- crates/dst/src/engine/generation.rs | 3 + crates/dst/src/engine/migrations.rs | 10 + crates/dst/src/engine/workload.rs | 71 +++-- crates/dst/src/rng.rs | 8 + crates/dst/src/schema.rs | 425 +++++++++++++--------------- 5 files changed, 265 insertions(+), 252 deletions(-) diff --git a/crates/dst/src/engine/generation.rs b/crates/dst/src/engine/generation.rs index 70204fa91b1..e5a1e89a4a5 100644 --- a/crates/dst/src/engine/generation.rs +++ b/crates/dst/src/engine/generation.rs @@ -1,3 +1,5 @@ +//! Workload value and migration generation for the engine DST driver. + use spacetimedb_lib::{AlgebraicValue, ProductValue}; use spacetimedb_runtime::sim::Rng; use spacetimedb_sats::ArrayValue; @@ -8,6 +10,7 @@ use super::row::Row; use crate::rng::{choice, Choice, WeightedChoice}; use crate::schema::Type; +/// Read-only generation context for one model state. pub(crate) struct GenCtx<'a> { rng: &'a Rng, model: &'a Model, diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs index 91b10333026..087699cda53 100644 --- a/crates/dst/src/engine/migrations.rs +++ b/crates/dst/src/engine/migrations.rs @@ -1,3 +1,9 @@ +//! Schema migration candidates for the engine DST driver. +//! +//! This module deliberately generates only changes expected to be handled by +//! engine auto-migration. If a generated schema requires a manual migration +//! plan, the engine driver treats that as a harness bug. + use super::model::{ColumnDomain, Model}; use crate::rng::{choice, choose_index, Choice, WeightedChoice}; use crate::schema::{ @@ -11,11 +17,13 @@ const MAX_EVENT_COLUMNS: usize = 32; const MAX_TABLE_COLUMNS: usize = 32; const MAX_TABLES: usize = 128; +/// A fully materialized target schema for one migration interaction. #[derive(Debug, Clone, PartialEq, Eq)] pub struct Migration { schema: SchemaPlan, } +/// A deterministic schema edit used while building a migration candidate. #[derive(Debug, Clone, PartialEq, Eq)] pub enum SchemaRewrite { AddTable { table: TablePlan }, @@ -23,6 +31,7 @@ pub enum SchemaRewrite { AlterTable { table: String, ops: Vec }, } +/// Table-local migration operations generated by the DST harness. #[derive(Debug, Clone, PartialEq, Eq)] pub enum TableMigrationOp { ChangeAccess, @@ -60,6 +69,7 @@ pub enum TableMigrationOp { ReschemaEventTable, } +/// Weighted categories of auto-migration surfaces to probe. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum MigrationChoice { AddTable, diff --git a/crates/dst/src/engine/workload.rs b/crates/dst/src/engine/workload.rs index f77933a3401..a1e986c8a9c 100644 --- a/crates/dst/src/engine/workload.rs +++ b/crates/dst/src/engine/workload.rs @@ -1,3 +1,5 @@ +//! Workload interaction generation for the engine DST driver. + use std::fmt::{Debug, Error, Formatter}; use super::generation::GenCtx; @@ -9,6 +11,7 @@ use crate::rng::{choice, pick_choice, Choice}; use crate::schema::SchemaPlan; use spacetimedb_runtime::sim::Rng; +/// One generated action for the engine target to execute. #[derive(Debug, Clone)] pub enum Interaction { BeginMutTx, @@ -19,6 +22,7 @@ pub enum Interaction { Replay, } +/// Counts of emitted workload interactions, reported at the end of each run. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct InteractionCounts { pub total: usize, @@ -45,6 +49,7 @@ impl InteractionCounts { } } +/// Observable result of executing an interaction against the engine. #[derive(Debug, Clone, PartialEq, Eq)] pub enum Observation { BeganMutTx, @@ -61,6 +66,7 @@ pub enum InsertOutcome { UniqueConstraintViolation, } +/// Runtime-tunable weights for top-level workload actions. #[derive(Debug, Clone, Copy)] pub struct InteractionWeights { pub insert: u64, @@ -82,27 +88,7 @@ impl Default for InteractionWeights { } } -#[derive(Debug, Clone, Copy)] -enum InteractionChoice { - Insert, - Delete, - CommitTx, - Migrate, - Replay, -} - -impl InteractionWeights { - fn choices(self) -> [Choice; 5] { - [ - choice(self.insert, InteractionChoice::Insert), - choice(self.delete, InteractionChoice::Delete), - choice(self.commit_tx, InteractionChoice::CommitTx), - choice(self.migrate, InteractionChoice::Migrate), - choice(self.replay, InteractionChoice::Replay), - ] - } -} - +/// Stateful iterator that emits interactions and mirrors them into the model. pub struct WorkloadGen { rng: Rng, model: Model, @@ -128,6 +114,39 @@ impl WorkloadGen { self.stats } + pub fn next_interaction(&mut self) -> Interaction { + let choice = self.pick_interaction_choice(); + let interaction = self.interaction_from_choice(choice); + + self.model.apply(&interaction); + self.stats.record(&interaction); + + interaction + } +} + +#[derive(Debug, Clone, Copy)] +enum InteractionChoice { + Insert, + Delete, + CommitTx, + Migrate, + Replay, +} + +impl InteractionWeights { + fn choices(self) -> [Choice; 5] { + [ + choice(self.insert, InteractionChoice::Insert), + choice(self.delete, InteractionChoice::Delete), + choice(self.commit_tx, InteractionChoice::CommitTx), + choice(self.migrate, InteractionChoice::Migrate), + choice(self.replay, InteractionChoice::Replay), + ] + } +} + +impl WorkloadGen { fn schema(&self) -> &SchemaPlan { self.model.schema() } @@ -139,16 +158,6 @@ impl WorkloadGen { }) } - pub fn next_interaction(&mut self) -> Interaction { - let choice = self.pick_interaction_choice(); - let interaction = self.interaction_from_choice(choice); - - self.model.apply(&interaction); - self.stats.record(&interaction); - - interaction - } - fn interaction_from_choice(&mut self, choice: InteractionChoice) -> Interaction { if !self.model.in_mut_tx() { return match choice { diff --git a/crates/dst/src/rng.rs b/crates/dst/src/rng.rs index 084ed909544..626d53b87d5 100644 --- a/crates/dst/src/rng.rs +++ b/crates/dst/src/rng.rs @@ -1,5 +1,8 @@ +//! Deterministic random-selection helpers shared by DST generators. + use spacetimedb_runtime::sim::Rng; +/// A weighted value in a deterministic choice table. #[derive(Clone, Copy)] pub(crate) struct Choice { weight: u64, @@ -12,10 +15,12 @@ impl Choice { } } +/// Construct a weighted choice entry. pub(crate) const fn choice(weight: u64, value: T) -> Choice { Choice { weight, value } } +/// Pick one value from `choices`, using each entry's relative weight. pub(crate) fn pick_choice(rng: &Rng, choices: &[Choice]) -> T { let total: u64 = choices.iter().map(|choice| choice.weight).sum(); @@ -34,6 +39,7 @@ pub(crate) fn pick_choice(rng: &Rng, choices: &[Choice]) -> T { unreachable!("selected value is always inside total weight") } +/// Static weighted choices for enum-like generator cases. pub(crate) trait WeightedChoice: Copy + 'static { const CHOICES: &'static [Choice]; @@ -42,10 +48,12 @@ pub(crate) trait WeightedChoice: Copy + 'static { } } +/// Pick an index from `0..len`, or return `None` for an empty collection. pub(crate) fn choose_index(rng: &Rng, len: usize) -> Option { (len > 0).then(|| rng.index(len)) } +/// Pick a value in the inclusive range `lo..=hi`. pub(crate) fn range_inclusive(rng: &Rng, lo: usize, hi: usize) -> usize { if lo >= hi { return lo; diff --git a/crates/dst/src/schema.rs b/crates/dst/src/schema.rs index 5ced13ed45a..b8d70463de2 100644 --- a/crates/dst/src/schema.rs +++ b/crates/dst/src/schema.rs @@ -1,3 +1,5 @@ +//! Schema plans and raw module lowering for the engine DST harness. + use crate::rng; use spacetimedb_lib::db::raw_def::v10::*; use spacetimedb_lib::db::raw_def::v9::{RawIndexAlgorithm, TableAccess, TableType}; @@ -7,10 +9,29 @@ use spacetimedb_sats::{ AlgebraicType, AlgebraicValue, ArrayType, ArrayValue, ProductType, ProductTypeElement, SumType, SumTypeVariant, }; +/// Generate the default engine DST schema. +/// +/// The initial schema is intentionally random and valid, not repaired into a +/// fixed coverage fixture. Long runs are expected to discover more surfaces via +/// migrations. pub fn default_schema(rng: Rng) -> SchemaPlan { SchemaGenerator::new(rng, SchemaProfile::engine_dst()).gen_schema() } +/// Lower a generated schema plan into the raw module format used by the engine. +pub fn to_raw_def(schema: &SchemaPlan) -> RawModuleDefV10 { + let mut builder = RawModuleDefV10Builder::new(); + builder.set_case_conversion_policy(CaseConversionPolicy::None); + + for table in &schema.tables { + to_raw_def_table(&mut builder, table); + } + + let mut raw = builder.finish(); + apply_sequence_bounds(schema, &mut raw); + raw +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum Type { Bool, @@ -22,6 +43,7 @@ pub enum Type { } impl Type { + /// Representative column types used by migration candidate generation. pub const ALL: &'static [Type] = &[ Type::Bool, Type::I64, @@ -68,118 +90,7 @@ impl Type { } } -pub struct SchemaDecisions; - -impl SchemaDecisions { - pub fn range(rng: &Rng, (lo, hi): (usize, usize)) -> usize { - rng::range_inclusive(rng, lo, hi) - } - - pub fn index(rng: &Rng, len: usize) -> usize { - rng::choose_index(rng, len).expect("len must be non-zero") - } - - pub fn choose_index(rng: &Rng, len: usize) -> Option { - rng::choose_index(rng, len) - } - - pub fn sample_probability(rng: &Rng, probability: f64) -> bool { - rng.sample_probability(probability) - } - - pub fn gen_table_name(rng: &Rng, tables: &[TablePlan]) -> String { - loop { - let name = format!("tbl_{}", Self::gen_ident(rng)); - if tables.iter().all(|table| table.name != name) { - return name; - } - } - } - - pub fn gen_column_name(rng: &Rng, seen: &[String]) -> String { - loop { - let name = Self::gen_ident(rng); - if !seen.contains(&name) { - return name; - } - } - } - - fn gen_ident(rng: &Rng) -> String { - const CHARS: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789_"; - const FIRST: &[u8] = b"abcdefghijklmnopqrstuvwxyz_"; - let len = 4 + (rng.next_u64() as usize % 12); - let mut s = String::with_capacity(len); - s.push(FIRST[Self::index(rng, FIRST.len())] as char); - for _ in 1..len { - s.push(CHARS[Self::index(rng, CHARS.len())] as char); - } - s - } -} - -pub struct SchemaNames; - -impl SchemaNames { - pub fn fresh_column_name(table: &TablePlan, base: &str) -> String { - if table.columns.iter().all(|column| column.name != base) { - return base.into(); - } - - for suffix in 0.. { - let candidate = format!("{base}_{suffix}"); - if table.columns.iter().all(|column| column.name != candidate) { - return candidate; - } - } - - unreachable!("unbounded suffix search must find a unique column name") - } - - pub fn fresh_table_name(tables: &[TablePlan], base: &str) -> String { - if tables.iter().all(|table| table.name != base) { - return base.into(); - } - - for suffix in 0.. { - let candidate = format!("{base}_{suffix}"); - if tables.iter().all(|table| table.name != candidate) { - return candidate; - } - } - - unreachable!("unbounded suffix search must find a unique table name") - } - - pub fn index_name(table: &TablePlan, index: &IndexPlan) -> String { - format!( - "{}_{}_idx", - table.name, - index - .columns - .iter() - .map(|&c| table.columns[c].name.as_str()) - .collect::>() - .join("_") - ) - } - - pub fn constraint_name(table: &TablePlan, constraint: &UniqueConstraintPlan) -> String { - format!( - "{}_{}_key", - table.name, - constraint - .columns - .iter() - .map(|&c| table.columns[c].name.as_str()) - .collect::>() - .join("_") - ) - } -} - -// Schema plan — the canonical source of truth. -// This Schema should be able to translate to valid `RawModuleDefV10`. +/// Canonical schema representation for the DST model and engine target. #[derive(Debug, Clone, PartialEq, Eq)] pub struct SchemaPlan { pub tables: Vec, @@ -268,6 +179,11 @@ impl SequencePlan { }) } + /// Build a small bounded sequence whose max is an already-observed value. + /// + /// This intentionally creates a high-risk migration surface: if migration + /// prechecks accept an unsafe sequence, later inserts can collide with + /// existing unique values. pub fn with_existing_value_as_max(column: usize, ty: Type, existing_value: i128) -> Option { const DOMAIN_SIZE: i128 = 3; @@ -280,89 +196,6 @@ impl SequencePlan { } } -// Lowering into RawModuleDefV10. -pub fn to_raw_def(schema: &SchemaPlan) -> RawModuleDefV10 { - let mut builder = RawModuleDefV10Builder::new(); - builder.set_case_conversion_policy(CaseConversionPolicy::None); - - for table in &schema.tables { - to_raw_def_table(&mut builder, table); - } - - let mut raw = builder.finish(); - apply_sequence_bounds(schema, &mut raw); - raw -} - -fn apply_sequence_bounds(schema: &SchemaPlan, raw: &mut RawModuleDefV10) { - for (table_plan, raw_table) in schema.tables.iter().zip(raw.tables_mut_for_tests().iter_mut()) { - for (sequence_plan, raw_sequence) in table_plan.sequences.iter().zip(raw_table.sequences.iter_mut()) { - raw_sequence.start = sequence_plan.start; - raw_sequence.min_value = sequence_plan.min_value; - raw_sequence.max_value = sequence_plan.max_value; - raw_sequence.increment = sequence_plan.increment; - } - } -} - -fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { - let product_type = ProductType { - elements: table - .columns - .iter() - .map(|col| ProductTypeElement { - name: Some(col.name.clone().into()), - algebraic_type: col.ty.to_algebraic(), - }) - .collect(), - }; - - let mut tbl = builder.build_table_with_new_type_for_tests(table.name.clone(), product_type, true); - - tbl = tbl.with_type(TableType::User); - tbl = tbl.with_event(table.is_event); - tbl = tbl.with_access(if table.is_public { - TableAccess::Public - } else { - TableAccess::Private - }); - // Primary key. - if let Some(pk) = table.primary_key { - tbl = tbl.with_primary_key(ColId(pk as u16)); - } - - // Unique constraints — all of them, including PK-matching. - for constraint in &table.unique_constraints { - let col_list: ColList = constraint.columns.iter().map(|&c| ColId(c as u16)).collect(); - tbl = tbl.with_unique_constraint(col_list); - } - - // Indexes. - for index in &table.indexes { - let col_list: ColList = index.columns.iter().map(|&c| ColId(c as u16)).collect(); - - let algorithm = match index.algorithm { - IndexAlgorithm::BTree => RawIndexAlgorithm::BTree { columns: col_list }, - IndexAlgorithm::Hash => RawIndexAlgorithm::Hash { columns: col_list }, - }; - - tbl = tbl.with_index_no_accessor_name(algorithm, SchemaNames::index_name(table, index)); - } - - // Sequences — all of them. - for seq in &table.sequences { - tbl = tbl.with_column_sequence(ColId(seq.column as u16)); - } - - // AddColumns needs defaults when existing rows are present. Supplying stable - // defaults for all columns lets the engine keep only the newly-added tail. - for (col_id, column) in table.columns.iter().enumerate() { - tbl = tbl.with_default_column_value(ColId(col_id as u16), column.ty.default_value()); - } - - tbl.finish(); -} - /// Controls the shape of generated schemas. #[derive(Debug, Clone)] pub struct SchemaProfile { @@ -438,6 +271,74 @@ impl Default for TypeWeights { } } +/// Random schema generator used by initial schema creation and add-table migrations. +pub struct SchemaGenerator { + rng: Rng, + profile: SchemaProfile, +} + +impl SchemaGenerator { + pub fn new(rng: Rng, profile: SchemaProfile) -> Self { + Self { rng, profile } + } + + /// Generate one table compatible with an existing schema. + /// + /// This is used by migration generation so add-table migrations use the same + /// shape policy as initial schema generation. Randomness happens before the + /// rewrite is applied; the rewrite itself remains deterministic. + pub fn gen_table_for_schema(&self, schema: &SchemaPlan, is_event: bool) -> TablePlan { + let mut sum_available = !schema_has_sum_column(schema); + self.gen_table(&schema.tables, is_event, &mut sum_available) + } + + /// Generate a complete schema from this generator's profile. + pub fn gen_schema(&self) -> SchemaPlan { + let table_count = SchemaDecisions::range(&self.rng, self.profile.table_count); + let mut tables: Vec = Vec::with_capacity(table_count); + let mut sum_available = true; + for table_idx in 0..table_count { + let must_be_data = table_idx + 1 == table_count && !tables.iter().any(|table| !table.is_event); + let is_event = !must_be_data && matches!(self.gen_table_kind(), TableKind::Event); + tables.push(self.gen_table(&tables, is_event, &mut sum_available)); + } + SchemaPlan { tables } + } +} + +/// Stable naming helpers for generated migration artifacts. +pub struct SchemaNames; + +impl SchemaNames { + pub fn fresh_column_name(table: &TablePlan, base: &str) -> String { + if table.columns.iter().all(|column| column.name != base) { + return base.into(); + } + + for suffix in 0.. { + let candidate = format!("{base}_{suffix}"); + if table.columns.iter().all(|column| column.name != candidate) { + return candidate; + } + } + + unreachable!("unbounded suffix search must find a unique column name") + } + + pub fn index_name(table: &TablePlan, index: &IndexPlan) -> String { + format!( + "{}_{}_idx", + table.name, + index + .columns + .iter() + .map(|&c| table.columns[c].name.as_str()) + .collect::>() + .join("_") + ) + } +} + #[derive(Debug, Clone, Copy)] enum TableKind { Data, @@ -486,16 +387,53 @@ impl TypeWeights { } } -pub struct SchemaGenerator { - rng: Rng, - profile: SchemaProfile, -} +struct SchemaDecisions; -impl SchemaGenerator { - pub fn new(rng: Rng, profile: SchemaProfile) -> Self { - Self { rng, profile } +impl SchemaDecisions { + fn range(rng: &Rng, (lo, hi): (usize, usize)) -> usize { + rng::range_inclusive(rng, lo, hi) + } + + fn index(rng: &Rng, len: usize) -> usize { + rng::choose_index(rng, len).expect("len must be non-zero") + } + + fn sample_probability(rng: &Rng, probability: f64) -> bool { + rng.sample_probability(probability) + } + + fn gen_table_name(rng: &Rng, tables: &[TablePlan]) -> String { + loop { + let name = format!("tbl_{}", Self::gen_ident(rng)); + if tables.iter().all(|table| table.name != name) { + return name; + } + } + } + + fn gen_column_name(rng: &Rng, seen: &[String]) -> String { + loop { + let name = Self::gen_ident(rng); + if !seen.contains(&name) { + return name; + } + } } + fn gen_ident(rng: &Rng) -> String { + const CHARS: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789_"; + const FIRST: &[u8] = b"abcdefghijklmnopqrstuvwxyz_"; + let len = 4 + (rng.next_u64() as usize % 12); + let mut s = String::with_capacity(len); + s.push(FIRST[Self::index(rng, FIRST.len())] as char); + for _ in 1..len { + s.push(CHARS[Self::index(rng, CHARS.len())] as char); + } + s + } +} + +impl SchemaGenerator { fn gen_columns(&self, sum_available: &mut bool) -> Vec { let n = SchemaDecisions::range(&self.rng, self.profile.columns); let mut names = Vec::with_capacity(n); @@ -551,7 +489,7 @@ impl SchemaGenerator { result.push(UniqueConstraintPlan { columns: cols }); } } - // Ensure PK has a matching unique constraint. + // A primary key always has a matching unique constraint. if let Some(pk) = pk { if !seen.iter().any(|cols| cols.len() == 1 && cols[0] == *pk) { result.push(UniqueConstraintPlan { columns: vec![*pk] }); @@ -566,11 +504,9 @@ impl SchemaGenerator { unique_constraints: &[UniqueConstraintPlan], pk: &Option, ) -> Vec { - // Every unique constraint and PK needs a matching index. let mut seen_cols: Vec> = Vec::new(); let mut indexes: Vec = Vec::new(); - // Index for PK. if let Some(pk) = pk { seen_cols.push(vec![*pk]); indexes.push(IndexPlan { @@ -579,7 +515,6 @@ impl SchemaGenerator { }); } - // Indexes for unique constraints. for constraint in unique_constraints { if seen_cols.contains(&constraint.columns) { continue; @@ -591,7 +526,6 @@ impl SchemaGenerator { }); } - // Additional random indexes. let n = SchemaDecisions::range(&self.rng, self.profile.indexes); for _ in 0..n { let num_cols = 1 + SchemaDecisions::index(&self.rng, columns.len().min(3)); @@ -618,11 +552,6 @@ impl SchemaGenerator { indexes } - pub fn gen_table_for_schema(&self, schema: &SchemaPlan, is_event: bool) -> TablePlan { - let mut sum_available = !schema_has_sum_column(schema); - self.gen_table(&schema.tables, is_event, &mut sum_available) - } - fn gen_table(&self, existing_tables: &[TablePlan], is_event: bool, sum_available: &mut bool) -> TablePlan { let columns = self.gen_columns(sum_available); let name = SchemaDecisions::gen_table_name(&self.rng, existing_tables); @@ -680,20 +609,74 @@ impl SchemaGenerator { let choices = self.profile.table_kind_weights.choices(); rng::pick_choice(&self.rng, &choices) } +} - pub fn gen_schema(&self) -> SchemaPlan { - let table_count = SchemaDecisions::range(&self.rng, self.profile.table_count); - let mut tables: Vec = Vec::with_capacity(table_count); - let mut sum_available = true; - for table_idx in 0..table_count { - let must_be_data = table_idx + 1 == table_count && !tables.iter().any(|table| !table.is_event); - let is_event = !must_be_data && matches!(self.gen_table_kind(), TableKind::Event); - tables.push(self.gen_table(&tables, is_event, &mut sum_available)); +fn apply_sequence_bounds(schema: &SchemaPlan, raw: &mut RawModuleDefV10) { + for (table_plan, raw_table) in schema.tables.iter().zip(raw.tables_mut_for_tests().iter_mut()) { + for (sequence_plan, raw_sequence) in table_plan.sequences.iter().zip(raw_table.sequences.iter_mut()) { + raw_sequence.start = sequence_plan.start; + raw_sequence.min_value = sequence_plan.min_value; + raw_sequence.max_value = sequence_plan.max_value; + raw_sequence.increment = sequence_plan.increment; } - SchemaPlan { tables } } } +fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { + let product_type = ProductType { + elements: table + .columns + .iter() + .map(|col| ProductTypeElement { + name: Some(col.name.clone().into()), + algebraic_type: col.ty.to_algebraic(), + }) + .collect(), + }; + + let mut tbl = builder.build_table_with_new_type_for_tests(table.name.clone(), product_type, true); + + tbl = tbl.with_type(TableType::User); + tbl = tbl.with_event(table.is_event); + tbl = tbl.with_access(if table.is_public { + TableAccess::Public + } else { + TableAccess::Private + }); + + if let Some(pk) = table.primary_key { + tbl = tbl.with_primary_key(ColId(pk as u16)); + } + + for constraint in &table.unique_constraints { + let col_list: ColList = constraint.columns.iter().map(|&c| ColId(c as u16)).collect(); + tbl = tbl.with_unique_constraint(col_list); + } + + for index in &table.indexes { + let col_list: ColList = index.columns.iter().map(|&c| ColId(c as u16)).collect(); + + let algorithm = match index.algorithm { + IndexAlgorithm::BTree => RawIndexAlgorithm::BTree { columns: col_list }, + IndexAlgorithm::Hash => RawIndexAlgorithm::Hash { columns: col_list }, + }; + + tbl = tbl.with_index_no_accessor_name(algorithm, SchemaNames::index_name(table, index)); + } + + for seq in &table.sequences { + tbl = tbl.with_column_sequence(ColId(seq.column as u16)); + } + + // AddColumns needs defaults when existing rows are present. Supplying stable + // defaults for all columns lets the engine keep only the newly-added tail. + for (col_id, column) in table.columns.iter().enumerate() { + tbl = tbl.with_default_column_value(ColId(col_id as u16), column.ty.default_value()); + } + + tbl.finish(); +} + fn schema_has_sum_column(schema: &SchemaPlan) -> bool { schema .tables From d68103b584b3b9cdfbede9799383937f5bbcd44d Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Tue, 14 Jul 2026 16:31:37 +0530 Subject: [PATCH 10/10] better error messages --- crates/dst/src/engine.rs | 23 ++++++++++++++--------- crates/dst/src/engine/migrations.rs | 17 +++++++++++++++-- crates/dst/src/engine/model.rs | 8 +++----- crates/dst/src/engine/properties.rs | 18 +++++++++++------- crates/dst/src/engine/workload.rs | 2 +- crates/dst/src/traits.rs | 23 ++++++++++++++++------- 6 files changed, 60 insertions(+), 31 deletions(-) diff --git a/crates/dst/src/engine.rs b/crates/dst/src/engine.rs index fd0822ff1aa..43713c4fce6 100644 --- a/crates/dst/src/engine.rs +++ b/crates/dst/src/engine.rs @@ -185,11 +185,13 @@ impl EngineTarget { }) } - fn is_unique_constraint_violation(error: &DBError) -> bool { - matches!( - error, - DBError::Datastore(DatastoreError::Index(IndexError::UniqueConstraintViolation(_))) - ) + fn unique_constraint_violation_details(error: &DBError) -> Option { + match error { + DBError::Datastore(DatastoreError::Index(IndexError::UniqueConstraintViolation(violation))) => { + Some(violation.to_string()) + } + _ => None, + } } fn commit_delta_from_tx_data(&self, tx_data: &TxData) -> CommitDelta { @@ -277,11 +279,14 @@ impl EngineTarget { .ok_or_else(|| anyhow::anyhow!("insert without active mutable transaction"))?; let outcome = match db.insert(tx, table_id, &bytes) { Ok((_generated_columns, row, _flags)) => InsertOutcome::Accepted(row.to_product_value()), - // Generated rows can intentionally hit unique constraints; the oracle validates that rejection. - Err(error) if Self::is_unique_constraint_violation(&error) => { - InsertOutcome::UniqueConstraintViolation + Err(error) => { + // Generated rows can intentionally hit unique constraints; the oracle validates that rejection. + if let Some(details) = Self::unique_constraint_violation_details(&error) { + InsertOutcome::UniqueConstraintViolation { details } + } else { + return Err(error.into()); + } } - Err(error) => return Err(error.into()), }; Ok(Observation::Inserted { outcome }) } diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs index 087699cda53..f120c4cbb9b 100644 --- a/crates/dst/src/engine/migrations.rs +++ b/crates/dst/src/engine/migrations.rs @@ -628,6 +628,7 @@ fn addable_sequence_boundary_probes( .filter_map(|(column, column_plan)| { let domain = column_domain(column)?; if column_plan.ty != Type::U64 + || table.sequences.iter().any(|sequence| sequence.column == column) || domain.sequenced || !domain.single_column_indexed || !domain.single_column_unique @@ -635,12 +636,24 @@ fn addable_sequence_boundary_probes( return None; } - let max_value = domain.positive_i128_value_above(2)?; - SequencePlan::with_existing_value_as_max(column, column_plan.ty, max_value) + domain + .integral_values() + .filter_map(|max_value| SequencePlan::with_existing_value_as_max(column, column_plan.ty, max_value)) + .find(|sequence| added_sequence_precheck_range_is_clear(&domain, sequence)) }) .collect() } +fn added_sequence_precheck_range_is_clear(domain: &ColumnDomain, sequence: &SequencePlan) -> bool { + let min = sequence.min_value.unwrap_or(1); + let max = sequence.max_value.unwrap_or(i128::MAX); + + // The engine's add-sequence precheck rejects existing values in `min..max`. + // The boundary probe intentionally places an existing value at `max`, so + // only values below the exclusive upper bound make the migration invalid. + domain.integral_values().all(|value| value < min || value >= max) +} + fn column_domain(model: &Model, table: &TablePlan, column: usize) -> Option { let column_name = &table.columns.get(column)?.name; model.column_domain_by_name(&table.name, column_name) diff --git a/crates/dst/src/engine/model.rs b/crates/dst/src/engine/model.rs index e6a194765b9..43f4eecd897 100644 --- a/crates/dst/src/engine/model.rs +++ b/crates/dst/src/engine/model.rs @@ -41,10 +41,6 @@ impl ColumnDomain { _ => None, }) } - - pub(crate) fn positive_i128_value_above(&self, min_exclusive: i128) -> Option { - self.integral_values().find(|value| *value > min_exclusive) - } } // Keep mutable transactions as an overlay: committed rows stay shared, while @@ -159,7 +155,9 @@ impl Model { if self.violates_unique_constraint(*table, row) { return Observation::Inserted { - outcome: InsertOutcome::UniqueConstraintViolation, + outcome: InsertOutcome::UniqueConstraintViolation { + details: "model unique constraint".into(), + }, }; } diff --git a/crates/dst/src/engine/properties.rs b/crates/dst/src/engine/properties.rs index 9511c37de48..81695f4ef0f 100644 --- a/crates/dst/src/engine/properties.rs +++ b/crates/dst/src/engine/properties.rs @@ -65,7 +65,7 @@ impl EngineOracle { ( Interaction::Insert { .. }, Observation::Inserted { - outcome: InsertOutcome::UniqueConstraintViolation, + outcome: InsertOutcome::UniqueConstraintViolation { .. }, }, ) => self.model.apply(interaction), (Interaction::Insert { .. }, _) => anyhow::bail!("insert produced unexpected observation"), @@ -92,7 +92,7 @@ impl EngineProperty for InsertMatches { fn check( &self, - _interaction: &Interaction, + interaction: &Interaction, observation: &Observation, expected: &Observation, ) -> anyhow::Result<()> { @@ -107,12 +107,16 @@ impl EngineProperty for InsertMatches { (InsertOutcome::Accepted(row), InsertOutcome::Accepted(expected)) => { anyhow::ensure!(row == expected, "insert_matches: accepted row diverged from model"); } - (InsertOutcome::UniqueConstraintViolation, InsertOutcome::UniqueConstraintViolation) => {} - (InsertOutcome::Accepted(_), InsertOutcome::UniqueConstraintViolation) => { - anyhow::bail!("insert_matches: target accepted row rejected by model"); + (InsertOutcome::UniqueConstraintViolation { .. }, InsertOutcome::UniqueConstraintViolation { .. }) => {} + (InsertOutcome::Accepted(_), InsertOutcome::UniqueConstraintViolation { .. }) => { + anyhow::bail!( + "insert_matches: target accepted row rejected by model\ninteraction: {interaction:#?}\ntarget: {observation:#?}\nmodel: {expected:#?}" + ); } - (InsertOutcome::UniqueConstraintViolation, InsertOutcome::Accepted(_)) => { - anyhow::bail!("insert_matches: target rejected row accepted by model"); + (InsertOutcome::UniqueConstraintViolation { .. }, InsertOutcome::Accepted(_)) => { + anyhow::bail!( + "insert_matches: target rejected row accepted by model\ninteraction: {interaction:#?}\ntarget: {observation:#?}\nmodel: {expected:#?}" + ); } } diff --git a/crates/dst/src/engine/workload.rs b/crates/dst/src/engine/workload.rs index a1e986c8a9c..efa3fbe54f0 100644 --- a/crates/dst/src/engine/workload.rs +++ b/crates/dst/src/engine/workload.rs @@ -63,7 +63,7 @@ pub enum Observation { #[derive(Debug, Clone, PartialEq, Eq)] pub enum InsertOutcome { Accepted(Row), - UniqueConstraintViolation, + UniqueConstraintViolation { details: String }, } /// Runtime-tunable weights for top-level workload actions. diff --git a/crates/dst/src/traits.rs b/crates/dst/src/traits.rs index 5abfc770092..af4cfb89d46 100644 --- a/crates/dst/src/traits.rs +++ b/crates/dst/src/traits.rs @@ -1,6 +1,9 @@ -use std::panic::{resume_unwind, AssertUnwindSafe}; +use std::{ + fmt::Debug, + panic::{resume_unwind, AssertUnwindSafe}, +}; -use anyhow::Error; +use anyhow::{Context, Error}; use futures::FutureExt; use spacetimedb_runtime::sim::Rng; @@ -26,8 +29,8 @@ pub type TestSuiteParts = ( ); pub trait TestSuite { - type Interaction: std::fmt::Debug; - type Interactions: Iterator + std::fmt::Debug; + type Interaction: Debug; + type Interactions: Iterator + Debug; type Target: TargetDriver; type Properties: Properties>::Observation>; @@ -43,9 +46,15 @@ pub trait TestSuite { let (mut interactions, mut target, mut properties) = self.build(rng).await?; let result = AssertUnwindSafe(async { - for interaction in interactions.by_ref().take(max_interactions) { - let observation = target.execute(&interaction).await?; - properties.observe(&interaction, &observation)?; + for (step, interaction) in interactions.by_ref().take(max_interactions).enumerate() { + let observation = target + .execute(&interaction) + .await + .with_context(|| format!("DST target failed at interaction #{step}: {interaction:?}"))?; + + properties + .observe(&interaction, &observation) + .with_context(|| format!("DST property failed at interaction #{step}: {interaction:?}"))?; } Ok(())