SmtAncestor: move backpressure to the getting side

Author: msmouse

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

@@ -28,7 +28,7 @@ impl AsyncConcurrentDropper {
         Self {
             name,
             num_tasks_tracker: Arc::new(NumTasksTracker::new(max_tasks)),
-            thread_pool: ThreadPool::new(num_threads),
+            thread_pool: ThreadPool::with_name(format!("{}_conc_dropper", name), num_threads),
         }
     }
 

storage/aptosdb/src/state_store/mod.rs

@@ -227,7 +227,7 @@ impl StateDb {
 
 impl DbReader for StateStore {
     fn get_buffered_state_base(&self) -> Result<SparseMerkleTree<StateValue>> {
-        Ok(self.smt_ancestors.lock().get_youngest()?)
+        Ok(self.smt_ancestors.lock().get_youngest())
     }
 
     /// Returns the latest state snapshot strictly before `next_version` if any.

storage/scratchpad/src/sparse_merkle/ancestors.rs

@@ -3,96 +3,44 @@
 
 #![forbid(unsafe_code)]
 
-use crate::{sparse_merkle::metrics::TIMER, SparseMerkleTree};
+use crate::{
+    sparse_merkle::{dropper::SUBTREE_DROPPER, metrics::TIMER},
+    SparseMerkleTree,
+};
 use aptos_crypto::hash::CryptoHash;
-use aptos_drop_helper::async_drop_queue::AsyncDropQueue;
 use aptos_infallible::Mutex;
-use aptos_logger::prelude::*;
 use aptos_metrics_core::TimerHelper;
-use std::{
-    collections::VecDeque,
-    sync::{Arc, MutexGuard},
-    time::Duration,
-};
+use std::sync::Arc;
 
-type Result<V> = std::result::Result<V, Error>;
-
-/// Keep the oldest SMTs in a central place so that:
-/// 1. elsewhere, dropping the SMT will be fast since this indirectly holds a ref to every tree in
-///    the entire forest.
-/// 2. this must be invoked somewhere in the critical path so that it can provide some
-///    back pressure to slow things down to prevent from leaking memory
+/// A container to track the ancestor of the SMTs that represent a committed state (older state is
+/// guaranteed to be found in persisted storage.
+/// When being queried, back pressure (a slow down) is provided in order to make sure not too many
+/// SMTs are being kept in memory.
 #[derive(Clone, Debug)]
 pub struct SmtAncestors<V: Clone + Send + Sync + 'static> {
-    /// Keep a queue of ancestors, in hope that the when the oldest is being evicted, it's the
-    /// the last ref of it, which means the drop latency doesn't impact other code paths.
-    ancestors: Arc<Mutex<VecDeque<SparseMerkleTree<V>>>>,
-    /// Drop the oldest ancestor asynchronously in good cases, with limited backlog, providing
-    /// back pressure when the drops are slow in order to avoid memory leak.
-    drop_queue: Arc<AsyncDropQueue>,
+    youngest: Arc<Mutex<SparseMerkleTree<V>>>,
 }
 
 impl<V: CryptoHash + Clone + Send + Sync + 'static> SmtAncestors<V> {
-    const MAX_PENDING_DROPS: usize = 4;
-    const NUM_ANCESTORS: usize = 8;
+    const MAX_PENDING_DROPS: usize = 8;
 
     pub fn new(ancestor: SparseMerkleTree<V>) -> Self {
         Self {
-            ancestors: Arc::new(Mutex::new(VecDeque::from(vec![ancestor]))),
-            drop_queue: Arc::new(AsyncDropQueue::new(
-                "smt_ancestors",
-                Self::MAX_PENDING_DROPS,
-            )),
+            youngest: Arc::new(Mutex::new(ancestor)),
         }
     }
 
-    fn ancestors(&self) -> MutexGuard<VecDeque<SparseMerkleTree<V>>> {
-        self.ancestors.lock()
-    }
+    pub fn get_youngest(&self) -> SparseMerkleTree<V> {
+        let _timer = TIMER.timer_with(&["get_youngest_ancestor"]);
 
-    pub fn get_youngest(&self) -> Result<SparseMerkleTree<V>> {
-        self.ancestors()
-            .back()
-            .map(SparseMerkleTree::clone)
-            .ok_or(Error::NotFound)
-    }
-
-    pub fn add(&self, youngest: SparseMerkleTree<V>) {
-        let _timer = TIMER.timer_with(&["add_smt_ancestor"]);
+        // The back pressure is on the getting side (which is the execution side) so that it's less
+        // likely for a lot of blocks locking the same old base SMT.
+        SUBTREE_DROPPER.wait_for_backlog_drop(Self::MAX_PENDING_DROPS);
 
-        let mut ancestors = self.ancestors();
-        ancestors.push_back(youngest);
-
-        if ancestors.len() > Self::NUM_ANCESTORS {
-            let oldest = ancestors.pop_front().unwrap();
-            oldest.log_generation("evict_ancestor");
-            if !oldest.is_the_only_ref() {
-                sample!(
-                    SampleRate::Duration(Duration::from_secs(1)),
-                    error!(
-                        "Oldest SMT being tracked by SmtAncestors is still referenced elsewhere. Potential memory leak.",
-                    )
-                );
-            } else {
-                self.drop_queue.enqueue_drop(oldest);
-            }
-        }
+        self.youngest.lock().clone()
     }
 
-    pub fn replace_with(&self, other: SmtAncestors<V>) {
-        let Self {
-            ancestors,
-            drop_queue: _drop_queue,
-        } = other;
-
-        *self.ancestors.lock() = Arc::into_inner(ancestors)
-            .expect("Not the only ref.")
-            .into_inner();
+    pub fn add(&self, youngest: SparseMerkleTree<V>) {
+        *self.youngest.lock() = youngest;
     }
 }
-
-#[derive(Debug, thiserror::Error, Eq, PartialEq)]
-pub enum Error {
-    #[error("Ancestor SMT not found.")]
-    NotFound,
-}

storage/scratchpad/src/sparse_merkle/mod.rs

@@ -598,10 +598,6 @@ where
             None
         }
     }
-
-    pub(crate) fn is_the_only_ref(&self) -> bool {
-        Arc::strong_count(&self.inner) == 1
-    }
 }
 
 /// In tests and benchmark, reference to ancestors are manually managed

storage/scratchpad/src/sparse_merkle/sparse_merkle_test.rs

@@ -400,6 +400,7 @@ fn test_update() {
 
     // Now prune smt1.
     drop(smt1);
+    SUBTREE_DROPPER.wait_for_backlog_drop(0);
 
     // Verify oldest ancestor
     assert_eq_pointee(&smt2.get_oldest_ancestor(), &smt2);