AsyncConcurrentDropper: wait_for_backlog_drop

Author: msmouse

crates/aptos-drop-helper/src/async_concurrent_dropper.rs

@@ -1,10 +1,13 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::metrics::TIMER;
+use crate::metrics::{GAUGE, TIMER};
 use aptos_infallible::Mutex;
-use aptos_metrics_core::TimerHelper;
-use std::sync::mpsc::{channel, Receiver, Sender};
+use aptos_metrics_core::{IntGaugeHelper, TimerHelper};
+use std::sync::{
+    mpsc::{channel, Receiver, Sender},
+    Arc, Condvar,
+};
 use threadpool::ThreadPool;
 
 /// A helper to send things to a thread pool for asynchronous dropping.
@@ -15,26 +18,17 @@ use threadpool::ThreadPool;
 ///      to another thing being waiting for a slot to be available.
 pub struct AsyncConcurrentDropper {
     name: &'static str,
-    token_tx: Sender<()>,
-    token_rx: Mutex<Receiver<()>>,
+    num_tasks_tracker: Arc<NumTasksTracker>,
     /// use dedicated threadpool to minimize the possibility of dead lock
     thread_pool: ThreadPool,
 }
 
 impl AsyncConcurrentDropper {
-    pub fn new(name: &'static str, max_async_drops: usize, num_threads: usize) -> Self {
-        let (token_tx, token_rx) = channel();
-        for _ in 0..max_async_drops {
-            token_tx
-                .send(())
-                .expect("DropHelper: Failed to buffer initial tokens.");
-        }
-        let thread_pool = ThreadPool::new(num_threads);
+    pub fn new(name: &'static str, max_tasks: usize, num_threads: usize) -> Self {
         Self {
             name,
-            token_tx,
-            token_rx: Mutex::new(token_rx),
-            thread_pool,
+            num_tasks_tracker: Arc::new(NumTasksTracker::new(max_tasks)),
+            thread_pool: ThreadPool::new(num_threads),
         }
     }
 
@@ -48,13 +42,19 @@ impl AsyncConcurrentDropper {
         rx
     }
 
+    pub fn wait_for_backlog_drop(&self, no_more_than: usize) {
+        let _timer = TIMER.timer_with(&[self.name, "wait_for_backlog_drop"]);
+        self.num_tasks_tracker.wait_for_backlog_drop(no_more_than);
+    }
+
     fn schedule_drop_impl<V: Send + 'static>(&self, v: V, notif_sender_opt: Option<Sender<()>>) {
         let _timer = TIMER.timer_with(&[self.name, "enqueue_drop"]);
+        let num_tasks = self.num_tasks_tracker.inc();
+        GAUGE.set_with(&[self.name, "num_tasks"], num_tasks as i64);
 
-        self.token_rx.lock().recv().unwrap();
-
-        let token_tx = self.token_tx.clone();
         let name = self.name;
+        let num_tasks_tracker = self.num_tasks_tracker.clone();
+
         self.thread_pool.execute(move || {
             let _timer = TIMER.timer_with(&[name, "real_drop"]);
 
@@ -64,15 +64,54 @@ impl AsyncConcurrentDropper {
                 sender.send(()).ok();
             }
 
-            token_tx.send(()).ok();
+            num_tasks_tracker.dec();
         })
     }
 }
 
+struct NumTasksTracker {
+    lock: Mutex<usize>,
+    cvar: Condvar,
+    max_tasks: usize,
+}
+
+impl NumTasksTracker {
+    fn new(max_tasks: usize) -> Self {
+        Self {
+            lock: Mutex::new(0),
+            cvar: Condvar::new(),
+            max_tasks,
+        }
+    }
+
+    fn inc(&self) -> usize {
+        let mut num_tasks = self.lock.lock();
+        while *num_tasks >= self.max_tasks {
+            num_tasks = self.cvar.wait(num_tasks).expect("lock poisoned.");
+        }
+        *num_tasks += 1;
+        *num_tasks
+    }
+
+    fn dec(&self) {
+        let mut num_tasks = self.lock.lock();
+        *num_tasks -= 1;
+        self.cvar.notify_all();
+    }
+
+    fn wait_for_backlog_drop(&self, no_more_than: usize) {
+        let mut num_tasks = self.lock.lock();
+        while *num_tasks > no_more_than {
+            num_tasks = self.cvar.wait(num_tasks).expect("lock poisoned.");
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::AsyncConcurrentDropper;
-    use std::{thread::sleep, time::Duration};
+    use std::{sync::Arc, thread::sleep, time::Duration};
+    use threadpool::ThreadPool;
 
     struct SlowDropper;
 
@@ -117,4 +156,43 @@ mod tests {
         s.schedule_drop(SlowDropper);
         assert!(now.elapsed() < Duration::from_millis(400));
     }
+
+    fn async_wait(
+        thread_pool: &ThreadPool,
+        dropper: &Arc<AsyncConcurrentDropper>,
+        no_more_than: usize,
+    ) {
+        let dropper = Arc::clone(dropper);
+        thread_pool.execute(move || dropper.wait_for_backlog_drop(no_more_than));
+    }
+
+    #[test]
+    fn test_wait_for_backlog_drop() {
+        let s = Arc::new(AsyncConcurrentDropper::new("test", 8, 4));
+        let t = ThreadPool::new(4);
+        let now = std::time::Instant::now();
+        for _ in 0..8 {
+            s.schedule_drop(SlowDropper);
+        }
+        assert!(now.elapsed() < Duration::from_millis(200));
+        s.wait_for_backlog_drop(8);
+        assert!(now.elapsed() < Duration::from_millis(200));
+        async_wait(&t, &s, 8);
+        async_wait(&t, &s, 8);
+        async_wait(&t, &s, 7);
+        async_wait(&t, &s, 4);
+        t.join();
+        assert!(now.elapsed() > Duration::from_millis(200));
+        assert!(now.elapsed() < Duration::from_millis(400));
+        s.wait_for_backlog_drop(4);
+        assert!(now.elapsed() < Duration::from_millis(400));
+        async_wait(&t, &s, 3);
+        async_wait(&t, &s, 2);
+        async_wait(&t, &s, 1);
+        t.join();
+        assert!(now.elapsed() > Duration::from_millis(400));
+        assert!(now.elapsed() < Duration::from_millis(600));
+        s.wait_for_backlog_drop(0);
+        assert!(now.elapsed() < Duration::from_millis(600));
+    }
 }

crates/aptos-drop-helper/src/metrics.rs

@@ -1,7 +1,9 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-use aptos_metrics_core::{exponential_buckets, register_histogram_vec, HistogramVec};
+use aptos_metrics_core::{
+    exponential_buckets, register_histogram_vec, register_int_gauge_vec, HistogramVec, IntGaugeVec,
+};
 use once_cell::sync::Lazy;
 
 pub static TIMER: Lazy<HistogramVec> = Lazy::new(|| {
@@ -13,3 +15,12 @@ pub static TIMER: Lazy<HistogramVec> = Lazy::new(|| {
     )
     .unwrap()
 });
+
+pub static GAUGE: Lazy<IntGaugeVec> = Lazy::new(|| {
+    register_int_gauge_vec!(
+        "aptos_drop_helper_gauges",
+        "Various gauges to help debugging.",
+        &["helper_name", "name"],
+    )
+    .unwrap()
+});