From 167a106055091c74bbcfc2876fa4277476ac28cc Mon Sep 17 00:00:00 2001 From: mrnaveira <47919901+mrnaveira@users.noreply.github.com> Date: Tue, 14 Jan 2025 22:17:37 +0000 Subject: [PATCH] refactor generic put_cf method --- .../state_store_rocksdb/src/model/block.rs | 34 +++++++++---------- .../state_store_rocksdb/src/model/encoding.rs | 2 +- .../state_store_rocksdb/src/model/model.rs | 10 +++--- .../src/model/state_transition.rs | 14 +++++--- .../src/model/state_tree_shard_versions.rs | 10 ------ .../state_store_rocksdb/src/model/substate.rs | 20 +++++++---- 6 files changed, 45 insertions(+), 45 deletions(-) diff --git a/dan_layer/state_store_rocksdb/src/model/block.rs b/dan_layer/state_store_rocksdb/src/model/block.rs index f66473a0e..ae1898325 100644 --- a/dan_layer/state_store_rocksdb/src/model/block.rs +++ b/dan_layer/state_store_rocksdb/src/model/block.rs @@ -101,23 +101,6 @@ impl RocksdbModel for BlockModel { Self::key_from_block_id(item.id()) } - fn column_families() -> Vec<&'static str> { - vec![ParentIdColumnFamily::name(), EpochHeightColumnFamily::name(), IsCommittedColumnFamily::name()] - } - - fn put_cf(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, cf_name: &str, value: &Self::Item) -> Result<(), RocksDbStorageError> { - // In each CF value We store the BlockId of the block, so we can reference it back - let cf_value = value.id().as_bytes(); - match cf_name { - ParentIdColumnFamily::NAME => ParentIdColumnFamily::put(db, tx, operation, value, cf_value)?, - EpochHeightColumnFamily::NAME => EpochHeightColumnFamily::put(db, tx, operation, value, cf_value)?, - IsCommittedColumnFamily::NAME => IsCommittedColumnFamily::put(db, tx, operation, value, cf_value)?, - _ => (), - } - - Ok(()) - } - fn get_cf(db: Arc<TransactionDB>, tx: &Transaction<'_, TransactionDB>, cf_name: &str, operation: &'static str, key_prefix: &str, ordering: Ordering) -> Result<Option<Self::Item>, RocksDbStorageError> { let cf = db.cf_handle(cf_name).unwrap(); @@ -151,10 +134,25 @@ impl RocksdbModel for BlockModel { Ok(Some(value)) } + fn column_families() -> Vec<&'static str> { + vec![ParentIdColumnFamily::name(), EpochHeightColumnFamily::name(), IsCommittedColumnFamily::name()] + } + + fn put_in_cfs(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, value: &Self::Item) -> Result<(), RocksDbStorageError> { + // In each CF value We store the BlockId of the block, so we can reference it back + let cf_value = value.id().as_bytes(); + + ParentIdColumnFamily::put(db.clone(), tx, operation, value, cf_value)?; + EpochHeightColumnFamily::put(db.clone(), tx, operation, value, cf_value)?; + IsCommittedColumnFamily::put(db, tx, operation, value, cf_value)?; + + Ok(()) + } + fn delete_from_cfs(db: Arc<TransactionDB>, tx: &Transaction<'_, TransactionDB>, operation: &'static str, item: &Self::Item) -> Result<(), RocksDbStorageError> { ParentIdColumnFamily::delete(db.clone(), tx, operation, item)?; EpochHeightColumnFamily::delete(db.clone(), tx, operation, item)?; - IsCommittedColumnFamily::delete(db.clone(), tx, operation, item)?; + IsCommittedColumnFamily::delete(db, tx, operation, item)?; Ok(()) } } diff --git a/dan_layer/state_store_rocksdb/src/model/encoding.rs b/dan_layer/state_store_rocksdb/src/model/encoding.rs index 30f7e5499..c7a9eabf2 100644 --- a/dan_layer/state_store_rocksdb/src/model/encoding.rs +++ b/dan_layer/state_store_rocksdb/src/model/encoding.rs @@ -20,7 +20,7 @@ // WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE // USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -use serde::{de::DeserializeOwned, Deserialize, Serialize}; +use serde::{de::DeserializeOwned, Serialize}; use crate::error::RocksDbStorageError; diff --git a/dan_layer/state_store_rocksdb/src/model/model.rs b/dan_layer/state_store_rocksdb/src/model/model.rs index cd918bc89..201f7ca8f 100644 --- a/dan_layer/state_store_rocksdb/src/model/model.rs +++ b/dan_layer/state_store_rocksdb/src/model/model.rs @@ -93,14 +93,16 @@ pub trait RocksdbModel { })?; // insert the main key in each of the column families - for cf_name in Self::column_families() { - Self::put_cf(db.clone(), tx, operation, &cf_name, &value)?; - } + Self::put_in_cfs(db.clone(), tx, operation, &value)?; Ok(()) } - fn put_cf(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, cf_name: &str, value: &Self::Item) -> Result<(), RocksDbStorageError>; + fn put_in_cfs(_db: Arc<TransactionDB>, _tx: &mut Transaction<'_, TransactionDB>, _operation: &'static str, _value: &Self::Item) -> Result<(), RocksDbStorageError> { + // It's up to concrete models to override this method + // We provide a default implementation to simplify all the models that do not have column families + Ok(()) + } fn key_exists(tx: &Transaction<'_, TransactionDB>, operation: &'static str, key: &str) -> Result<bool, RocksDbStorageError> { let value = tx.get(&key) diff --git a/dan_layer/state_store_rocksdb/src/model/state_transition.rs b/dan_layer/state_store_rocksdb/src/model/state_transition.rs index 329d629c2..78001f157 100644 --- a/dan_layer/state_store_rocksdb/src/model/state_transition.rs +++ b/dan_layer/state_store_rocksdb/src/model/state_transition.rs @@ -72,15 +72,19 @@ impl RocksdbModel for StateTransitionModel { vec![ShardColumnFamily::name()] } - fn put_cf(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, cf_name: &str, value: &Self::Item) -> Result<(), RocksDbStorageError> { + fn put_in_cfs(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, value: &Self::Item) -> Result<(), RocksDbStorageError> { // In each CF value We store the key to the main collection, so we can retrieve the actual value let main_key = Self::key(value); let main_key_bytes = main_key.as_bytes(); - match cf_name { - ShardColumnFamily::NAME => ShardColumnFamily::put(db, tx, operation, value, main_key_bytes)?, - _ => (), - } + ShardColumnFamily::put(db, tx, operation, value, main_key_bytes)?; + + Ok(()) + } + + fn delete_from_cfs(db: Arc<TransactionDB>, tx: &Transaction<'_, TransactionDB>, operation: &'static str, item: &Self::Item) -> Result<(), RocksDbStorageError> { + ShardColumnFamily::delete(db.clone(), tx, operation, item)?; + Ok(()) } } diff --git a/dan_layer/state_store_rocksdb/src/model/state_tree_shard_versions.rs b/dan_layer/state_store_rocksdb/src/model/state_tree_shard_versions.rs index ed18d0c45..c691043f3 100644 --- a/dan_layer/state_store_rocksdb/src/model/state_tree_shard_versions.rs +++ b/dan_layer/state_store_rocksdb/src/model/state_tree_shard_versions.rs @@ -20,15 +20,10 @@ // WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE // USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -use std::sync::Arc; - -use rocksdb::{Transaction, TransactionDB}; use serde::{Deserialize, Serialize}; use tari_dan_common_types::shard::Shard; use tari_state_tree::Version; -use crate::error::RocksDbStorageError; - use super::model::RocksdbModel; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -59,9 +54,4 @@ impl RocksdbModel for StateTreeShardVersionModel { fn column_families() -> Vec<&'static str> { vec![] } - - fn put_cf(_db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, _operation: &'static str, _cf_name: &str, _value: &Self::Item) -> Result<(), RocksDbStorageError> { - // No column families for this model - Ok(()) - } } \ No newline at end of file diff --git a/dan_layer/state_store_rocksdb/src/model/substate.rs b/dan_layer/state_store_rocksdb/src/model/substate.rs index d489b5d17..663865c0c 100644 --- a/dan_layer/state_store_rocksdb/src/model/substate.rs +++ b/dan_layer/state_store_rocksdb/src/model/substate.rs @@ -133,17 +133,23 @@ impl RocksdbModel for SubstateModel { vec![VersionColumnFamily::name(), CreatedByTxColumnFamily::name(), DestroyedByTxColumnFamily::name()] } - fn put_cf(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, cf_name: &str, value: &Self::Item) -> Result<(), RocksDbStorageError> { + fn put_in_cfs(db: Arc<TransactionDB>, tx: &mut Transaction<'_, TransactionDB>, operation: &'static str, value: &Self::Item) -> Result<(), RocksDbStorageError> { // In each CF value We store the key to the main collection, so we can retrieve the actual value let main_key = Self::key(value); let main_key_bytes = main_key.as_bytes(); - match cf_name { - VersionColumnFamily::NAME => VersionColumnFamily::put(db, tx, operation, value, main_key_bytes)?, - CreatedByTxColumnFamily::NAME => CreatedByTxColumnFamily::put(db, tx, operation, value, main_key_bytes)?, - DestroyedByTxColumnFamily::NAME => DestroyedByTxColumnFamily::put(db, tx, operation, value, main_key_bytes)?, - _ => (), - } + VersionColumnFamily::put(db.clone(), tx, operation, value, main_key_bytes)?; + CreatedByTxColumnFamily::put(db.clone(), tx, operation, value, main_key_bytes)?; + DestroyedByTxColumnFamily::put(db, tx, operation, value, main_key_bytes)?; + + Ok(()) + } + + fn delete_from_cfs(db: Arc<TransactionDB>, tx: &Transaction<'_, TransactionDB>, operation: &'static str, item: &Self::Item) -> Result<(), RocksDbStorageError> { + VersionColumnFamily::delete(db.clone(), tx, operation, item)?; + CreatedByTxColumnFamily::delete(db.clone(), tx, operation, item)?; + DestroyedByTxColumnFamily::delete(db, tx, operation, item)?; + Ok(()) } }