Add a metric for subscription queue lag (#39344)

GitOrigin-RevId: 1736b8fd2cee971dbb8dcb9db68a14b24d5a2353

Author: goffrie

crates/database/src/metrics.rs

@@ -15,6 +15,7 @@ use metrics::{
     log_distribution,
     log_distribution_with_labels,
     register_convex_counter,
+    register_convex_gauge,
     register_convex_histogram,
     CancelableTimer,
     IntoLabel,
@@ -1182,3 +1183,19 @@ register_convex_counter!(
 pub fn log_index_too_large_blocking_writes(index_type: SearchType) {
     log_counter_with_labels(&INDEX_TOO_LARGE_BLOCKING_WRITES, 1, vec![index_type.tag()]);
 }
+
+register_convex_histogram!(
+    SUBSCRIPTION_QUEUE_LAG_SECONDS,
+    "How long subscription requests wait in the subscription worker queue",
+);
+pub fn log_subscription_queue_lag(seconds: f64) {
+    log_distribution(&SUBSCRIPTION_QUEUE_LAG_SECONDS, seconds);
+}
+
+register_convex_gauge!(
+    SUBSCRIPTION_QUEUE_LENGTH_INFO,
+    "The number of items in subscription queues",
+);
+pub fn log_subscription_queue_length_delta(delta: i64) {
+    SUBSCRIPTION_QUEUE_LENGTH_INFO.add(delta as f64);
+}

crates/database/src/subscription.rs

@@ -105,6 +105,7 @@ impl SubscriptionsClient {
             TrySendError::Full(..) => metrics::subscriptions_worker_full_error().into(),
             TrySendError::Closed(..) => metrics::shutdown_error(),
         })?;
+        metrics::log_subscription_queue_length_delta(1);
         Ok(subscription)
     }
 
@@ -160,6 +161,7 @@ pub enum SubscriptionsWorker {}
 impl SubscriptionsWorker {
     pub(crate) fn start<RT: Runtime>(log: LogOwner, runtime: RT) -> SubscriptionsClient {
         let (tx, rx) = mpsc::channel(*SUBSCRIPTIONS_WORKER_QUEUE_SIZE);
+        let rx = CountingReceiver(rx);
 
         let log_reader = log.reader();
         let mut manager = SubscriptionManager::new(log);
@@ -174,8 +176,25 @@ impl SubscriptionsWorker {
     }
 }
 
+struct CountingReceiver(mpsc::Receiver<SubscriptionRequest>);
+impl Drop for CountingReceiver {
+    fn drop(&mut self) {
+        self.0.close();
+        metrics::log_subscription_queue_length_delta(-(self.0.len() as i64));
+    }
+}
+impl CountingReceiver {
+    async fn recv(&mut self) -> Option<SubscriptionRequest> {
+        let r = self.0.recv().await;
+        if r.is_some() {
+            metrics::log_subscription_queue_length_delta(-1);
+        }
+        r
+    }
+}
+
 impl SubscriptionManager {
-    async fn run_worker(&mut self, mut rx: mpsc::Receiver<SubscriptionRequest>) {
+    async fn run_worker(&mut self, mut rx: CountingReceiver) {
         tracing::info!("Starting subscriptions worker");
         loop {
             futures::select_biased! {
@@ -259,6 +278,7 @@ impl SubscriptionManager {
         mut token: Token,
         sender: SubscriptionSender,
     ) -> anyhow::Result<SubscriberId> {
+        metrics::log_subscription_queue_lag(self.log.max_ts().secs_since_f64(token.ts()));
         // The client may not have fully refreshed their token past our
         // processed timestamp, so finish the job for them if needed.
         //
@@ -713,7 +733,10 @@ mod tests {
     };
 
     use crate::{
-        subscription::SubscriptionManager,
+        subscription::{
+            CountingReceiver,
+            SubscriptionManager,
+        },
         ReadSet,
         Token,
     };
@@ -1133,8 +1156,8 @@ mod tests {
         to_notify
     }
 
-    fn disconnected_rx<T>() -> mpsc::Receiver<T> {
-        mpsc::channel(1).1
+    fn disconnected_rx() -> CountingReceiver {
+        CountingReceiver(mpsc::channel(1).1)
     }
 
     #[tokio::test]

crates/database/src/write_log.rs

@@ -383,7 +383,7 @@ impl LogOwner {
 
     pub fn max_ts(&self) -> Timestamp {
         let snapshot = { self.inner.lock().log.clone() };
-        block_in_place(|| snapshot.max_ts())
+        snapshot.max_ts()
     }
 
     pub fn refresh_token(