convex-rs PR 10: Fix for a panic that could occur in a subscribe call (#41227)

If a query was subscribed to multiple times by the same client, an unsubscribe followed by a subsequent subscribe (if in the right/wrong order) would trigger bad internal state and a panic.

This was due to a using a value that was incremented and decremented as part of the SubscriberId. That could lead to a SubscriberId for a new subscription matching a previous one, which is invalid state.

Co-authored-by: Christian Wyglendowski <[email protected]>
GitOrigin-RevId: 04150d72ffba18410d332086d32d18ce33a7c67b

Author: 2 people

crates/convex/src/base_client/mod.rs

@@ -63,6 +63,11 @@ struct LocalQuery {
     canonicalized_udf_path: CanonicalizedUdfPath,
     args: BTreeMap<String, Value>,
     num_subscribers: usize, // TODO: remove
+    /// A unique index value for each subscription to this query.
+    ///
+    /// Must be incremented each time a new subscription is added, and never
+    /// decremented.
+    subscription_index: usize,
 }
 
 #[derive(Clone, Debug)]
@@ -115,9 +120,11 @@ impl LocalSyncState {
         let query_token = serialize_path_and_args(udf_path.clone(), args.clone());
 
         if let Some(existing_entry) = self.query_set.get_mut(&query_token) {
+            // This is a new subscription to an existing query.
             existing_entry.num_subscribers += 1;
+            existing_entry.subscription_index += 1;
             let query_id = existing_entry.id;
-            let subscription = SubscriberId(query_id, existing_entry.num_subscribers - 1);
+            let subscription = SubscriberId(query_id, existing_entry.subscription_index);
             let prev = self.latest_results.subscribers.insert(subscription);
             assert!(prev.is_none(), "INTERNAL BUG: Subscriber ID already taken.");
             return (None, subscription);
@@ -147,6 +154,7 @@ impl LocalSyncState {
             canonicalized_udf_path,
             args,
             num_subscribers: 1,
+            subscription_index: 0,
         };
 
         self.query_set.insert(query_token.clone(), query);

crates/convex/src/client/mod.rs

@@ -484,6 +484,7 @@ pub mod tests {
             SyncProtocol,
         },
         value::Value,
+        QuerySubscription,
     };
 
     impl ConvexClient {
@@ -769,6 +770,40 @@ pub mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_client_subscribe_unsubscribe_subscribe() -> anyhow::Result<()> {
+        let (mut client, mut test_protocol) = ConvexClient::with_test_protocol().await?;
+        let subscription1b: QuerySubscription;
+        {
+            // This subscription goes out of scope and unsubscribes at the end of this
+            // block. The internal num_subscribers value gets decremented.
+            let _ignored = client.subscribe("getValue1", btreemap! {}).await?;
+            subscription1b = client.subscribe("getValue1", btreemap! {}).await?;
+        }
+        // In the buggy scenario, this subscription gets an ID via num_subscribers ID
+        // that matches subscription1b. That triggers a panic.
+        let subscription1c = client.subscribe("getValue1", btreemap! {}).await?;
+        test_protocol.take_sent().await;
+        let mut watch = client.watch_all();
+
+        test_protocol
+            .fake_server_response(
+                fake_transition(StateVersion::initial(), vec![(QueryId::new(0), 10.into())]).0,
+            )
+            .await?;
+
+        let results = watch.next().await.expect("Watch should have results");
+        assert_eq!(
+            results.get(&subscription1b),
+            Some(&FunctionResult::Value(10.into()))
+        );
+        assert_eq!(
+            results.get(&subscription1c),
+            Some(&FunctionResult::Value(10.into()))
+        );
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_client_consistent_view_watch() -> anyhow::Result<()> {
         let (mut client, mut test_protocol) = ConvexClient::with_test_protocol().await?;