Fix query overwrites in the subscription manager

Author: joshua-spacetime

crates/core/src/subscription/module_subscription_actor.rs

@@ -2331,48 +2331,147 @@ mod tests {
         Ok(())
     }
 
-    /// Test that one client unsubscribing does not affect another
+    /// Test that one client subscribing does not affect another
     #[tokio::test]
-    async fn test_unsubscribe() -> anyhow::Result<()> {
+    async fn test_subscribe_distinct_queries_same_plan() -> 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 = ModuleSubscriptions::for_test_enclosing_runtime(db.clone());
 
-        let u_id = db.create_table_for_test(
+        let u_id = db.create_table_for_test_with_the_works(
             "u",
             &[
                 ("i", AlgebraicType::U64),
                 ("a", AlgebraicType::U64),
                 ("b", AlgebraicType::U64),
             ],
             &[0.into()],
+            // The join column for this table does not have to be unique,
+            // because pruning only requires us to probe the join index on `v`.
+            &[],
+            StAccess::Public,
         )?;
-        let v_id = db.create_table_for_test(
+        let v_id = db.create_table_for_test_with_the_works(
             "v",
             &[
                 ("i", AlgebraicType::U64),
                 ("x", AlgebraicType::U64),
                 ("y", AlgebraicType::U64),
             ],
             &[0.into(), 1.into()],
+            &[0.into()],
+            StAccess::Public,
         )?;
 
         commit_tx(&db, &subs, [], [(v_id, product![1u64, 1u64, 1u64])])?;
 
         let mut query_ids = 0;
 
+        // Both clients subscribe to the same query modulo whitespace
         subscribe_multi(
             &subs,
             &["select u.* from u join v on u.i = v.i where v.x = 1"],
             tx_for_a,
             &mut query_ids,
         )?;
+        subscribe_multi(
+            &subs,
+            &["select u.* from u join v on u.i = v.i where v.x =  1"],
+            tx_for_b.clone(),
+            &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(_),
+                ..
+            }))
+        );
+
+        // Insert a new row into `u`
+        commit_tx(&db, &subs, [], [(u_id, product![1u64, 0u64, 0u64])])?;
+
+        assert_tx_update_for_table(
+            &mut rx_for_a,
+            u_id,
+            &ProductType::from([AlgebraicType::U64, AlgebraicType::U64, AlgebraicType::U64]),
+            [product![1u64, 0u64, 0u64]],
+            [],
+        )
+        .await;
+
+        assert_tx_update_for_table(
+            &mut rx_for_b,
+            u_id,
+            &ProductType::from([AlgebraicType::U64, AlgebraicType::U64, AlgebraicType::U64]),
+            [product![1u64, 0u64, 0u64]],
+            [],
+        )
+        .await;
+
+        Ok(())
+    }
+
+    /// Test that one client unsubscribing does not affect another
+    #[tokio::test]
+    async fn test_unsubscribe_distinct_queries_same_plan() -> 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 = ModuleSubscriptions::for_test_enclosing_runtime(db.clone());
+
+        let u_id = db.create_table_for_test_with_the_works(
+            "u",
+            &[
+                ("i", AlgebraicType::U64),
+                ("a", AlgebraicType::U64),
+                ("b", AlgebraicType::U64),
+            ],
+            &[0.into()],
+            // The join column for this table does not have to be unique,
+            // because pruning only requires us to probe the join index on `v`.
+            &[],
+            StAccess::Public,
+        )?;
+        let v_id = db.create_table_for_test_with_the_works(
+            "v",
+            &[
+                ("i", AlgebraicType::U64),
+                ("x", AlgebraicType::U64),
+                ("y", AlgebraicType::U64),
+            ],
+            &[0.into(), 1.into()],
+            &[0.into()],
+            StAccess::Public,
+        )?;
+
+        commit_tx(&db, &subs, [], [(v_id, product![1u64, 1u64, 1u64])])?;
+
+        let mut query_ids = 0;
+
         subscribe_multi(
             &subs,
             &["select u.* from u join v on u.i = v.i where v.x = 1"],
+            tx_for_a,
+            &mut query_ids,
+        )?;
+        subscribe_multi(
+            &subs,
+            &["select u.* from u join v on u.i = v.i where  v.x = 1"],
             tx_for_b.clone(),
             &mut query_ids,
         )?;

crates/core/src/subscription/module_subscription_manager.rs

@@ -399,7 +399,7 @@ impl QueriedTableIndexIds {
 /// See [`JoinEdge`] for more details.
 #[derive(Debug, Default)]
 pub struct JoinEdges {
-    edges: BTreeMap<JoinEdge, HashMap<AlgebraicValue, QueryHash>>,
+    edges: BTreeMap<JoinEdge, HashMap<AlgebraicValue, HashSet<QueryHash>>>,
 }
 
 impl JoinEdges {
@@ -408,18 +408,28 @@ impl JoinEdges {
         let mut inserted = false;
         for (edge, rhs_val) in qs.query.join_edges() {
             inserted = true;
-            self.edges.entry(edge).or_default().insert(rhs_val, qs.query.hash);
+            self.edges
+                .entry(edge)
+                .or_default()
+                .entry(rhs_val)
+                .or_default()
+                .insert(qs.query.hash);
         }
         inserted
     }
 
     /// If this query has any join edges, remove them from the map.
     fn remove_query(&mut self, query: &Query) {
         for (edge, rhs_val) in query.join_edges() {
-            if let Some(hashes) = self.edges.get_mut(&edge) {
-                hashes.remove(&rhs_val);
-                if hashes.is_empty() {
-                    self.edges.remove(&edge);
+            if let Some(values) = self.edges.get_mut(&edge) {
+                if let Some(hashes) = values.get_mut(&rhs_val) {
+                    hashes.remove(&query.hash);
+                    if hashes.is_empty() {
+                        values.remove(&rhs_val);
+                        if values.is_empty() {
+                            self.edges.remove(&edge);
+                        }
+                    }
                 }
             }
         }
@@ -436,6 +446,7 @@ impl JoinEdges {
         self.edges
             .range(JoinEdge::range_for_table(table_id))
             .filter_map(move |(edge, hashes)| find_rhs_val(edge, row).as_ref().and_then(|rhs_val| hashes.get(rhs_val)))
+            .flatten()
     }
 }