From 634e70ae66086fedad00c58ae984c7a78ebac811 Mon Sep 17 00:00:00 2001
From: Balaji Arun <balaji@aptoslabs.com>
Date: Mon, 29 Jan 2024 13:30:24 -0800
Subject: [PATCH] [dag] use DropGuard to drop handles

---
 consensus/src/dag/dag_driver.rs | 18 ++++--------------
 1 file changed, 4 insertions(+), 14 deletions(-)

diff --git a/consensus/src/dag/dag_driver.rs b/consensus/src/dag/dag_driver.rs
index c758b6a3b7905..facae42d4d950 100644
--- a/consensus/src/dag/dag_driver.rs
+++ b/consensus/src/dag/dag_driver.rs
@@ -29,7 +29,7 @@ use aptos_consensus_types::common::{Author, Payload, PayloadFilter};
 use aptos_crypto::hash::CryptoHash;
 use aptos_infallible::Mutex;
 use aptos_logger::{debug, error};
-use aptos_reliable_broadcast::ReliableBroadcast;
+use aptos_reliable_broadcast::{DropGuard, ReliableBroadcast};
 use aptos_time_service::{TimeService, TimeServiceTrait};
 use aptos_types::{block_info::Round, epoch_state::EpochState};
 use aptos_validator_transaction_pool as vtxn_pool;
@@ -53,7 +53,7 @@ pub(crate) struct DagDriver {
     payload_client: Arc<dyn PayloadClient>,
     reliable_broadcast: Arc<ReliableBroadcast<DAGMessage, ExponentialBackoff, DAGRpcResult>>,
     time_service: TimeService,
-    rb_handles: Mutex<BoundedVecDeque<(AbortHandle, u64)>>,
+    rb_handles: Mutex<BoundedVecDeque<(DropGuard, u64)>>,
     storage: Arc<dyn DAGStorage>,
     order_rule: Mutex<OrderRule>,
     fetch_requester: Arc<dyn TFetchRequester>,
@@ -323,14 +323,13 @@ impl DagDriver {
             debug!("Finish reliable broadcast for round {}", round);
         };
         tokio::spawn(Abortable::new(task, abort_registration));
-        if let Some((prev_handle, prev_round_timestamp)) = self
+        if let Some((_handle, prev_round_timestamp)) = self
             .rb_handles
             .lock()
-            .push_back((abort_handle, timestamp))
+            .push_back((DropGuard::new(abort_handle), timestamp))
         {
             // TODO: this observation is inaccurate.
             observe_round(prev_round_timestamp, RoundStage::Finished);
-            prev_handle.abort();
         }
     }
 }
@@ -361,15 +360,6 @@ impl RpcHandler for DagDriver {
     }
 }
 
-impl Drop for DagDriver {
-    fn drop(&mut self) {
-        let abort_handles = self.rb_handles.lock();
-        for (handle, _) in abort_handles.iter() {
-            handle.abort()
-        }
-    }
-}
-
 struct BoundedVecDeque<T> {
     inner: VecDeque<T>,
     capacity: usize,