perf(pubsub): concurrent acks/nacks (#5140)

Most of the work for #3975 

Send ack/nack RPCs in a background task. The keen eye will notice that
the lease loop is no longer blocked on any `await` points! (Well, except
for shutdown, where we do want to await these RPCs).

I will send a follow up PR to garbage collect the pending lease
extensions in `next_event()` instead of `extend()`. It seems cleaner to
me.

Then I will do some benchmarking.

Author: dbolduc

src/pubsub/src/subscriber/lease_loop.rs

@@ -63,7 +63,7 @@ impl LeaseLoop {
                     biased;
                     event = state.next_event() => {
                         match event {
-                            LeaseEvent::Flush => state.flush().await,
+                            LeaseEvent::Flush => state.flush(),
                             LeaseEvent::Extend => state.extend(),
                         }
                     },

src/pubsub/src/subscriber/lease_state.rs

@@ -26,6 +26,7 @@ use tokio::sync::oneshot::Sender;
 // Use a `tokio::time::Instant` to facilitate time-based unit testing.
 use tokio::task::JoinSet;
 use tokio::time::{Duration, Instant, Interval, interval_at};
+use tokio_util::task::TaskTracker;
 
 // An ack ID is less than 200 bytes. The limit for a request is 512kB. It should
 // be safe to fit 2500 Ack IDs in a single RPC.
@@ -117,6 +118,9 @@ where
     // How long messages can be kept under lease
     max_lease: Duration,
 
+    // In flight acks and nacks.
+    pending_acks_nacks: TaskTracker,
+
     // In flight lease extension operations.
     //
     // These are held separate from pending acks/nacks because we do not need to
@@ -149,6 +153,7 @@ where
             flush_interval,
             extend_interval,
             max_lease: options.max_lease,
+            pending_acks_nacks: TaskTracker::new(),
             pending_extends: JoinSet::new(),
         }
     }
@@ -201,22 +206,29 @@ where
     }
 
     /// Flush pending acks/nacks
-    pub(super) async fn flush(&mut self) {
-        // TODO(#3975) - await these concurrently.
+    pub(super) fn flush(&mut self) {
         let (to_ack, to_nack) = self.leases.drain();
         if !to_ack.is_empty() {
-            self.leaser.ack(to_ack).await;
+            let leaser = self.leaser.clone();
+            self.pending_acks_nacks
+                .spawn(async move { leaser.ack(to_ack).await });
         }
         if !to_nack.is_empty() {
-            self.leaser.nack(to_nack).await;
+            let leaser = self.leaser.clone();
+            self.pending_acks_nacks
+                .spawn(async move { leaser.nack(to_nack).await });
         }
 
         let (to_ack, to_nack) = self.eo_leases.drain();
         if !to_ack.is_empty() {
-            self.leaser.confirmed_ack(to_ack).await;
+            let leaser = self.leaser.clone();
+            self.pending_acks_nacks
+                .spawn(async move { leaser.confirmed_ack(to_ack).await });
         }
         if !to_nack.is_empty() {
-            self.leaser.nack(to_nack).await;
+            let leaser = self.leaser.clone();
+            self.pending_acks_nacks
+                .spawn(async move { leaser.nack(to_nack).await });
         }
     }
 
@@ -246,18 +258,20 @@ where
     ///
     /// This flushes all pending acks and nacks all other messages.
     pub(super) async fn shutdown(mut self) {
-        // TODO(#3975) - await these concurrently.
-
         // Note that if `WaitForProcessing` was selected by the application,
         // there are no messages under lease. They have all been processed.
         self.leases.evict();
         let (to_ack, to_nack) = self.leases.drain();
         if !to_ack.is_empty() {
-            self.leaser.ack(to_ack).await;
+            let leaser = self.leaser.clone();
+            self.pending_acks_nacks
+                .spawn(async move { leaser.ack(to_ack).await });
         }
         if !to_nack.is_empty() {
             // TODO(#4847) - this nack needs to be broken into batches.
-            self.leaser.nack(to_nack).await;
+            let leaser = self.leaser.clone();
+            self.pending_acks_nacks
+                .spawn(async move { leaser.nack(to_nack).await });
         }
 
         // TODO(#5109) - evicting exactly-once leases is ok, but not ideal.
@@ -269,9 +283,15 @@ where
         let (_, to_nack) = self.eo_leases.drain();
         if !to_nack.is_empty() {
             // TODO(#4847) - this nack needs to be broken into batches.
-            self.leaser.nack(to_nack).await;
+            let leaser = self.leaser.clone();
+            self.pending_acks_nacks
+                .spawn(async move { leaser.nack(to_nack).await });
         }
 
+        // Wait for pending acks/nacks to complete.
+        self.pending_acks_nacks.close();
+        self.pending_acks_nacks.wait().await;
+
         // Wait for pending lease extensions to complete. This is not useful in
         // practice, because we are nacking all the messages, but it simplifies
         // our tests.
@@ -334,6 +354,16 @@ pub(super) mod tests {
         let _ = pending_extends.join_all().await;
     }
 
+    async fn flush_and_await<L>(state: &mut LeaseState<L>)
+    where
+        L: Leaser + Clone + Send + 'static,
+    {
+        state.flush();
+        let pending_acks_nacks = std::mem::take(&mut state.pending_acks_nacks);
+        pending_acks_nacks.close();
+        pending_acks_nacks.wait().await;
+    }
+
     #[tokio::test(start_paused = true)]
     async fn basic_add_ack_nack() {
         let mock = MockLeaser::new();
@@ -529,7 +559,7 @@ pub(super) mod tests {
         // no messages under lease management.
         let mut state = LeaseState::new(Arc::new(mock), LeaseOptions::default());
         extend_and_await(&mut state).await;
-        state.flush().await;
+        state.flush();
         state.shutdown().await;
     }
 
@@ -588,7 +618,7 @@ pub(super) mod tests {
             state.eo_leases
         );
 
-        state.flush().await;
+        flush_and_await(&mut state).await;
         assert_eq!(
             TestLeases {
                 under_lease: test_ids(20..100),
@@ -621,6 +651,43 @@ pub(super) mod tests {
         );
     }
 
+    #[tokio::test(start_paused = true)]
+    async fn pending_acks_nacks_size_management() {
+        let mut mock = MockLeaser::new();
+        mock.expect_ack()
+            .times(1)
+            .withf(|v| *v == vec![test_id(1)])
+            .returning(|_| ());
+        mock.expect_nack()
+            .times(1)
+            .withf(|v| *v == vec![test_id(2)])
+            .returning(|_| ());
+
+        let mut state = LeaseState::new(Arc::new(mock), LeaseOptions::default());
+
+        state.add(test_id(1), at_least_once_info());
+        state.process(Ack(test_id(1)));
+
+        state.flush();
+        // Yield execution so the ack attempt can execute.
+        tokio::task::yield_now().await;
+        assert!(
+            state.pending_acks_nacks.is_empty(),
+            "The ack task should have completed. We should not hold onto it."
+        );
+
+        state.add(test_id(2), at_least_once_info());
+        state.process(Nack(test_id(2)));
+
+        state.flush();
+        // Yield execution so the nack attempt can execute.
+        tokio::task::yield_now().await;
+        assert!(
+            state.pending_acks_nacks.is_empty(),
+            "The nack task should have completed. We should not hold onto it."
+        );
+    }
+
     #[tokio::test(start_paused = true)]
     async fn extend_at_least_once() {
         let mut seq = mockall::Sequence::new();
@@ -727,7 +794,7 @@ pub(super) mod tests {
         extend_and_await(&mut state).await;
 
         // Flush the acks and confirm.
-        state.flush().await;
+        flush_and_await(&mut state).await;
         state.confirm(ack_results);
 
         // We should not extend the confirmed acks.
@@ -872,7 +939,7 @@ pub(super) mod tests {
         // With MAX_IDS_PER_RPC pending acks, the batch is full. We should flush it now.
         assert_eq!(state.next_event().await, LeaseEvent::Flush);
         assert_eq!(start.elapsed(), Duration::ZERO);
-        state.flush().await;
+        flush_and_await(&mut state).await;
 
         // With the batch is not full. The next event should occur on the interval timer.
         for i in MAX_IDS_PER_RPC..(2 * MAX_IDS_PER_RPC - 1) {
@@ -910,7 +977,7 @@ pub(super) mod tests {
         }
         assert_eq!(state.next_event().await, LeaseEvent::Flush);
         assert_eq!(start.elapsed(), Duration::ZERO);
-        state.flush().await;
+        flush_and_await(&mut state).await;
 
         // With the batch is not full. The next event should occur on the interval timer.
         for i in MAX_IDS_PER_RPC..(2 * MAX_IDS_PER_RPC - 1) {
@@ -948,7 +1015,7 @@ pub(super) mod tests {
         // With MAX_IDS_PER_RPC pending nacks, the batch is full. We should flush it now.
         assert_eq!(state.next_event().await, LeaseEvent::Flush);
         assert_eq!(start.elapsed(), Duration::ZERO);
-        state.flush().await;
+        flush_and_await(&mut state).await;
 
         // With the batch is not full. The next event should occur on the interval timer.
         for i in MAX_IDS_PER_RPC..(2 * MAX_IDS_PER_RPC - 1) {
@@ -987,7 +1054,7 @@ pub(super) mod tests {
         // With MAX_IDS_PER_RPC pending nacks for exactly once leases, the batch is full. We should flush it now.
         assert_eq!(state.next_event().await, LeaseEvent::Flush);
         assert_eq!(start.elapsed(), Duration::ZERO);
-        state.flush().await;
+        flush_and_await(&mut state).await;
 
         // With the batch is not full. The next event should occur on the interval timer.
         for i in MAX_IDS_PER_RPC..(2 * MAX_IDS_PER_RPC - 1) {