test(1099): Multi-column index selection through query macro (#1001)

Cargo.toml

@@ -20,6 +20,7 @@ members = [
   "crates/testing",
   "crates/vm",
   "modules/benchmarks",
+  "modules/perf-test",
   "modules/rust-wasm-test",
   "modules/quickstart-chat",
   "modules/sdk-test",
@@ -63,6 +64,10 @@ rpath = false
 [profile.bench]
 debug = true
 
+[profile.test]
+opt-level = 1
+debug = true
+
 [profile.profiling]
 inherits = "release"
 debug = true

crates/table/src/pointer_map.rs

@@ -153,16 +153,19 @@ static_assert_size!(PointerMap, 80);
 // These are only used as sanity checks in the debug profile, and e.g., in tests.
 #[cfg(debug_assertions)]
 impl PointerMap {
+    #[allow(unused)]
     fn maintains_invariants(&self) -> bool {
         self.maintains_map_invariant() && self.maintains_colliders_invariant()
     }
 
+    #[allow(unused)]
     fn maintains_colliders_invariant(&self) -> bool {
         self.colliders.iter().enumerate().all(|(idx, slot)| {
             slot.len() >= 2 || slot.is_empty() && self.emptied_collider_slots.contains(&ColliderSlotIndex::new(idx))
         })
     }
 
+    #[allow(unused)]
     fn maintains_map_invariant(&self) -> bool {
         self.map.values().all(|poc| {
             let collider = poc.as_collider();
@@ -234,8 +237,6 @@ impl PointerMap {
     ///
     /// Handles any hash conflicts for `hash`.
     pub fn insert(&mut self, hash: RowHash, ptr: RowPointer) -> bool {
-        debug_assert!(self.maintains_invariants());
-
         let mut was_in_map = false;
 
         self.map
@@ -291,17 +292,13 @@ impl PointerMap {
             // 0 hashes so far.
             .or_insert(PtrOrCollider::ptr(ptr));
 
-        debug_assert!(self.maintains_invariants());
-
         was_in_map
     }
 
     /// Removes the association `hash -> ptr`.
     ///
     /// Returns whether the association was deleted.
     pub fn remove(&mut self, hash: RowHash, ptr: RowPointer) -> bool {
-        debug_assert!(self.maintains_invariants());
-
         let ret = 'fun: {
             let Entry::Occupied(mut entry) = self.map.entry(hash) else {
                 break 'fun false;
@@ -335,8 +332,6 @@ impl PointerMap {
             true
         };
 
-        debug_assert!(self.maintains_invariants());
-
         ret
     }
 }

crates/testing/src/modules.rs

@@ -197,6 +197,12 @@ pub static DEFAULT_CONFIG: Config = Config {
     fsync: FsyncPolicy::Never,
 };
 
+/// For performance tests, do not persist to disk.
+pub static IN_MEMORY_CONFIG: Config = Config {
+    storage: Storage::Disk,
+    fsync: FsyncPolicy::Never,
+};
+
 /// Used to parse output from module logs.
 ///
 /// Sync with: `core::database_logger::Record`. We can't use it

crates/testing/tests/standalone_integration_test.rs

@@ -1,6 +1,6 @@
 use serial_test::serial;
 use spacetimedb_testing::modules::{
-    CompilationMode, CompiledModule, LogLevel, LoggerRecord, ModuleHandle, DEFAULT_CONFIG,
+    CompilationMode, CompiledModule, LogLevel, LoggerRecord, ModuleHandle, DEFAULT_CONFIG, IN_MEMORY_CONFIG,
 };
 
 fn init() {
@@ -139,3 +139,44 @@ fn test_call_query_macro() {
         },
     );
 }
+
+#[test]
+#[serial]
+/// This test runs the index scan workloads in the `perf-test` module.
+/// Timing spans should be < 1ms if the correct index was used.
+/// Otherwise these workloads will degenerate into full table scans.
+fn test_index_scans() {
+    init();
+    CompiledModule::compile("perf-test", CompilationMode::Release).with_module_async(
+        IN_MEMORY_CONFIG,
+        |module| async move {
+            let json = r#"{"call": {"fn": "load_location_table", "args": []}}"#;
+            module.send(json.to_string()).await.unwrap();
+
+            let json = r#"{"call": {"fn": "test_index_scan_on_id", "args": []}}"#;
+            module.send(json.to_string()).await.unwrap();
+
+            let json = r#"{"call": {"fn": "test_index_scan_on_chunk", "args": []}}"#;
+            module.send(json.to_string()).await.unwrap();
+
+            let json = r#"{"call": {"fn": "test_index_scan_on_x_z_dimension", "args": []}}"#;
+            module.send(json.to_string()).await.unwrap();
+
+            // TODO(1011): Uncomment once multi-column prefix scans are supported
+            // let json = r#"{"call": {"fn": "test_index_scan_on_x_z", "args": []}}"#;
+            // module.send(json.to_string()).await.unwrap();
+
+            let logs = read_logs(&module).await;
+
+            // Each timing span should be < 1ms
+            let timing = |line: &str| {
+                line.starts_with("Timing span")
+                    && (line.ends_with("ns") || line.ends_with("us") || line.ends_with("µs"))
+            };
+            assert!(timing(&logs[0]));
+            assert!(timing(&logs[1]));
+            assert!(timing(&logs[2]));
+            // assert!(timing(&logs[3]));
+        },
+    );
+}

modules/perf-test/.gitignore

@@ -0,0 +1,17 @@
+# Generated by Cargo
+# will have compiled files and executables
+debug/
+target/
+
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+# MSVC Windows builds of rustc generate these, which store debugging information
+*.pdb
+
+# Spacetime ignore
+/.spacetime

modules/perf-test/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "perf-test-module"
+version = "0.1.0"
+edition.workspace = true
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+spacetimedb = { path = "../../crates/bindings" }
+log.workspace = true

modules/perf-test/README.md

@@ -0,0 +1,4 @@
+# `perf-test`
+
+A module with various performance workloads for SpacetimeDB,
+called from the `testing` crate.

modules/perf-test/src/lib.rs

@@ -0,0 +1,80 @@
+use spacetimedb::{query, spacetimedb, time_span::Span};
+
+#[spacetimedb(table)]
+#[spacetimedb(index(btree, name = "x", x))]
+#[spacetimedb(index(btree, name = "chunk", chunk))]
+#[spacetimedb(index(btree, name = "coordinates", x, z, dimension))]
+#[derive(Debug, PartialEq, Eq)]
+pub struct Location {
+    #[primarykey]
+    pub id: u64,
+    pub chunk: u64,
+    pub x: i32,
+    pub z: i32,
+    pub dimension: u32,
+}
+
+// 1000 chunks, 1200 rows per chunk = 1.2M rows
+const NUM_CHUNKS: u64 = 1000;
+const ROWS_PER_CHUNK: u64 = 1200;
+
+#[spacetimedb(reducer)]
+pub fn load_location_table() {
+    for chunk in 0u64..NUM_CHUNKS {
+        for i in 0u64..ROWS_PER_CHUNK {
+            let id = chunk * 1200 + i;
+            let x = 0i32;
+            let z = chunk as i32;
+            let dimension = id as u32;
+            let _ = Location::insert(Location {
+                id,
+                chunk,
+                x,
+                z,
+                dimension,
+            });
+        }
+    }
+}
+
+const ID: u64 = 989_987;
+const CHUNK: u64 = ID / ROWS_PER_CHUNK;
+
+#[spacetimedb(reducer)]
+/// Probing a single column index for a single row should be fast!
+pub fn test_index_scan_on_id() {
+    let span = Span::start("Index scan on {id}");
+    let location = Location::filter_by_id(&ID).unwrap();
+    span.end();
+    assert_eq!(ID, location.id);
+}
+
+#[spacetimedb(reducer)]
+/// Scanning a single column index for `ROWS_PER_CHUNK` rows should also be fast!
+pub fn test_index_scan_on_chunk() {
+    let span = Span::start("Index scan on {chunk}");
+    let n = Location::filter_by_chunk(&CHUNK).count();
+    span.end();
+    assert_eq!(n as u64, ROWS_PER_CHUNK);
+}
+
+#[spacetimedb(reducer)]
+/// Probing a multi-column index for a single row should be fast!
+pub fn test_index_scan_on_x_z_dimension() {
+    let z = CHUNK as i32;
+    let dimension = ID as u32;
+    let span = Span::start("Index scan on {x, z, dimension}");
+    let n = query!(|r: Location| r.x == 0 && r.z == z && r.dimension == dimension).count();
+    span.end();
+    assert_eq!(n, 1);
+}
+
+#[spacetimedb(reducer)]
+/// Probing a multi-column index for `ROWS_PER_CHUNK` rows should also be fast!
+pub fn test_index_scan_on_x_z() {
+    let z = CHUNK as i32;
+    let span = Span::start("Index scan on {x, z}");
+    let n = query!(|r: Location| r.x == 0 && r.z == z).count();
+    span.end();
+    assert_eq!(n as u64, ROWS_PER_CHUNK);
+}