Skip to content

Commit 91b33a7

Browse files
CentrilCentril
authored
Optimize TxData + DatabaseUpdate for fast construction (#4138)
# Description of Changes Rework `TxData` to: - store all information for a table together in a single type `TxDataTableEntry` rather than having several different maps. - be constructed as fast as possible. - fit every entry within a single cache line, i.e., `(TableId, TxDataTableEntry)` takes up 64 bytes. - fit a single table entry inline to optimize for small transactions. - expose methods `{inserts, deletes}_for_table` to make `DeltaTx` faster. Rework `DatabaseUpdate` to: - store a single `DatabaseTableUpdate` inline - make `from_writes` profit from the changes to `TxData`, avoiding the temporary hash map and allocating the necessary capacity from the start. # API and ABI breaking changes None # Expected complexity level and risk 3? Fairly simple changes, but in important places. # Testing Existing tests are changed to match the changes to `TxData`.
1 parent b181061 commit 91b33a7

9 files changed

Lines changed: 267 additions & 250 deletions

File tree

‎crates/core/src/db/durability.rs‎

Lines changed: 14 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,9 @@ use spacetimedb_commitlog::payload::{
66
txdata::{Mutations, Ops},
77
Txdata,
88
};
9-
use spacetimedb_data_structures::map::IntSet;
109
use spacetimedb_datastore::{execution_context::ReducerContext, traits::TxData};
1110
use spacetimedb_durability::{DurableOffset, TxOffset};
1211
use spacetimedb_lib::Identity;
13-
use spacetimedb_primitives::TableId;
1412
use tokio::{
1513
runtime,
1614
sync::{
@@ -208,32 +206,23 @@ impl DurabilityWorkerActor {
208206
return;
209207
}
210208

211-
let is_persistent_table = |table_id: &TableId| -> bool { !tx_data.is_ephemeral_table(table_id) };
212-
213-
let inserts: Box<_> = tx_data
214-
.inserts()
215-
// Skip ephemeral tables
216-
.filter(|(table_id, _)| is_persistent_table(table_id))
217-
.map(|(table_id, rowdata)| Ops {
218-
table_id: *table_id,
219-
rowdata: rowdata.clone(),
220-
})
209+
let mut inserts: Box<_> = tx_data
210+
.persistent_inserts()
211+
.map(|(table_id, rowdata)| Ops { table_id, rowdata })
221212
.collect();
222-
223-
let truncates: IntSet<TableId> = tx_data.truncates().collect();
224-
225-
let deletes: Box<_> = tx_data
226-
.deletes()
227-
.filter(|(table_id, _)| is_persistent_table(table_id))
228-
.map(|(table_id, rowdata)| Ops {
229-
table_id: *table_id,
230-
rowdata: rowdata.clone(),
231-
})
232-
// filter out deletes for tables that are truncated in the same transaction.
233-
.filter(|ops| !truncates.contains(&ops.table_id))
213+
// What we get from `tx_data` is not necessarily sorted,
214+
// but the durability layer expects by-table_id sorted data.
215+
// Unstable sorts are valid, there will only ever be one entry per table_id.
216+
inserts.sort_unstable_by_key(|ops| ops.table_id);
217+
218+
let mut deletes: Box<_> = tx_data
219+
.persistent_deletes()
220+
.map(|(table_id, rowdata)| Ops { table_id, rowdata })
234221
.collect();
222+
deletes.sort_unstable_by_key(|ops| ops.table_id);
235223

236-
let truncates: Box<_> = truncates.into_iter().filter(is_persistent_table).collect();
224+
let mut truncates: Box<[_]> = tx_data.persistent_truncates().collect();
225+
truncates.sort_unstable_by_key(|table_id| *table_id);
237226

238227
let inputs = reducer_context.map(|rcx| rcx.into());
239228

‎crates/core/src/host/module_host.rs‎

Lines changed: 12 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -37,13 +37,13 @@ use indexmap::IndexSet;
3737
use itertools::Itertools;
3838
use prometheus::{Histogram, IntGauge};
3939
use scopeguard::ScopeGuard;
40+
use smallvec::SmallVec;
4041
use spacetimedb_auth::identity::ConnectionAuthCtx;
4142
use spacetimedb_client_api_messages::energy::FunctionBudget;
4243
use spacetimedb_client_api_messages::websocket::{
4344
ByteListLen, Compression, OneOffTable, QueryUpdate, Subscribe, SubscribeMulti, SubscribeSingle,
4445
};
4546
use spacetimedb_data_structures::error_stream::ErrorStream;
46-
use spacetimedb_data_structures::map::{HashCollectionExt as _, IntMap};
4747
use spacetimedb_datastore::error::DatastoreError;
4848
use spacetimedb_datastore::execution_context::{Workload, WorkloadType};
4949
use spacetimedb_datastore::locking_tx_datastore::{MutTxId, ViewCallInfo};
@@ -54,11 +54,10 @@ use spacetimedb_expr::expr::CollectViews;
5454
use spacetimedb_lib::db::raw_def::v9::Lifecycle;
5555
use spacetimedb_lib::identity::{AuthCtx, RequestId};
5656
use spacetimedb_lib::metrics::ExecutionMetrics;
57-
use spacetimedb_lib::Timestamp;
58-
use spacetimedb_lib::{AlgebraicType, ConnectionId};
57+
use spacetimedb_lib::{ConnectionId, Timestamp};
5958
use spacetimedb_primitives::{ArgId, ProcedureId, TableId, ViewFnPtr, ViewId};
6059
use spacetimedb_query::compile_subscription;
61-
use spacetimedb_sats::{AlgebraicTypeRef, ProductValue};
60+
use spacetimedb_sats::{AlgebraicType, AlgebraicTypeRef, ProductValue};
6261
use spacetimedb_schema::auto_migrate::{AutoMigrateError, MigrationPolicy};
6362
use spacetimedb_schema::def::{ModuleDef, ProcedureDef, ReducerDef, TableDef, ViewDef};
6463
use spacetimedb_schema::schema::{Schema, TableSchema};
@@ -73,7 +72,7 @@ use tokio::sync::oneshot;
7372

7473
#[derive(Debug, Default, Clone, From)]
7574
pub struct DatabaseUpdate {
76-
pub tables: Vec<DatabaseTableUpdate>,
75+
pub tables: SmallVec<[DatabaseTableUpdate; 1]>,
7776
}
7877

7978
impl FromIterator<DatabaseTableUpdate> for DatabaseUpdate {
@@ -93,26 +92,15 @@ impl DatabaseUpdate {
9392
}
9493

9594
pub fn from_writes(tx_data: &TxData) -> Self {
96-
let mut map: IntMap<TableId, DatabaseTableUpdate> = IntMap::new();
97-
let new_update = |table_id, table_name: &str| DatabaseTableUpdate {
95+
let entries = tx_data.iter_table_entries();
96+
let mut tables = SmallVec::with_capacity(entries.len());
97+
tables.extend(entries.map(|(table_id, e)| DatabaseTableUpdate {
9898
table_id,
99-
table_name: table_name.into(),
100-
inserts: [].into(),
101-
deletes: [].into(),
102-
};
103-
for (table_id, table_name, rows) in tx_data.inserts_with_table_name() {
104-
map.entry(*table_id)
105-
.or_insert_with(|| new_update(*table_id, table_name))
106-
.inserts = rows.clone();
107-
}
108-
for (table_id, table_name, rows) in tx_data.deletes_with_table_name() {
109-
map.entry(*table_id)
110-
.or_insert_with(|| new_update(*table_id, table_name))
111-
.deletes = rows.clone();
112-
}
113-
DatabaseUpdate {
114-
tables: map.into_values().collect(),
115-
}
99+
table_name: e.table_name.clone(),
100+
inserts: e.inserts.clone(),
101+
deletes: e.deletes.clone(),
102+
}));
103+
DatabaseUpdate { tables }
116104
}
117105

118106
/// The number of rows in the payload

‎crates/core/src/sql/execute.rs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use crate::subscription::tx::DeltaTx;
1717
use crate::util::slow::SlowQueryLogger;
1818
use crate::vm::{check_row_limit, DbProgram, TxMode};
1919
use anyhow::anyhow;
20+
use smallvec::SmallVec;
2021
use spacetimedb_datastore::execution_context::Workload;
2122
use spacetimedb_datastore::locking_tx_datastore::state_view::StateView;
2223
use spacetimedb_datastore::traits::IsolationLevel;
@@ -42,7 +43,7 @@ pub struct StmtResult {
4243

4344
pub(crate) fn collect_result(
4445
result: &mut Vec<MemTable>,
45-
updates: &mut Vec<DatabaseTableUpdate>,
46+
updates: &mut SmallVec<[DatabaseTableUpdate; 1]>,
4647
r: CodeResult,
4748
) -> Result<(), DBError> {
4849
match r {
@@ -74,7 +75,7 @@ fn execute(
7475
p: &mut DbProgram<'_, '_>,
7576
ast: Vec<CrudExpr>,
7677
sql: &str,
77-
updates: &mut Vec<DatabaseTableUpdate>,
78+
updates: &mut SmallVec<[DatabaseTableUpdate; 1]>,
7879
) -> Result<Vec<MemTable>, DBError> {
7980
let slow_query_threshold = if let TxMode::Tx(tx) = p.tx {
8081
p.db.query_limit(tx)?.map(Duration::from_millis)
@@ -102,7 +103,7 @@ pub fn execute_sql(
102103
subs: Option<&ModuleSubscriptions>,
103104
) -> Result<Vec<MemTable>, DBError> {
104105
if CrudExpr::is_reads(&ast) {
105-
let mut updates = Vec::new();
106+
let mut updates = SmallVec::new();
106107
db.with_read_only(Workload::Sql, |tx| {
107108
execute(
108109
&mut DbProgram::new(db, &mut TxMode::Tx(tx), auth),
@@ -112,7 +113,7 @@ pub fn execute_sql(
112113
)
113114
})
114115
} else if subs.is_none() {
115-
let mut updates = Vec::new();
116+
let mut updates = SmallVec::new();
116117
db.with_auto_commit(Workload::Sql, |mut_tx| {
117118
execute(
118119
&mut DbProgram::new(db, &mut mut_tx.into(), auth),
@@ -123,7 +124,7 @@ pub fn execute_sql(
123124
})
124125
} else {
125126
let mut tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::Sql);
126-
let mut updates = Vec::with_capacity(ast.len());
127+
let mut updates = SmallVec::with_capacity(ast.len());
127128
let res = execute(
128129
&mut DbProgram::new(db, &mut (&mut tx).into(), auth.clone()),
129130
ast,
@@ -170,7 +171,7 @@ pub fn execute_sql_tx<'a>(
170171
return Ok(None);
171172
}
172173

173-
let mut updates = Vec::new(); // No subscription updates in this path, because it requires owning the tx.
174+
let mut updates = SmallVec::new(); // No subscription updates in this path, because it requires owning the tx.
174175
execute(&mut DbProgram::new(db, &mut tx, auth), ast, sql, &mut updates).map(Some)
175176
}
176177

‎crates/core/src/subscription/execution_unit.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@ use super::subscription::{IncrementalJoin, SupportedQuery};
33
use crate::db::relational_db::{RelationalDB, Tx};
44
use crate::error::DBError;
55
use crate::estimation;
6-
use crate::host::module_host::{DatabaseTableUpdate, DatabaseTableUpdateRelValue, UpdatesRelValue};
6+
use crate::host::module_host::DatabaseTableUpdate;
7+
use crate::host::module_host::{DatabaseTableUpdateRelValue, UpdatesRelValue};
78
use crate::messages::websocket::TableUpdate;
89
use crate::subscription::websocket_building::{BuildableWebsocketFormat, RowListBuilderSource};
910
use crate::util::slow::SlowQueryLogger;

‎crates/core/src/subscription/query.rs‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ mod tests {
162162
use crate::vm::tests::create_table_with_rows;
163163
use crate::vm::DbProgram;
164164
use itertools::Itertools;
165+
use smallvec::SmallVec;
165166
use spacetimedb_client_api_messages::websocket::{BsatnFormat, CompressableQueryUpdate, Compression};
166167
use spacetimedb_datastore::execution_context::Workload;
167168
use spacetimedb_lib::bsatn;
@@ -195,7 +196,7 @@ mod tests {
195196
let q = Expr::Crud(Box::new(CrudExpr::Query(query.clone())));
196197

197198
let mut result = Vec::with_capacity(1);
198-
let mut updates = Vec::new();
199+
let mut updates = SmallVec::new();
199200
collect_result(&mut result, &mut updates, run_ast(p, q, sources).into())?;
200201
Ok(result)
201202
}
@@ -445,12 +446,13 @@ mod tests {
445446
}
446447

447448
let update = DatabaseUpdate {
448-
tables: vec![DatabaseTableUpdate {
449+
tables: [DatabaseTableUpdate {
449450
table_id,
450451
table_name: "test".into(),
451452
deletes: deletes.into(),
452453
inserts: [].into(),
453-
}],
454+
}]
455+
.into(),
454456
};
455457

456458
db.commit_tx(tx)?;
@@ -535,7 +537,7 @@ mod tests {
535537
};
536538

537539
let update = DatabaseUpdate {
538-
tables: vec![data.clone()],
540+
tables: [data.clone()].into(),
539541
};
540542

541543
check_query_incr(&db, &tx, &s, &update, 1, &[row])?;
@@ -656,7 +658,7 @@ mod tests {
656658
};
657659

658660
let update = DatabaseUpdate {
659-
tables: vec![data1, data2],
661+
tables: smallvec::smallvec![data1, data2],
660662
};
661663

662664
let row_1 = product!(1u64, "health");
@@ -1161,9 +1163,9 @@ mod tests {
11611163
.collect::<Arc<_>>();
11621164

11631165
let tables = if inserts.is_empty() && deletes.is_empty() {
1164-
vec![]
1166+
smallvec::smallvec![]
11651167
} else {
1166-
vec![DatabaseTableUpdate {
1168+
smallvec::smallvec![DatabaseTableUpdate {
11671169
table_id,
11681170
table_name,
11691171
inserts,

‎crates/core/src/subscription/tx.rs‎

Lines changed: 10 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -46,21 +46,21 @@ impl DeltaTableIndexes {
4646
fn build_indexes_for_rows<'a>(
4747
tx: &'a TxId,
4848
meta: &'a QueriedTableIndexIds,
49-
rows: impl Iterator<Item = (&'a TableId, &'a Arc<[ProductValue]>)>,
49+
rows: impl Iterator<Item = (TableId, &'a Arc<[ProductValue]>)>,
5050
) -> HashMap<(TableId, IndexId), DeltaTableIndex> {
5151
let mut indexes: HashMap<(TableId, IndexId), DeltaTableIndex> = HashMap::new();
5252
for (table_id, rows) in rows {
53-
if let Some(schema) = tx.get_schema(*table_id) {
53+
if let Some(schema) = tx.get_schema(table_id) {
5454
// Fetch the column ids for each index
5555
let mut cols_for_index = vec![];
56-
for index_id in meta.index_ids_for_table(*table_id) {
56+
for index_id in meta.index_ids_for_table(table_id) {
5757
cols_for_index.push((index_id, schema.col_list_for_index_id(index_id)));
5858
}
5959
for (i, row) in rows.iter().enumerate() {
6060
for (index_id, col_list) in &cols_for_index {
6161
if !col_list.is_empty() {
6262
indexes
63-
.entry((*table_id, *index_id))
63+
.entry((table_id, *index_id))
6464
.or_default()
6565
.entry(row.project(col_list).unwrap())
6666
.or_default()
@@ -161,38 +161,24 @@ impl Datastore for DeltaTx<'_> {
161161
impl DeltaStore for DeltaTx<'_> {
162162
fn num_inserts(&self, table_id: TableId) -> usize {
163163
self.data
164-
.and_then(|data| {
165-
data.inserts()
166-
.find(|(id, ..)| **id == table_id)
167-
.map(|(.., rows)| rows.len())
168-
})
164+
.and_then(|data| data.inserts_for_table(table_id).map(|rows| rows.len()))
169165
.unwrap_or_default()
170166
}
171167

172168
fn num_deletes(&self, table_id: TableId) -> usize {
173169
self.data
174-
.and_then(|data| {
175-
data.deletes()
176-
.find(|(id, ..)| **id == table_id)
177-
.map(|(.., rows)| rows.len())
178-
})
170+
.and_then(|data| data.deletes_for_table(table_id).map(|rows| rows.len()))
179171
.unwrap_or_default()
180172
}
181173

182174
fn inserts_for_table(&self, table_id: TableId) -> Option<std::slice::Iter<'_, ProductValue>> {
183-
self.data.and_then(|data| {
184-
data.inserts()
185-
.find(|(id, ..)| **id == table_id)
186-
.map(|(.., rows)| rows.iter())
187-
})
175+
self.data
176+
.and_then(|data| data.inserts_for_table(table_id).map(|rows| rows.iter()))
188177
}
189178

190179
fn deletes_for_table(&self, table_id: TableId) -> Option<std::slice::Iter<'_, ProductValue>> {
191-
self.data.and_then(|data| {
192-
data.deletes()
193-
.find(|(id, ..)| **id == table_id)
194-
.map(|(.., rows)| rows.iter())
195-
})
180+
self.data
181+
.and_then(|data| data.deletes_for_table(table_id).map(|rows| rows.iter()))
196182
}
197183

198184
fn index_scan_range_for_delta(

‎crates/datastore/src/locking_tx_datastore/committed_state.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1066,7 +1066,7 @@ impl CommittedState {
10661066
);
10671067

10681068
// Record any truncated tables in the `TxData`.
1069-
tx_data.add_truncates(truncates);
1069+
tx_data.set_truncates(truncates);
10701070

10711071
// Merge read sets from the `MutTxId` into the `CommittedState`.
10721072
// It's important that this happens after applying the changes to `tx_data`,

0 commit comments

Comments
 (0)