Optimizing column entries bookkeeping (#393)

Author: taco-paco

magicblock-api/src/tickers.rs

@@ -160,13 +160,9 @@ pub fn init_system_metrics_ticker(
     fn set_accounts_count(bank: &Bank) {
         metrics::set_accounts_count(bank.accounts_db.get_accounts_count());
     }
+
     let ledger = ledger.clone();
     let bank = bank.clone();
-    try_set_ledger_storage_size(&ledger);
-    set_accounts_storage_size(&bank);
-    set_accounts_count(&bank);
-    try_set_ledger_counts(&ledger);
-
     tokio::task::spawn(async move {
         loop {
             tokio::select! {

magicblock-ledger/src/database/columns.rs

@@ -1,16 +1,9 @@
-use std::sync::atomic::{AtomicI64, Ordering};
-
 use byteorder::{BigEndian, ByteOrder};
-use log::*;
 use serde::{de::DeserializeOwned, Serialize};
 use solana_sdk::{clock::Slot, pubkey::Pubkey, signature::Signature};
 use solana_storage_proto::convert::generated;
 
 use super::meta;
-use crate::{
-    database::{iterator::IteratorMode, ledger_column::LedgerColumn},
-    errors::LedgerResult,
-};
 
 /// Column family for Transaction Status
 const TRANSACTION_STATUS_CF: &str = "transaction_status";
@@ -677,24 +670,3 @@ pub fn should_enable_compression<C: 'static + Column + ColumnName>() -> bool {
 // Column Queries
 // -----------------
 pub(crate) const DIRTY_COUNT: i64 = -1;
-pub fn count_column_using_cache<C: Column + ColumnName>(
-    column: &LedgerColumn<C>,
-    cached_value: &AtomicI64,
-) -> LedgerResult<i64> {
-    let cached = cached_value.load(Ordering::Relaxed);
-    // NOTE: a value of -1 indicates that the cached value is dirty
-    if cached != DIRTY_COUNT {
-        return Ok(cached);
-    }
-    column
-        .iter(IteratorMode::Start)
-        .map(Iterator::count)
-        .map(|val| if val > i64::MAX as usize {
-            // NOTE: this value is only used for metrics/diagnostics and
-            // aside from the fact that we will never encounter this case,
-            // it is good enough to return i64::MAX
-            error!("Column {} count is too large: {} for metrics, returning max.", C::NAME, val);
-            i64::MAX
-        } else { val as i64 })
-        .inspect(|updated| cached_value.store(*updated, Ordering::Relaxed))
-}

magicblock-ledger/src/database/db.rs

@@ -1,4 +1,8 @@
-use std::{marker::PhantomData, path::Path, sync::Arc};
+use std::{
+    marker::PhantomData,
+    path::Path,
+    sync::{atomic::AtomicI64, Arc},
+};
 
 use bincode::deserialize;
 use rocksdb::{ColumnFamily, DBRawIterator, LiveFile};
@@ -12,7 +16,10 @@ use super::{
     rocks_db::Rocks,
     write_batch::WriteBatch,
 };
-use crate::{errors::LedgerError, metrics::PerfSamplingStatus};
+use crate::{
+    database::columns::DIRTY_COUNT, errors::LedgerError,
+    metrics::PerfSamplingStatus,
+};
 
 #[derive(Debug)]
 pub struct Database {
@@ -90,6 +97,7 @@ impl Database {
             column_options: Arc::clone(&self.column_options),
             read_perf_status: PerfSamplingStatus::default(),
             write_perf_status: PerfSamplingStatus::default(),
+            entry_counter: AtomicI64::new(DIRTY_COUNT),
         }
     }
 

magicblock-ledger/src/database/ledger_column.rs

@@ -1,6 +1,13 @@
-use std::{marker::PhantomData, sync::Arc};
+use std::{
+    marker::PhantomData,
+    sync::{
+        atomic::{AtomicI64, Ordering},
+        Arc,
+    },
+};
 
 use bincode::{deserialize, serialize};
+use log::{error, warn};
 use prost::Message;
 use rocksdb::{properties as RocksProperties, ColumnFamily};
 use serde::de::DeserializeOwned;
@@ -14,7 +21,7 @@ use super::{
     rocks_db::Rocks,
 };
 use crate::{
-    database::write_batch::WriteBatch,
+    database::{columns::DIRTY_COUNT, write_batch::WriteBatch},
     errors::{LedgerError, LedgerResult},
     metrics::{
         maybe_enable_rocksdb_perf, report_rocksdb_read_perf,
@@ -34,6 +41,15 @@ where
     pub column_options: Arc<LedgerColumnOptions>,
     pub read_perf_status: PerfSamplingStatus,
     pub write_perf_status: PerfSamplingStatus,
+    // We are caching the column item counts since they are expensive to obtain.
+    // `-1` indicates that they are "dirty"    //
+    //     // We are using an i64 to make this work even though the counts are usize,
+    //     // however if we had 50,000 transactions/sec and 50ms slots for 100 years then:
+    //     //
+    //     // slots:   200 * 3600 * 24 * 365 * 100 =           630,720,000,000
+    //     // txs:  50,000 * 3600 * 24 * 365 * 100 =       157,680,000,000,000
+    //     // i64::MAX                             = 9,223,372,036,854,775,807
+    pub entry_counter: AtomicI64,
 }
 
 impl<C: Column + ColumnName> LedgerColumn<C> {
@@ -260,6 +276,39 @@ where
     pub fn flush(&self) -> LedgerResult<()> {
         self.backend.flush_cf(self.handle())
     }
+
+    pub fn count_column_using_cache(&self) -> LedgerResult<i64> {
+        let cached = self.entry_counter.load(Ordering::Relaxed);
+        if cached != DIRTY_COUNT {
+            return Ok(cached);
+        }
+
+        self
+            .iter(IteratorMode::Start)
+            .map(Iterator::count)
+            .map(|val| if val > i64::MAX as usize {
+                // NOTE: this value is only used for metrics/diagnostics and
+                // aside from the fact that we will never encounter this case,
+                // it is good enough to return i64::MAX
+                error!("Column {} count is too large: {} for metrics, returning max.", C::NAME, val);
+                i64::MAX
+            } else { val as i64 })
+            .inspect(|updated| self.entry_counter.store(*updated, Ordering::Relaxed))
+    }
+
+    /// Increases entries counter if it's not [`DIRTY_COUNT`]
+    /// Otherwise just skips it until it is set
+    #[inline(always)]
+    pub fn try_increase_entry_counter(&self, by: u64) {
+        try_increase_entry_counter(&self.entry_counter, by);
+    }
+
+    /// Decreases entries counter if it's not [`DIRTY_COUNT`]
+    /// Otherwise just skips it until it is set
+    #[inline(always)]
+    pub fn try_decrease_entry_counter(&self, by: u64) {
+        try_decrease_entry_counter(&self.entry_counter, by);
+    }
 }
 
 impl<C> LedgerColumn<C>
@@ -490,3 +539,68 @@ where
         })
     }
 }
+
+/// Increases entries counter if it's not [`DIRTY_COUNT`]
+/// Otherwise just skips it until it is set
+pub fn try_increase_entry_counter(entry_counter: &AtomicI64, by: u64) {
+    loop {
+        let prev = entry_counter.load(Ordering::Acquire);
+        if prev == DIRTY_COUNT {
+            return;
+        }
+
+        // In case value changed to [`DIRTY_COUNT`] in between
+        if entry_counter
+            .compare_exchange(
+                prev,
+                prev + by as i64,
+                Ordering::AcqRel,
+                Ordering::Relaxed,
+            )
+            .is_ok()
+        {
+            return;
+        }
+    }
+}
+
+/// Decreases entries counter if it's not [`DIRTY_COUNT`]
+/// Otherwise just skips it until it is set
+pub fn try_decrease_entry_counter(entry_counter: &AtomicI64, by: u64) {
+    loop {
+        let prev = entry_counter.load(Ordering::Acquire);
+        if prev == DIRTY_COUNT {
+            return;
+        }
+
+        let new = prev - by as i64;
+        if new >= 0 {
+            // In case value changed to [`DIRTY_COUNT`] in between
+            if entry_counter
+                .compare_exchange(
+                    prev,
+                    new,
+                    Ordering::AcqRel,
+                    Ordering::Relaxed,
+                )
+                .is_ok()
+            {
+                return;
+            }
+        } else {
+            warn!("Negative entry counter!");
+            // In case value fixed to valid one in between
+            if entry_counter
+                .compare_exchange(
+                    prev,
+                    DIRTY_COUNT,
+                    Ordering::AcqRel,
+                    Ordering::Relaxed,
+                )
+                .is_ok()
+            {
+                return;
+            }
+        }
+    }
+}