Add db-stats command.

It lists info about `beacon_fork_choice` db key/value sizes.

Co-authored-by: weekday <weekday@grandine.io>

database/src/lib.rs

@@ -9,7 +9,7 @@ use anyhow::Result;
 use bytesize::ByteSize;
 use im::OrdMap;
 use itertools::Either;
-use libmdbx::{DatabaseFlags, Environment, Geometry, WriteFlags};
+use libmdbx::{DatabaseFlags, Environment, Geometry, ObjectLength, Stat, WriteFlags};
 use log::info;
 use snap::raw::{Decoder, Encoder};
 use std_ext::ArcExt as _;
@@ -23,13 +23,20 @@ const MAX_NAMED_DATABASES: usize = 10;
 pub struct Database(DatabaseKind);
 
 impl Database {
-    pub fn persistent(name: &str, directory: impl AsRef<Path>, max_size: ByteSize) -> Result<Self> {
+    pub fn persistent(
+        name: &str,
+        directory: impl AsRef<Path>,
+        max_size: ByteSize,
+        read_only: bool,
+    ) -> Result<Self> {
         // If a database with the legacy name exists, keep using it.
         // Otherwise, create a new database with the specified name.
         // This check will not force existing users to resync.
         let legacy_name = directory.as_ref().to_str().ok_or(Error)?;
 
-        fs_err::create_dir_all(&directory)?;
+        if !read_only {
+            fs_err::create_dir_all(&directory)?;
+        }
 
         // TODO(Grandine Team): The call to `set_max_dbs` and `MAX_NAMED_DATABASES` should be
         //                      unnecessary if the default database is used.
@@ -48,7 +55,10 @@ impl Database {
 
         let database_name = if existing_db.is_err() {
             info!("database: {legacy_name} with name {name}");
-            transaction.create_db(Some(name), DatabaseFlags::default())?;
+            if !read_only {
+                transaction.create_db(Some(name), DatabaseFlags::default())?;
+            }
+
             name
         } else {
             info!("legacy database: {legacy_name}");
@@ -194,6 +204,54 @@ impl Database {
         .transpose()
     }
 
+    pub fn db_stats(&self) -> Result<Option<Stat>> {
+        match self.kind() {
+            DatabaseKind::Persistent {
+                database_name,
+                environment,
+            } => {
+                let transaction = environment.begin_ro_txn()?;
+                let database = transaction.open_db(Some(database_name))?;
+
+                Some(transaction.db_stat(&database)?)
+            }
+            DatabaseKind::InMemory { map: _ } => None,
+        }
+        .pipe(Ok)
+    }
+
+    pub fn iterate_all_keys_with_lengths(
+        &self,
+    ) -> Result<impl Iterator<Item = Result<(Cow<[u8]>, usize)>>> {
+        match self.kind() {
+            DatabaseKind::Persistent {
+                database_name,
+                environment,
+            } => {
+                let transaction = environment.begin_ro_txn()?;
+                let database = transaction.open_db(Some(database_name))?;
+
+                let mut cursor = transaction.cursor(&database)?;
+
+                core::iter::from_fn(move || cursor.next().transpose())
+                    .map(|result| {
+                        let (key, ObjectLength(length)) = result?;
+                        Ok((key, length))
+                    })
+                    .pipe(Either::Left)
+            }
+            DatabaseKind::InMemory { map } => {
+                let map = map.lock().expect("in-memory database mutex is poisoned");
+
+                map.clone()
+                    .into_iter()
+                    .map(|(key, value)| Ok((Cow::Owned(key.to_vec()), value.len())))
+                    .pipe(Either::Right)
+            }
+        }
+        .pipe(Ok)
+    }
+
     pub fn iterator_ascending(
         &self,
         range: RangeFrom<impl AsRef<[u8]>>,
@@ -574,6 +632,34 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(build_persistent_database)]
+    #[test_case(build_in_memory_database)]
+    fn test_all_keys_iterator_with_lengths(constructor: Constructor) -> Result<()> {
+        let database = constructor()?;
+        let values = database
+            .iterate_all_keys_with_lengths()?
+            .map(|result| {
+                let (key, length) = result?;
+                let key_string = core::str::from_utf8(key.as_ref())?;
+                Ok((key_string.to_owned(), length))
+            })
+            .collect::<Result<Vec<_>>>()?;
+
+        let compressed_len = compress(b"A")?.len();
+        assert_eq!(compressed_len, 3);
+
+        let expected = [
+            ("A".to_owned(), compressed_len),
+            ("B".to_owned(), compressed_len),
+            ("C".to_owned(), compressed_len),
+            ("E".to_owned(), compressed_len),
+        ];
+
+        assert_eq!(values, expected);
+
+        Ok(())
+    }
+
     // This covers a bug we introduced and fixed while implementing in-memory mode.
     #[test_case(build_persistent_database)]
     #[test_case(build_in_memory_database)]
@@ -667,7 +753,7 @@ mod tests {
     }
 
     fn build_persistent_database() -> Result<Database> {
-        let database = Database::persistent("test_db", TempDir::new()?, ByteSize::mib(1))?;
+        let database = Database::persistent("test_db", TempDir::new()?, ByteSize::mib(1), false)?;
         populate_database(&database)?;
         Ok(database)
     }

fork_choice_control/src/lib.rs

@@ -25,6 +25,10 @@ pub use crate::{
     misc::{MutatorRejectionReason, VerifyAggregateAndProofResult, VerifyAttestationResult},
     queries::{BlockWithRoot, ForkChoiceContext, ForkTip, Snapshot},
     specialized::{AdHocBenchController, BenchController},
+    storage::{
+        BlobSidecarByBlobId, BlockCheckpoint, BlockRootBySlot, FinalizedBlockByRoot, PrefixableKey,
+        SlotBlobId, SlotByStateRoot, StateByBlockRoot, StateCheckpoint, UnfinalizedBlockByRoot,
+    },
     storage::{StateLoadStrategy, Storage, DEFAULT_ARCHIVAL_EPOCH_INTERVAL},
     storage_tool::{export_state_and_blocks, replay_blocks},
     wait::Wait,

fork_choice_control/src/storage.rs

@@ -769,16 +769,29 @@ type StateStorage<'storage, P> = (
     UnfinalizedBlocks<'storage, P>,
 );
 
+pub trait PrefixableKey {
+    const PREFIX: &'static str;
+
+    #[must_use]
+    fn has_prefix(bytes: &[u8]) -> bool {
+        bytes.starts_with(Self::PREFIX.as_bytes())
+    }
+}
+
 #[derive(Ssz)]
 // A `bound_for_read` attribute like this must be added when deriving `SszRead` for any type that
 // contains a block or state. The name of the `C` type parameter is hardcoded in `ssz_derive`.
 #[ssz(bound_for_read = "BeaconState<P>: SszRead<C>", derive_hash = false)]
-struct StateCheckpoint<P: Preset> {
+pub struct StateCheckpoint<P: Preset> {
     block_root: H256,
     head_slot: Slot,
     state: Arc<BeaconState<P>>,
 }
 
+impl<P: Preset> PrefixableKey for StateCheckpoint<P> {
+    const PREFIX: &'static str = Self::KEY;
+}
+
 impl<P: Preset> StateCheckpoint<P> {
     // This was renamed from `cstate` for compatibility with old schema versions.
     const KEY: &'static str = "cstate2";
@@ -792,10 +805,14 @@ impl<P: Preset> StateCheckpoint<P> {
     derive_hash = false,
     transparent
 )]
-struct BlockCheckpoint<P: Preset> {
+pub struct BlockCheckpoint<P: Preset> {
     block: Arc<SignedBeaconBlock<P>>,
 }
 
+impl<P: Preset> PrefixableKey for BlockCheckpoint<P> {
+    const PREFIX: &'static str = Self::KEY;
+}
+
 impl<P: Preset> BlockCheckpoint<P> {
     const KEY: &'static str = "cblock";
 }
@@ -822,56 +839,47 @@ impl TryFrom<Cow<'_, [u8]>> for BlockRootBySlot {
     }
 }
 
-impl BlockRootBySlot {
+impl PrefixableKey for BlockRootBySlot {
     const PREFIX: &'static str = "r";
-
-    fn has_prefix(bytes: &[u8]) -> bool {
-        bytes.starts_with(Self::PREFIX.as_bytes())
-    }
 }
 
 #[derive(Display)]
 #[display("{}{_0:x}", Self::PREFIX)]
 pub struct FinalizedBlockByRoot(pub H256);
 
-impl FinalizedBlockByRoot {
+impl PrefixableKey for FinalizedBlockByRoot {
     const PREFIX: &'static str = "b";
-
-    #[cfg(test)]
-    fn has_prefix(bytes: &[u8]) -> bool {
-        bytes.starts_with(Self::PREFIX.as_bytes())
-    }
 }
 
 #[derive(Display)]
 #[display("{}{_0:x}", Self::PREFIX)]
 pub struct UnfinalizedBlockByRoot(pub H256);
 
-impl UnfinalizedBlockByRoot {
+impl PrefixableKey for UnfinalizedBlockByRoot {
     const PREFIX: &'static str = "b_nf";
 }
 
 #[derive(Display)]
 #[display("{}{_0:x}", Self::PREFIX)]
 pub struct StateByBlockRoot(pub H256);
 
-impl StateByBlockRoot {
+impl PrefixableKey for StateByBlockRoot {
     const PREFIX: &'static str = "s";
 }
 
 #[derive(Display)]
 #[display("{}{_0:x}", Self::PREFIX)]
 pub struct SlotByStateRoot(pub H256);
 
-impl SlotByStateRoot {
+impl PrefixableKey for SlotByStateRoot {
     const PREFIX: &'static str = "t";
 }
 
 #[derive(Display)]
 #[display("{}{_0:x}{_1}", Self::PREFIX)]
 pub struct BlobSidecarByBlobId(pub H256, pub BlobIndex);
 
-impl BlobSidecarByBlobId {
+impl PrefixableKey for BlobSidecarByBlobId {
     const PREFIX: &'static str = "o";
 
     #[cfg(test)]
@@ -884,12 +892,8 @@ impl BlobSidecarByBlobId {
 #[display("{}{_0:020}{_1:x}{_2}", Self::PREFIX)]
 pub struct SlotBlobId(pub Slot, pub H256, pub BlobIndex);
 
-impl SlotBlobId {
+impl PrefixableKey for SlotBlobId {
     const PREFIX: &'static str = "i";
-
-    fn has_prefix(bytes: &[u8]) -> bool {
-        bytes.starts_with(Self::PREFIX.as_bytes())
-    }
 }
 
 #[derive(Debug, Error)]
@@ -930,7 +934,7 @@ mod tests {
     #[test]
     #[allow(clippy::similar_names)]
     fn test_prune_old_blob_sidecars() -> Result<()> {
-        let database = Database::persistent("test_db", TempDir::new()?, ByteSize::mib(10))?;
+        let database = Database::persistent("test_db", TempDir::new()?, ByteSize::mib(10), false)?;
 
         let storage = Storage::<Mainnet>::new(
             Arc::new(Config::mainnet()),

grandine/src/commands.rs

@@ -6,6 +6,14 @@ use types::phase0::primitives::Slot;
 #[derive(Clone, Subcommand)]
 #[cfg_attr(test, derive(PartialEq, Eq, Debug))]
 pub enum GrandineCommand {
+    /// Show `beacon_fork_choice` database element sizes
+    /// (example: grandine db-stats)
+    DbStats {
+        /// Custom database path
+        #[clap(short, long)]
+        path: Option<PathBuf>,
+    },
+
     /// Export blocks and state to ssz files within slot range for debugging
     /// (example: grandine export --from 0 --to 5)
     Export {