refactor(state): split the database module (#3568)

* refactor(state): split the disk_format module

* refactor(ci): add the new disk_db file to the state CI list

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

.github/workflows/test.yml

@@ -163,6 +163,7 @@ jobs:
         with:
           files: |
             /zebra-state/**/disk_format.rs
+            /zebra-state/**/disk_db.rs
             /zebra-state/**/finalized_state.rs
             /zebra-state/**/constants.rs
 

zebra-state/src/service/finalized_state.rs

@@ -1,10 +1,5 @@
 //! The primary implementation of the `zebra_state::Service` built upon rocksdb
 
-mod disk_format;
-
-#[cfg(test)]
-mod tests;
-
 use std::{
     borrow::Borrow,
     collections::HashMap,
@@ -26,11 +21,26 @@ use zebra_chain::{
     value_balance::ValueBalance,
 };
 
-use crate::{service::check, BoxError, Config, FinalizedBlock, HashOrHeight};
+use crate::{
+    service::{
+        check,
+        finalized_state::{
+            disk_db::{ReadDisk, WriteDisk},
+            disk_format::{FromDisk, IntoDisk, TransactionLocation},
+        },
+        QueuedFinalized,
+    },
+    BoxError, Config, FinalizedBlock, HashOrHeight,
+};
 
-use self::disk_format::{DiskDeserialize, DiskSerialize, FromDisk, IntoDisk, TransactionLocation};
+mod disk_db;
+mod disk_format;
 
-use super::QueuedFinalized;
+#[cfg(any(test, feature = "proptest-impl"))]
+mod arbitrary;
+
+#[cfg(test)]
+mod tests;
 
 /// The finalized part of the chain state, stored in the db.
 pub struct FinalizedState {

zebra-state/src/service/finalized_state/arbitrary.rs

@@ -0,0 +1,86 @@
+//! Arbitrary value generation and test harnesses for the finalized state.
+
+#![allow(dead_code)]
+
+use std::sync::Arc;
+
+use proptest::prelude::*;
+
+use zebra_chain::block;
+
+use crate::service::finalized_state::disk_format::{FromDisk, IntoDisk, TransactionLocation};
+
+impl Arbitrary for TransactionLocation {
+    type Parameters = ();
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        (any::<block::Height>(), any::<u32>())
+            .prop_map(|(height, index)| Self { height, index })
+            .boxed()
+    }
+
+    type Strategy = BoxedStrategy<Self>;
+}
+
+pub fn round_trip<T>(input: T) -> T
+where
+    T: IntoDisk + FromDisk,
+{
+    let bytes = input.as_bytes();
+    T::from_bytes(bytes)
+}
+
+pub fn assert_round_trip<T>(input: T)
+where
+    T: IntoDisk + FromDisk + Clone + PartialEq + std::fmt::Debug,
+{
+    let before = input.clone();
+    let after = round_trip(input);
+    assert_eq!(before, after);
+}
+
+pub fn round_trip_ref<T>(input: &T) -> T
+where
+    T: IntoDisk + FromDisk,
+{
+    let bytes = input.as_bytes();
+    T::from_bytes(bytes)
+}
+
+pub fn assert_round_trip_ref<T>(input: &T)
+where
+    T: IntoDisk + FromDisk + Clone + PartialEq + std::fmt::Debug,
+{
+    let before = input;
+    let after = round_trip_ref(input);
+    assert_eq!(before, &after);
+}
+
+pub fn round_trip_arc<T>(input: Arc<T>) -> T
+where
+    T: IntoDisk + FromDisk,
+{
+    let bytes = input.as_bytes();
+    T::from_bytes(bytes)
+}
+
+pub fn assert_round_trip_arc<T>(input: Arc<T>)
+where
+    T: IntoDisk + FromDisk + Clone + PartialEq + std::fmt::Debug,
+{
+    let before = input.clone();
+    let after = round_trip_arc(input);
+    assert_eq!(*before, after);
+}
+
+/// The round trip test covers types that are used as value field in a rocksdb
+/// column family. Only these types are ever deserialized, and so they're the only
+/// ones that implement both `IntoDisk` and `FromDisk`.
+pub fn assert_value_properties<T>(input: T)
+where
+    T: IntoDisk + FromDisk + Clone + PartialEq + std::fmt::Debug,
+{
+    assert_round_trip_ref(&input);
+    assert_round_trip_arc(Arc::new(input.clone()));
+    assert_round_trip(input);
+}

zebra-state/src/service/finalized_state/disk_db.rs

@@ -0,0 +1,92 @@
+//! Module defining access to RocksDB via accessor traits.
+//!
+//! This module makes sure that:
+//! - all disk writes happen inside a RocksDB transaction, and
+//! - format-specific invariants are maintained.
+
+use std::fmt::Debug;
+
+use crate::service::finalized_state::disk_format::{FromDisk, IntoDisk};
+
+/// Helper trait for inserting (Key, Value) pairs into rocksdb with a consistently
+/// defined format
+pub trait WriteDisk {
+    /// Serialize and insert the given key and value into a rocksdb column family,
+    /// overwriting any existing `value` for `key`.
+    fn zs_insert<K, V>(&mut self, cf: &rocksdb::ColumnFamily, key: K, value: V)
+    where
+        K: IntoDisk + Debug,
+        V: IntoDisk;
+
+    /// Remove the given key form rocksdb column family if it exists.
+    fn zs_delete<K>(&mut self, cf: &rocksdb::ColumnFamily, key: K)
+    where
+        K: IntoDisk + Debug;
+}
+
+impl WriteDisk for rocksdb::WriteBatch {
+    fn zs_insert<K, V>(&mut self, cf: &rocksdb::ColumnFamily, key: K, value: V)
+    where
+        K: IntoDisk + Debug,
+        V: IntoDisk,
+    {
+        let key_bytes = key.as_bytes();
+        let value_bytes = value.as_bytes();
+        self.put_cf(cf, key_bytes, value_bytes);
+    }
+
+    fn zs_delete<K>(&mut self, cf: &rocksdb::ColumnFamily, key: K)
+    where
+        K: IntoDisk + Debug,
+    {
+        let key_bytes = key.as_bytes();
+        self.delete_cf(cf, key_bytes);
+    }
+}
+
+/// Helper trait for retrieving values from rocksdb column familys with a consistently
+/// defined format
+pub trait ReadDisk {
+    /// Returns the value for `key` in the rocksdb column family `cf`, if present.
+    fn zs_get<K, V>(&self, cf: &rocksdb::ColumnFamily, key: &K) -> Option<V>
+    where
+        K: IntoDisk,
+        V: FromDisk;
+
+    /// Check if a rocksdb column family `cf` contains the serialized form of `key`.
+    fn zs_contains<K>(&self, cf: &rocksdb::ColumnFamily, key: &K) -> bool
+    where
+        K: IntoDisk;
+}
+
+impl ReadDisk for rocksdb::DB {
+    fn zs_get<K, V>(&self, cf: &rocksdb::ColumnFamily, key: &K) -> Option<V>
+    where
+        K: IntoDisk,
+        V: FromDisk,
+    {
+        let key_bytes = key.as_bytes();
+
+        // We use `get_pinned_cf` to avoid taking ownership of the serialized
+        // value, because we're going to deserialize it anyways, which avoids an
+        // extra copy
+        let value_bytes = self
+            .get_pinned_cf(cf, key_bytes)
+            .expect("expected that disk errors would not occur");
+
+        value_bytes.map(V::from_bytes)
+    }
+
+    fn zs_contains<K>(&self, cf: &rocksdb::ColumnFamily, key: &K) -> bool
+    where
+        K: IntoDisk,
+    {
+        let key_bytes = key.as_bytes();
+
+        // We use `get_pinned_cf` to avoid taking ownership of the serialized
+        // value, because we don't use the value at all. This avoids an extra copy.
+        self.get_pinned_cf(cf, key_bytes)
+            .expect("expected that disk errors would not occur")
+            .is_some()
+    }
+}

zebra-state/src/service/finalized_state/disk_format.rs

@@ -1,4 +1,5 @@
-//! Module defining exactly how to move types in and out of rocksdb
+//! Module defining the serialization format for finalized data.
+
 use std::{collections::BTreeMap, convert::TryInto, fmt::Debug, sync::Arc};
 
 use bincode::Options;
@@ -16,6 +17,9 @@ use zebra_chain::{
     value_balance::ValueBalance,
 };
 
+#[cfg(test)]
+mod tests;
+
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub struct TransactionLocation {
     pub height: block::Height,
@@ -378,206 +382,3 @@ impl FromDisk for NonEmptyHistoryTree {
         .expect("deserialization format should match the serialization format used by IntoDisk")
     }
 }
-
-/// Helper trait for inserting (Key, Value) pairs into rocksdb with a consistently
-/// defined format
-pub trait DiskSerialize {
-    /// Serialize and insert the given key and value into a rocksdb column family,
-    /// overwriting any existing `value` for `key`.
-    fn zs_insert<K, V>(&mut self, cf: &rocksdb::ColumnFamily, key: K, value: V)
-    where
-        K: IntoDisk + Debug,
-        V: IntoDisk;
-
-    /// Remove the given key form rocksdb column family if it exists.
-    fn zs_delete<K>(&mut self, cf: &rocksdb::ColumnFamily, key: K)
-    where
-        K: IntoDisk + Debug;
-}
-
-impl DiskSerialize for rocksdb::WriteBatch {
-    fn zs_insert<K, V>(&mut self, cf: &rocksdb::ColumnFamily, key: K, value: V)
-    where
-        K: IntoDisk + Debug,
-        V: IntoDisk,
-    {
-        let key_bytes = key.as_bytes();
-        let value_bytes = value.as_bytes();
-        self.put_cf(cf, key_bytes, value_bytes);
-    }
-
-    fn zs_delete<K>(&mut self, cf: &rocksdb::ColumnFamily, key: K)
-    where
-        K: IntoDisk + Debug,
-    {
-        let key_bytes = key.as_bytes();
-        self.delete_cf(cf, key_bytes);
-    }
-}
-
-/// Helper trait for retrieving values from rocksdb column familys with a consistently
-/// defined format
-pub trait DiskDeserialize {
-    /// Returns the value for `key` in the rocksdb column family `cf`, if present.
-    fn zs_get<K, V>(&self, cf: &rocksdb::ColumnFamily, key: &K) -> Option<V>
-    where
-        K: IntoDisk,
-        V: FromDisk;
-
-    /// Check if a rocksdb column family `cf` contains the serialized form of `key`.
-    fn zs_contains<K>(&self, cf: &rocksdb::ColumnFamily, key: &K) -> bool
-    where
-        K: IntoDisk;
-}
-
-impl DiskDeserialize for rocksdb::DB {
-    fn zs_get<K, V>(&self, cf: &rocksdb::ColumnFamily, key: &K) -> Option<V>
-    where
-        K: IntoDisk,
-        V: FromDisk,
-    {
-        let key_bytes = key.as_bytes();
-
-        // We use `get_pinned_cf` to avoid taking ownership of the serialized
-        // value, because we're going to deserialize it anyways, which avoids an
-        // extra copy
-        let value_bytes = self
-            .get_pinned_cf(cf, key_bytes)
-            .expect("expected that disk errors would not occur");
-
-        value_bytes.map(V::from_bytes)
-    }
-
-    fn zs_contains<K>(&self, cf: &rocksdb::ColumnFamily, key: &K) -> bool
-    where
-        K: IntoDisk,
-    {
-        let key_bytes = key.as_bytes();
-
-        // We use `get_pinned_cf` to avoid taking ownership of the serialized
-        // value, because we don't use the value at all. This avoids an extra copy.
-        self.get_pinned_cf(cf, key_bytes)
-            .expect("expected that disk errors would not occur")
-            .is_some()
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use proptest::{arbitrary::any, prelude::*};
-
-    impl Arbitrary for TransactionLocation {
-        type Parameters = ();
-
-        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
-            (any::<block::Height>(), any::<u32>())
-                .prop_map(|(height, index)| Self { height, index })
-                .boxed()
-        }
-
-        type Strategy = BoxedStrategy<Self>;
-    }
-
-    fn round_trip<T>(input: T) -> T
-    where
-        T: IntoDisk + FromDisk,
-    {
-        let bytes = input.as_bytes();
-        T::from_bytes(bytes)
-    }
-
-    fn assert_round_trip<T>(input: T)
-    where
-        T: IntoDisk + FromDisk + Clone + PartialEq + std::fmt::Debug,
-    {
-        let before = input.clone();
-        let after = round_trip(input);
-        assert_eq!(before, after);
-    }
-
-    fn round_trip_ref<T>(input: &T) -> T
-    where
-        T: IntoDisk + FromDisk,
-    {
-        let bytes = input.as_bytes();
-        T::from_bytes(bytes)
-    }
-
-    fn assert_round_trip_ref<T>(input: &T)
-    where
-        T: IntoDisk + FromDisk + Clone + PartialEq + std::fmt::Debug,
-    {
-        let before = input;
-        let after = round_trip_ref(input);
-        assert_eq!(before, &after);
-    }
-
-    fn round_trip_arc<T>(input: Arc<T>) -> T
-    where
-        T: IntoDisk + FromDisk,
-    {
-        let bytes = input.as_bytes();
-        T::from_bytes(bytes)
-    }
-
-    fn assert_round_trip_arc<T>(input: Arc<T>)
-    where
-        T: IntoDisk + FromDisk + Clone + PartialEq + std::fmt::Debug,
-    {
-        let before = input.clone();
-        let after = round_trip_arc(input);
-        assert_eq!(*before, after);
-    }
-
-    /// The round trip test covers types that are used as value field in a rocksdb
-    /// column family. Only these types are ever deserialized, and so they're the only
-    /// ones that implement both `IntoDisk` and `FromDisk`.
-    fn assert_value_properties<T>(input: T)
-    where
-        T: IntoDisk + FromDisk + Clone + PartialEq + std::fmt::Debug,
-    {
-        assert_round_trip_ref(&input);
-        assert_round_trip_arc(Arc::new(input.clone()));
-        assert_round_trip(input);
-    }
-
-    #[test]
-    fn roundtrip_transaction_location() {
-        zebra_test::init();
-        proptest!(|(val in any::<TransactionLocation>())| assert_value_properties(val));
-    }
-
-    #[test]
-    fn roundtrip_block_hash() {
-        zebra_test::init();
-        proptest!(|(val in any::<block::Hash>())| assert_value_properties(val));
-    }
-
-    #[test]
-    fn roundtrip_block_height() {
-        zebra_test::init();
-        proptest!(|(val in any::<block::Height>())| assert_value_properties(val));
-    }
-
-    #[test]
-    fn roundtrip_block() {
-        zebra_test::init();
-
-        proptest!(|(val in any::<Block>())| assert_value_properties(val));
-    }
-
-    #[test]
-    fn roundtrip_transparent_output() {
-        zebra_test::init();
-
-        proptest!(|(val in any::<transparent::Utxo>())| assert_value_properties(val));
-    }
-
-    #[test]
-    fn roundtrip_value_balance() {
-        zebra_test::init();
-
-        proptest!(|(val in any::<ValueBalance::<NonNegative>>())| assert_value_properties(val));
-    }
-}

zebra-state/src/service/finalized_state/disk_format/tests.rs

@@ -0,0 +1,3 @@
+//! Tests for the finalized disk format.
+
+mod prop;

zebra-state/src/service/finalized_state/disk_format/tests/prop.rs

@@ -0,0 +1,53 @@
+//! Randomised tests for the finalized disk format.
+
+use proptest::{arbitrary::any, prelude::*};
+
+use zebra_chain::{
+    amount::NonNegative,
+    block::{self, Block},
+    transparent,
+    value_balance::ValueBalance,
+};
+
+use crate::service::finalized_state::{
+    arbitrary::assert_value_properties, disk_format::TransactionLocation,
+};
+
+#[test]
+fn roundtrip_transaction_location() {
+    zebra_test::init();
+    proptest!(|(val in any::<TransactionLocation>())| assert_value_properties(val));
+}
+
+#[test]
+fn roundtrip_block_hash() {
+    zebra_test::init();
+    proptest!(|(val in any::<block::Hash>())| assert_value_properties(val));
+}
+
+#[test]
+fn roundtrip_block_height() {
+    zebra_test::init();
+    proptest!(|(val in any::<block::Height>())| assert_value_properties(val));
+}
+
+#[test]
+fn roundtrip_block() {
+    zebra_test::init();
+
+    proptest!(|(val in any::<Block>())| assert_value_properties(val));
+}
+
+#[test]
+fn roundtrip_transparent_output() {
+    zebra_test::init();
+
+    proptest!(|(val in any::<transparent::Utxo>())| assert_value_properties(val));
+}
+
+#[test]
+fn roundtrip_value_balance() {
+    zebra_test::init();
+
+    proptest!(|(val in any::<ValueBalance::<NonNegative>>())| assert_value_properties(val));
+}

zebra-state/src/service/finalized_state/tests.rs

@@ -1,2 +1,4 @@
+//! Finalized state tests.
+
 mod prop;
 mod vectors;

zebra-state/src/service/finalized_state/tests/prop.rs

@@ -1,3 +1,5 @@
+//! Randomised property tests for the finalized state.
+
 use std::env;
 
 use zebra_chain::{block::Height, parameters::NetworkUpgrade};

zebra-state/src/service/finalized_state/tests/vectors.rs

@@ -1,3 +1,5 @@
+//! Fixed test vectors for the finalized state.
+
 use halo2::arithmetic::FieldExt;
 use halo2::pasta::pallas;
 use hex::FromHex;