Fix subscription manager metadata for joins (#2705)

Author: joshua-spacetime

crates/core/src/subscription/module_subscription_actor.rs

@@ -1802,6 +1802,34 @@ mod tests {
         assert_eq!(metrics.delta_queries_evaluated, 1);
         assert_eq!(metrics.delta_queries_matched, 1);
 
+        // Modify a matching row in `u`
+        let metrics = commit_tx(
+            &db,
+            &subs,
+            [(u_id, product![1u64, 2u64, 2u64])],
+            [(u_id, product![1u64, 2u64, 3u64])],
+        )?;
+
+        assert_tx_update_for_table(
+            &mut rx_for_b,
+            u_id,
+            &ProductType::from([AlgebraicType::U64, AlgebraicType::U64, AlgebraicType::U64]),
+            [product![1u64, 2u64, 3u64]],
+            [product![1u64, 2u64, 2u64]],
+        )
+        .await;
+
+        // We should have evaluated all of the queries
+        assert_eq!(metrics.delta_queries_evaluated, 4);
+        assert_eq!(metrics.delta_queries_matched, 1);
+
+        // Insert a non-matching row in `u`
+        let metrics = commit_tx(&db, &subs, [], [(u_id, product![3u64, 0u64, 0u64])])?;
+
+        // We should have evaluated all of the queries
+        assert_eq!(metrics.delta_queries_evaluated, 4);
+        assert_eq!(metrics.delta_queries_matched, 0);
+
         Ok(())
     }
 

crates/core/src/subscription/module_subscription_manager.rs

@@ -309,11 +309,13 @@ pub struct SubscriptionManager {
     // Queries for which there is at least one subscriber.
     queries: HashMap<QueryHash, QueryState>,
 
-    // Inverted index from tables to queries that read from them.
-    // Note, a query is present in either `tables` or `search_args` but not both.
+    // If a query reads from a table,
+    // but does not have a simple equality filter on that table,
+    // we map the table to the query in this inverted index.
     tables: IntMap<TableId, HashSet<QueryHash>>,
 
-    // For queries that have simple equality filters,
+    // If a query reads from a table,
+    // and has a simple equality filter on that table,
     // we map the filter values to the query in this lookup table.
     search_args: SearchArguments,
 }
@@ -619,7 +621,7 @@ impl SubscriptionManager {
     }
 
     // Update the mapping from table id to related queries by inserting the given query.
-    // If this query has search arguments, that mapping is updated instead.
+    // Also add any search arguments the query may have.
     // This takes a ref to the table map instead of `self` to avoid borrowing issues.
     fn insert_query(
         tables: &mut IntMap<TableId, HashSet<QueryHash>>,
@@ -629,15 +631,15 @@ impl SubscriptionManager {
         // If this is new, we need to update the table to query mapping.
         if !query_state.has_subscribers() {
             let hash = query_state.query.hash();
-            let mut has_search_args = false;
+            let mut table_ids = query_state.query.table_ids().collect::<HashSet<_>>();
             for (table_id, col_id, arg) in query_state.search_args() {
-                has_search_args = true;
+                table_ids.remove(&table_id);
                 search_args.insert_query(table_id, col_id, arg, hash);
             }
-            if !has_search_args {
-                for table_id in query_state.query.table_ids() {
-                    tables.entry(table_id).or_default().insert(hash);
-                }
+            // Update the `tables` map if this query reads from a table,
+            // but does not have a search argument for that table.
+            for table_id in table_ids {
+                tables.entry(table_id).or_default().insert(hash);
             }
         }
     }
@@ -714,9 +716,10 @@ impl SubscriptionManager {
         {
             queries.insert(hash);
         }
-        queries
-            .into_iter()
-            .chain(self.tables.get(&table_update.table_id).into_iter().flatten())
+        for hash in self.tables.get(&table_update.table_id).into_iter().flatten() {
+            queries.insert(hash);
+        }
+        queries.into_iter()
     }
 
     /// This method takes a set of delta tables,
@@ -1519,7 +1522,7 @@ mod tests {
     }
 
     #[test]
-    fn test_lookup_queries_for_search_arg() -> ResultTest<()> {
+    fn test_search_args_for_selects() -> ResultTest<()> {
         let db = TestDB::durable()?;
 
         let table_id = create_table(&db, "t")?;
@@ -1581,6 +1584,75 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_search_args_for_join() -> ResultTest<()> {
+        let db = TestDB::durable()?;
+
+        let schema = [("id", AlgebraicType::U8), ("a", AlgebraicType::U8)];
+
+        let t_id = db.create_table_for_test("t", &schema, &[0.into()])?;
+        let s_id = db.create_table_for_test("s", &schema, &[0.into()])?;
+
+        let client = Arc::new(client(0));
+        let mut subscriptions = SubscriptionManager::default();
+
+        let plan = compile_plan(&db, "select t.* from t join s on t.id = s.id where s.a = 1")?;
+        let hash = plan.hash;
+
+        subscriptions.add_subscription_multi(client.clone(), vec![plan], QueryId::new(0))?;
+
+        // Do we need to evaluate the above join query for this table update?
+        // Yes, because the above query does not filter on `t`.
+        // Therefore we must evaluate it for any update on `t`.
+        let table_update = DatabaseTableUpdate {
+            table_id: t_id,
+            table_name: "t".into(),
+            inserts: [product![0u8, 0u8]].into(),
+            deletes: [].into(),
+        };
+
+        let hashes = subscriptions
+            .queries_for_table_update(&table_update)
+            .cloned()
+            .collect::<Vec<_>>();
+
+        assert_eq!(hashes, vec![hash]);
+
+        // Do we need to evaluate the above join query for this table update?
+        // Yes, because `s.a = 1`.
+        let table_update = DatabaseTableUpdate {
+            table_id: s_id,
+            table_name: "s".into(),
+            inserts: [product![0u8, 1u8]].into(),
+            deletes: [].into(),
+        };
+
+        let hashes = subscriptions
+            .queries_for_table_update(&table_update)
+            .cloned()
+            .collect::<Vec<_>>();
+
+        assert_eq!(hashes, vec![hash]);
+
+        // Do we need to evaluate the above join query for this table update?
+        // No, because `s.a != 1`.
+        let table_update = DatabaseTableUpdate {
+            table_id: s_id,
+            table_name: "s".into(),
+            inserts: [product![0u8, 2u8]].into(),
+            deletes: [].into(),
+        };
+
+        let hashes = subscriptions
+            .queries_for_table_update(&table_update)
+            .cloned()
+            .collect::<Vec<_>>();
+
+        assert!(hashes.is_empty());
+
+        Ok(())
+    }
+
     #[test]
     fn test_subscribe_fails_with_duplicate_request_id() -> ResultTest<()> {
         let db = TestDB::durable()?;