From 130aa35d79df01f511b7b6ccd5c1e0918001c114 Mon Sep 17 00:00:00 2001
From: aldenhu <msmouse@gmail.com>
Date: Tue, 15 Oct 2024 23:36:05 +0000
Subject: [PATCH] LedgerUpdateOutput drop with default dropper

---
 Cargo.lock                                    |  2 ++
 crates/aptos-drop-helper/Cargo.toml           |  1 +
 crates/aptos-drop-helper/src/lib.rs           | 27 ++++++++++++++++++-
 execution/executor-types/Cargo.toml           |  1 +
 .../src/ledger_update_output.rs               |  5 ++--
 5 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0a54c0826e3c9..c7d752ed1fca8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1348,6 +1348,7 @@ version = "0.1.0"
 dependencies = [
  "aptos-infallible",
  "aptos-metrics-core",
+ "derive_more",
  "once_cell",
  "threadpool",
 ]
@@ -1547,6 +1548,7 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "aptos-crypto",
+ "aptos-drop-helper",
  "aptos-scratchpad",
  "aptos-secure-net",
  "aptos-storage-interface",
diff --git a/crates/aptos-drop-helper/Cargo.toml b/crates/aptos-drop-helper/Cargo.toml
index 54dc6fc482c16..936ef297eeebc 100644
--- a/crates/aptos-drop-helper/Cargo.toml
+++ b/crates/aptos-drop-helper/Cargo.toml
@@ -15,5 +15,6 @@ rust-version = { workspace = true }
 [dependencies]
 aptos-infallible = { workspace = true }
 aptos-metrics-core = { workspace = true }
+derive_more = { workspace = true }
 once_cell = { workspace = true }
 threadpool = { workspace = true }
diff --git a/crates/aptos-drop-helper/src/lib.rs b/crates/aptos-drop-helper/src/lib.rs
index ce91987bf6525..169aae9c41fe3 100644
--- a/crates/aptos-drop-helper/src/lib.rs
+++ b/crates/aptos-drop-helper/src/lib.rs
@@ -2,7 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::async_concurrent_dropper::AsyncConcurrentDropper;
+use derive_more::{Deref, DerefMut};
 use once_cell::sync::Lazy;
+use std::mem::ManuallyDrop;
 
 pub mod async_concurrent_dropper;
 pub mod async_drop_queue;
@@ -11,8 +13,31 @@ mod metrics;
 pub static DEFAULT_DROPPER: Lazy<AsyncConcurrentDropper> =
     Lazy::new(|| AsyncConcurrentDropper::new("default", 32, 8));
 
-/// Arc<T: ArcAsyncDrop> will be `Send + 'static`, which is requried to be able to drop Arc<T>
+/// Arc<T: ArcAsyncDrop> will be `Send + 'static`, which is required to be able to drop Arc<T>
 /// in another thread
 pub trait ArcAsyncDrop: Send + Sync + 'static {}
 
 impl<T: Send + Sync + 'static> ArcAsyncDrop for T {}
+
+#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Deref, DerefMut)]
+#[repr(transparent)]
+pub struct DropHelper<T: Send + 'static> {
+    #[deref]
+    #[deref_mut]
+    inner: ManuallyDrop<T>,
+}
+
+impl<T: Send + 'static> DropHelper<T> {
+    pub fn new(inner: T) -> Self {
+        Self {
+            inner: ManuallyDrop::new(inner),
+        }
+    }
+}
+
+impl<T: Send + 'static> Drop for DropHelper<T> {
+    fn drop(&mut self) {
+        let inner = unsafe { ManuallyDrop::take(&mut self.inner) };
+        DEFAULT_DROPPER.schedule_drop(inner);
+    }
+}
diff --git a/execution/executor-types/Cargo.toml b/execution/executor-types/Cargo.toml
index 5c1517f7f5973..b27f6f09efbdc 100644
--- a/execution/executor-types/Cargo.toml
+++ b/execution/executor-types/Cargo.toml
@@ -15,6 +15,7 @@ rust-version = { workspace = true }
 [dependencies]
 anyhow = { workspace = true }
 aptos-crypto = { workspace = true }
+aptos-drop-helper = { workspace = true }
 aptos-scratchpad = { workspace = true }
 aptos-secure-net = { workspace = true }
 aptos-storage-interface = { workspace = true }
diff --git a/execution/executor-types/src/ledger_update_output.rs b/execution/executor-types/src/ledger_update_output.rs
index 0ccaa3dadebc8..f0928c8f4bf8f 100644
--- a/execution/executor-types/src/ledger_update_output.rs
+++ b/execution/executor-types/src/ledger_update_output.rs
@@ -6,6 +6,7 @@
 use crate::StateComputeResult;
 use anyhow::{ensure, Result};
 use aptos_crypto::HashValue;
+use aptos_drop_helper::DropHelper;
 use aptos_storage_interface::cached_state_view::ShardedStateCache;
 use aptos_types::{
     contract_event::ContractEvent,
@@ -24,7 +25,7 @@ use std::sync::Arc;
 #[derive(Clone, Debug, Default, Deref)]
 pub struct LedgerUpdateOutput {
     #[deref]
-    inner: Arc<Inner>,
+    inner: Arc<DropHelper<Inner>>,
 }
 
 impl LedgerUpdateOutput {
@@ -76,7 +77,7 @@ impl LedgerUpdateOutput {
 
     fn new_impl(inner: Inner) -> Self {
         Self {
-            inner: Arc::new(inner),
+            inner: Arc::new(DropHelper::new(inner)),
         }
     }