Add splay to subscription invalidations (#38941)

This PR adds splay to subscription invalidation. Code push invalidates all user queries because the modules change, which creates a spike in load. This should smooth the spike in load by adding a delay before sending the message to wake up the subscription. The delay is random, from a range proportional to how many subscribers there are. Also adds metrics and logs.

GitOrigin-RevId: 6d5b430a930d3d8a2acc61b6a93febe3bf8a7935

Author: emmaling27

crates/common/src/knobs.rs

@@ -1347,3 +1347,14 @@ pub static SUBSCRIPTIONS_WORKER_QUEUE_SIZE: LazyLock<usize> =
 /// search query fails because indexes are bootstrapping.
 pub static SEARCH_INDEXES_UNAVAILABLE_RETRY_DELAY: LazyLock<Duration> =
     LazyLock::new(|| Duration::from_secs(env_config("SEARCH_INDEXES_UNAVAILABLE_RETRY_DELAY", 3)));
+
+/// The maximum number of subscriptions that can be invalidated immediately. If
+/// there are more, they will be splayed out.
+pub static SUBSCRIPTION_INVALIDATION_DELAY_THRESHOLD: LazyLock<usize> =
+    LazyLock::new(|| env_config("SUBSCRIPTION_INVALIDATION_DELAY_THRESHOLD", 200));
+
+/// How much to splay subscription invalidations. More precisely, this is the
+/// number used to multiply by the number of subscriptions that need to be
+/// invalidated to determine the delay before invalidating them.
+pub static SUBSCRIPTION_INVALIDATION_DELAY_MULTIPLIER: LazyLock<u64> =
+    LazyLock::new(|| env_config("SUBSCRIPTION_INVALIDATION_DELAY_MULTIPLIER", 5));

crates/database/src/metrics.rs

@@ -1078,3 +1078,11 @@ register_convex_histogram!(
 pub fn log_list_snapshot_page_documents(num_docs: usize) {
     log_distribution(&LIST_SNAPSHOT_PAGE_DOCUMENTS, num_docs as f64);
 }
+
+register_convex_histogram!(
+    SUBSCRIPTION_INVALIDATION_UPDATES,
+    "Number of subscriptions invalidated when advancing the log",
+);
+pub fn log_subscriptions_invalidated(num: usize) {
+    log_distribution(&SUBSCRIPTION_INVALIDATION_UPDATES, num as f64);
+}

crates/database/src/subscription.rs

@@ -14,6 +14,7 @@ use std::{
         },
         Arc,
     },
+    time::Duration,
 };
 
 use ::metrics::Timer;
@@ -24,7 +25,11 @@ use common::{
         DocumentIndexKeys,
     },
     errors::report_error,
-    knobs::SUBSCRIPTIONS_WORKER_QUEUE_SIZE,
+    knobs::{
+        SUBSCRIPTIONS_WORKER_QUEUE_SIZE,
+        SUBSCRIPTION_INVALIDATION_DELAY_MULTIPLIER,
+        SUBSCRIPTION_INVALIDATION_DELAY_THRESHOLD,
+    },
     runtime::{
         block_in_place,
         Runtime,
@@ -58,7 +63,10 @@ use tokio::sync::{
 use value::ResolvedDocumentId;
 
 use crate::{
-    metrics,
+    metrics::{
+        self,
+        log_subscriptions_invalidated,
+    },
     reads::ReadSet,
     write_log::{
         LogOwner,
@@ -116,6 +124,20 @@ impl Drop for SubscriptionSender {
     }
 }
 
+impl SubscriptionSender {
+    fn drop_with_delay(self, delay: Option<Duration>) {
+        self.valid_ts.store(-1, Ordering::SeqCst);
+        if let Some(delay) = delay {
+            tokio::spawn(async move {
+                tokio::time::sleep(delay).await;
+                _ = self.valid_tx.send(SubscriptionState::Invalid);
+            });
+        } else {
+            _ = self.valid_tx.send(SubscriptionState::Invalid);
+        }
+    }
+}
+
 enum SubscriptionRequest {
     Subscribe {
         token: Token,
@@ -314,9 +336,27 @@ impl SubscriptionManager {
                 }
             }
             // Then, invalidate all the remaining subscriptions.
+            let num_subscriptions_invalidated = to_notify.len();
+            let should_splay_invalidations =
+                num_subscriptions_invalidated > *SUBSCRIPTION_INVALIDATION_DELAY_THRESHOLD;
+            if should_splay_invalidations {
+                tracing::info!(
+                    "Splaying subscription invalidations since there are {} subscriptions to \
+                     invalidate. The threshold is {}",
+                    num_subscriptions_invalidated,
+                    *SUBSCRIPTION_INVALIDATION_DELAY_THRESHOLD
+                );
+            }
             for subscriber_id in to_notify {
-                self._remove(subscriber_id);
+                let delay = should_splay_invalidations.then(|| {
+                    Duration::from_millis(rand::random_range(
+                        0..=num_subscriptions_invalidated as u64
+                            * *SUBSCRIPTION_INVALIDATION_DELAY_MULTIPLIER,
+                    ))
+                });
+                self._remove(subscriber_id, delay);
             }
+            log_subscriptions_invalidated(num_subscriptions_invalidated);
 
             assert!(self.processed_ts <= next_ts);
             self.processed_ts = next_ts;
@@ -369,13 +409,14 @@ impl SubscriptionManager {
         if self.get_subscriber(key).is_none() {
             return;
         }
-        self._remove(key.id);
+        self._remove(key.id, None);
     }
 
-    fn _remove(&mut self, id: SubscriberId) {
+    fn _remove(&mut self, id: SubscriberId, delay: Option<Duration>) {
         let entry = self.subscribers.remove(id);
         self.subscriptions.remove(id, &entry.reads);
         // dropping `entry.sender` will invalidate the subscription
+        entry.sender.drop_with_delay(delay);
     }
 }
 
@@ -779,7 +820,7 @@ mod tests {
                 .subscribe_for_testing(token.clone())
                 .unwrap();
             subscriptions.push(subscriber);
-            subscription_manager._remove(id);
+            subscription_manager._remove(id, None);
         }
 
         assert!(
@@ -817,7 +858,7 @@ mod tests {
         let (_subscription, id) = subscription_manager
             .subscribe_for_testing(token.clone())
             .unwrap();
-        subscription_manager._remove(id);
+        subscription_manager._remove(id, None);
 
         assert!(notify_subscribed_tokens(
             &mut id_generator,
@@ -891,7 +932,7 @@ mod tests {
                     );
                 }
                 for (_, id) in &subscriptions {
-                    subscription_manager._remove(*id);
+                    subscription_manager._remove(*id, None);
                 }
                 let notifications = notify_subscribed_tokens(
                     &mut id_generator,
@@ -916,7 +957,7 @@ mod tests {
                 for token in &tokens {
                     let (_subscription, id) = subscription_manager
                         .subscribe_for_testing(token.clone()).unwrap();
-                    subscription_manager._remove(id);
+                    subscription_manager._remove(id, None);
                 }
                 let notifications = notify_subscribed_tokens(
                     &mut id_generator,