wip: improve gc cells

Author: pashinov

storage/src/store/shard_state/cell_storage.rs

@@ -50,6 +50,10 @@ impl CellStorage {
         StoreContext::new(&self.db, &self.raw_cells_cache, capacity)
     }
 
+    pub fn create_remove_ctx(&self) -> RemoveContext {
+        RemoveContext::new(&self.db, &self.raw_cells_cache)
+    }
+
     pub fn apply_temp_cell(&self, root: &HashBytes) -> Result<()> {
         const MAX_NEW_CELLS_BATCH_SIZE: usize = 10000;
 
@@ -554,6 +558,148 @@ impl StoreContext {
     }
 }
 
+#[derive(Clone)]
+struct RemovedCell {
+    old_rc: i64,
+    removes: u32,
+    refs: Vec<HashBytes>,
+}
+
+impl<'a> RemovedCell {
+    fn remove(&'a mut self) -> Result<Option<&'a [HashBytes]>, CellStorageError> {
+        self.removes += 1;
+        if self.removes as i64 <= self.old_rc {
+            Ok(self.next_refs())
+        } else {
+            Err(CellStorageError::CounterMismatch)
+        }
+    }
+
+    fn next_refs(&'a self) -> Option<&'a [HashBytes]> {
+        if self.old_rc > self.removes as i64 {
+            None
+        } else {
+            Some(&self.refs)
+        }
+    }
+}
+
+pub struct RemoveContext {
+    db: CellsDb,
+    raw_cache: Arc<RawCellsCache>,
+    transaction: FastDashMap<HashBytes, RemovedCell>,
+}
+
+impl RemoveContext {
+    fn new(db: &CellsDb, raw_cache: &Arc<RawCellsCache>) -> Self {
+        Self {
+            db: db.clone(),
+            raw_cache: raw_cache.clone(),
+            transaction: FastDashMap::with_capacity_and_hasher_and_shard_amount(
+                128,
+                Default::default(),
+                512,
+            ),
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        self.transaction.len()
+    }
+
+    pub fn remove_cell(&self, hash: &HashBytes) -> Result<(), CellStorageError> {
+        let mut stack = Vec::with_capacity(16);
+        stack.push(vec![*hash]);
+
+        let mut buffer = Vec::with_capacity(4);
+
+        // While some cells left
+        'outer: loop {
+            let Some(iter) = stack.last_mut() else {
+                break;
+            };
+
+            for cell_id in iter.iter() {
+                // Process the current cell.
+                let refs = match self.transaction.entry(*cell_id) {
+                    Entry::Occupied(mut v) => v.get_mut().remove()?.map(|v| v.to_vec()),
+                    Entry::Vacant(v) => {
+                        let old_rc =
+                            self.raw_cache
+                                .get_rc_for_delete(&self.db, cell_id, &mut buffer)?;
+                        debug_assert!(old_rc > 0);
+
+                        v.insert(RemovedCell {
+                            old_rc,
+                            removes: 1,
+                            refs: buffer.clone(),
+                        })
+                        .next_refs()
+                        .map(|v| v.to_vec())
+                    }
+                };
+
+                if let Some(refs) = refs {
+                    // And proceed to its refs if any.
+                    stack.push(refs);
+                    continue 'outer;
+                }
+            }
+
+            // Drop the current cell when all of its children were processed.
+            stack.pop();
+        }
+
+        // Clear big chunks of data before finalization
+        drop(stack);
+
+        Ok(())
+    }
+
+    pub fn finalize(self, batch: &mut WriteBatch) -> usize {
+        std::thread::scope(|s| {
+            let number_shards = self.transaction._shard_count();
+            // safety: we hold only read locks
+            let shards = unsafe { (0..number_shards).map(|i| self.transaction._get_read_shard(i)) };
+            let cache = &self.raw_cache;
+
+            // todo: clamp to number of cpus x2
+            for shard in shards {
+                // spawned threads will be joined at the end of the scope, so we don't need to store them
+                s.spawn(move || {
+                    for (key, item) in shard {
+                        let value = item.get();
+
+                        let new_rc = value.old_rc - value.removes as i64;
+                        cache.on_remove_cell(key, new_rc);
+                    }
+                });
+            }
+
+            let batch_update = s.spawn(|| {
+                let cells_cf = &self.db.cells.cf();
+
+                let total = self.transaction.len();
+
+                for kv in self.transaction.iter() {
+                    let key = kv.key();
+                    let value = kv.value();
+
+                    batch.merge_cf(
+                        cells_cf,
+                        key.as_slice(),
+                        refcount::encode_negative_refcount(value.removes),
+                    );
+                }
+
+                total
+            });
+
+            batch_update.join().expect("thread panicked")
+        })
+    }
+}
+
 #[derive(thiserror::Error, Debug)]
 pub enum CellStorageError {
     #[error("Cell not found in cell db")]

storage/src/store/shard_state/mod.rs

@@ -14,6 +14,7 @@ use tycho_block_util::state::*;
 use tycho_util::metrics::HistogramGuard;
 use tycho_util::FastHashMap;
 use weedb::rocksdb;
+use weedb::rocksdb::WriteBatch;
 
 use self::cell_storage::*;
 use self::store_state_raw::StoreStateContext;
@@ -293,13 +294,22 @@ impl ShardStateStorage {
                 let key = key.to_vec();
 
                 let (total, inner_alloc) = tokio::task::spawn_blocking(move || {
+                    let start = std::time::Instant::now();
                     let (stats, mut batch) = cell_storage.remove_cell(&alloc, &root_hash)?;
 
+                    let finished = start.elapsed().as_millis();
+
                     batch.delete_cf(&db.shard_states.get_unbounded_cf().bound(), key);
                     db.raw()
                         .rocksdb()
                         .write_opt(batch, db.cells.write_config())?;
 
+                    tracing::info!(
+                        ms = finished,
+                        full_ms = start.elapsed().as_millis(),
+                        "remove_cell_time"
+                    );
+
                     Ok::<_, anyhow::Error>((stats, alloc))
                 })
                 .await??;
@@ -359,6 +369,10 @@ impl ShardStateStorage {
         // Iterate all states and remove outdated
         let mut removed_states = 0usize;
         let mut removed_cells = 0usize;
+
+        let mut cur_block_id = BlockId::default();
+        let mut data_roots = Vec::with_capacity(16);
+
         loop {
             let _hist = HistogramGuard::begin("tycho_storage_state_data_gc_time");
             let (key, value) = match iter.item() {
@@ -380,24 +394,61 @@ impl ShardStateStorage {
                 continue;
             }
 
+            if cur_block_id == block_id {
+                data_roots.push(root_hash);
+
+                iter.next();
+                continue;
+            }
+
+            cur_block_id = block_id;
+            data_roots.clear();
+
             alloc.reset();
 
             {
                 let _guard = self.gc_lock.lock().await;
 
                 let db = self.cells_db.clone();
                 let cell_storage = self.cell_storage.clone();
+
                 let key = key.to_vec();
+                let data_roots = data_roots.clone();
 
                 let (total, inner_alloc) = tokio::task::spawn_blocking(move || {
-                    let (stats, mut batch) = cell_storage.remove_cell(&alloc, &root_hash)?;
+                    let start = std::time::Instant::now();
+
+                    let ctx = cell_storage.create_remove_ctx();
+                    rayon::scope(|s| {
+                        for data_root in data_roots.iter() {
+                            s.spawn(|_| {
+                                ctx.remove_cell(data_root).unwrap();
+                            });
+                        }
+                    });
+
+                    let mut batch = WriteBatch::with_capacity_bytes(ctx.len() * (32 + 8 + 8));
+
+                    let total = ctx.finalize(&mut batch);
+
+                    let finished = start.elapsed().as_millis();
+
+                    batch.delete_cf(
+                        &db.shard_state_data.get_unbounded_cf().bound(),
+                        key.as_slice(),
+                    );
 
-                    batch.delete_cf(&db.shard_state_data.get_unbounded_cf().bound(), key);
                     db.raw()
                         .rocksdb()
                         .write_opt(batch, db.cells.write_config())?;
 
-                    Ok::<_, anyhow::Error>((stats, alloc))
+                    tracing::info!(
+                        ms = finished,
+                        full_ms = start.elapsed().as_millis(),
+                        "accounts remove_cell_time"
+                    );
+
+                    Ok::<_, anyhow::Error>((total, alloc))
                 })
                 .await??;