Prune queries from subscription evaluation (#2641)

Author: joshua-spacetime

crates/core/src/host/wasm_common/module_host_actor.rs

@@ -583,7 +583,7 @@ impl<T: WasmInstance> WasmModuleInstance<T> {
             request_id,
             timer,
         };
-        let event = match self
+        let (event, _) = match self
             .info
             .subscriptions
             .commit_and_broadcast_event(client.as_deref(), event, tx)

crates/core/src/subscription/execution_unit.rs

@@ -37,7 +37,7 @@ use std::time::Duration;
 /// as is the case for incremental joins.
 /// And we want to associate a hash with the entire unit of execution,
 /// rather than an individual plan.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct QueryHash {
     data: [u8; 32],
 }
@@ -49,8 +49,15 @@ impl From<QueryHash> for u256 {
 }
 
 impl QueryHash {
+    /// The zero value of a QueryHash
     pub const NONE: Self = Self { data: [0; 32] };
 
+    /// The min value of a QueryHash
+    pub const MIN: Self = Self::NONE;
+
+    /// The max value of a QueryHash
+    pub const MAX: Self = Self { data: [0xFFu8; 32] };
+
     pub fn from_bytes(bytes: &[u8]) -> Self {
         Self {
             data: blake3::hash(bytes).into(),

crates/core/src/subscription/module_subscription_actor.rs

@@ -712,7 +712,7 @@ impl ModuleSubscriptions {
         caller: Option<&ClientConnectionSender>,
         mut event: ModuleEvent,
         tx: MutTx,
-    ) -> Result<Result<Arc<ModuleEvent>, WriteConflict>, DBError> {
+    ) -> Result<Result<(Arc<ModuleEvent>, ExecutionMetrics), WriteConflict>, DBError> {
         // Take a read lock on `subscriptions` before committing tx
         // else it can result in subscriber receiving duplicate updates.
         let subscriptions = self.subscriptions.read();
@@ -742,10 +742,16 @@ impl ModuleSubscriptions {
             .unwrap_or_else(|| DeltaTx::from(&*read_tx));
 
         let event = Arc::new(event);
+        let mut metrics = ExecutionMetrics::default();
 
         match &event.status {
             EventStatus::Committed(_) => {
-                subscriptions.eval_updates(&read_tx, event.clone(), caller, &self.relational_db.database_identity())
+                metrics.merge(subscriptions.eval_updates(
+                    &read_tx,
+                    event.clone(),
+                    caller,
+                    &self.relational_db.database_identity(),
+                ));
             }
             EventStatus::Failed(_) => {
                 if let Some(client) = caller {
@@ -761,7 +767,7 @@ impl ModuleSubscriptions {
             EventStatus::OutOfEnergy => {} // ?
         }
 
-        Ok(Ok(event))
+        Ok(Ok((event, metrics)))
     }
 }
 
@@ -798,6 +804,7 @@ mod tests {
     use spacetimedb_lib::bsatn::ToBsatn;
     use spacetimedb_lib::db::auth::StAccess;
     use spacetimedb_lib::identity::AuthCtx;
+    use spacetimedb_lib::metrics::ExecutionMetrics;
     use spacetimedb_lib::{bsatn, ConnectionId, ProductType, ProductValue, Timestamp};
     use spacetimedb_lib::{error::ResultTest, AlgebraicType, Identity};
     use spacetimedb_primitives::TableId;
@@ -1079,19 +1086,19 @@ mod tests {
         subs: &ModuleSubscriptions,
         deletes: impl IntoIterator<Item = (TableId, ProductValue)>,
         inserts: impl IntoIterator<Item = (TableId, ProductValue)>,
-    ) -> anyhow::Result<()> {
+    ) -> anyhow::Result<ExecutionMetrics> {
         let mut tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::ForTests);
         for (table_id, row) in deletes {
             tx.delete_product_value(table_id, &row)?;
         }
         for (table_id, row) in inserts {
             db.insert(&mut tx, table_id, &bsatn::to_vec(&row)?)?;
         }
-        assert!(matches!(
-            subs.commit_and_broadcast_event(None, module_event(), tx),
-            Ok(Ok(_))
-        ));
-        Ok(())
+
+        let Ok(Ok((_, metrics))) = subs.commit_and_broadcast_event(None, module_event(), tx) else {
+            panic!("Encountered an error in `commit_and_broadcast_event`");
+        };
+        Ok(metrics)
     }
 
     #[test]
@@ -1689,6 +1696,119 @@ mod tests {
         Ok(())
     }
 
+    /// Test that we do not evaluate queries that we know will not match table update rows
+    #[tokio::test]
+    async fn test_query_pruning() -> anyhow::Result<()> {
+        // Establish a connection for each client
+        let (tx_for_a, mut rx_for_a) = client_connection(client_id_from_u8(1));
+        let (tx_for_b, mut rx_for_b) = client_connection(client_id_from_u8(2));
+
+        let db = relational_db()?;
+        let subs = module_subscriptions(db.clone());
+
+        let u_id = db.create_table_for_test(
+            "u",
+            &[
+                ("i", AlgebraicType::U64),
+                ("a", AlgebraicType::U64),
+                ("b", AlgebraicType::U64),
+            ],
+            &[0.into()],
+        )?;
+        let v_id = db.create_table_for_test(
+            "v",
+            &[
+                ("i", AlgebraicType::U64),
+                ("x", AlgebraicType::U64),
+                ("y", AlgebraicType::U64),
+            ],
+            &[0.into(), 1.into()],
+        )?;
+
+        commit_tx(
+            &db,
+            &subs,
+            [],
+            [
+                (u_id, product![0u64, 1u64, 1u64]),
+                (u_id, product![1u64, 2u64, 2u64]),
+                (u_id, product![2u64, 3u64, 3u64]),
+                (v_id, product![0u64, 4u64, 4u64]),
+                (v_id, product![1u64, 5u64, 5u64]),
+            ],
+        )?;
+
+        let mut query_ids = 0;
+
+        // Returns (i: 0, a: 1, b: 1)
+        subscribe_multi(
+            &subs,
+            &[
+                "select u.* from u join v on u.i = v.i where v.x = 4",
+                "select u.* from u join v on u.i = v.i where v.x = 6",
+            ],
+            tx_for_a,
+            &mut query_ids,
+        )?;
+
+        // Returns (i: 1, a: 2, b: 2)
+        subscribe_multi(
+            &subs,
+            &[
+                "select u.* from u join v on u.i = v.i where v.x = 5",
+                "select u.* from u join v on u.i = v.i where v.x = 7",
+            ],
+            tx_for_b,
+            &mut query_ids,
+        )?;
+
+        // Wait for both subscriptions
+        assert!(matches!(
+            rx_for_a.recv().await,
+            Some(SerializableMessage::Subscription(SubscriptionMessage {
+                result: SubscriptionResult::SubscribeMulti(_),
+                ..
+            }))
+        ));
+        assert!(matches!(
+            rx_for_b.recv().await,
+            Some(SerializableMessage::Subscription(SubscriptionMessage {
+                result: SubscriptionResult::SubscribeMulti(_),
+                ..
+            }))
+        ));
+
+        // Modify a single row in `v`
+        let metrics = commit_tx(
+            &db,
+            &subs,
+            [(v_id, product![1u64, 5u64, 5u64])],
+            [(v_id, product![1u64, 5u64, 6u64])],
+        )?;
+
+        // We should only have evaluated a single query
+        assert_eq!(metrics.delta_queries_evaluated, 1);
+        assert_eq!(metrics.delta_queries_matched, 1);
+
+        // Insert a new row into `v`
+        let metrics = commit_tx(&db, &subs, [], [(v_id, product![2u64, 6u64, 6u64])])?;
+
+        assert_tx_update_for_table(
+            &mut rx_for_a,
+            u_id,
+            &ProductType::from([AlgebraicType::U64, AlgebraicType::U64, AlgebraicType::U64]),
+            [product![2u64, 3u64, 3u64]],
+            [],
+        )
+        .await;
+
+        // We should only have evaluated a single query
+        assert_eq!(metrics.delta_queries_evaluated, 1);
+        assert_eq!(metrics.delta_queries_matched, 1);
+
+        Ok(())
+    }
+
     /// Asserts that a subscription holds a tx handle for the entire length of its evaluation.
     #[test]
     fn test_tx_subscription_ordering() -> ResultTest<()> {