Write the dag relationship in batch

flexidag/src/consensusdb/consensus_relations.rs

@@ -1,14 +1,15 @@
 use super::schema::{KeyCodec, ValueCodec};
 use super::{
     db::DBStorage,
-    prelude::{BatchDbWriter, CachedDbAccess, DirectDbWriter, StoreError},
+    prelude::{CachedDbAccess, StoreError},
 };
 use crate::define_schema;
-use rocksdb::WriteBatch;
 use starcoin_crypto::HashValue as Hash;
+use starcoin_storage::batch::{WriteBatch, WriteBatchData, WriteBatchWithColumn};
+use starcoin_storage::storage::{InnerStore, WriteOp};
 use starcoin_types::blockhash::{BlockHashes, BlockLevel};
+use std::collections::HashMap;
 use std::sync::Arc;
-
 /// Reader API for `RelationsStore`.
 pub trait RelationsStoreReader {
     fn get_parents(&self, hash: Hash) -> Result<BlockHashes, StoreError>;
@@ -90,41 +91,6 @@ impl DbRelationsStore {
     pub fn clone_with_new_cache(&self, cache_size: usize) -> Self {
         Self::new(Arc::clone(&self.db), self.level, cache_size)
     }
-
-    pub fn insert_batch(
-        &mut self,
-        batch: &mut WriteBatch,
-        hash: Hash,
-        parents: BlockHashes,
-    ) -> Result<(), StoreError> {
-        if self.has(hash)? {
-            return Err(StoreError::KeyAlreadyExists(hash.to_string()));
-        }
-
-        // Insert a new entry for `hash`
-        self.parents_access
-            .write(BatchDbWriter::new(batch, &self.db), hash, parents.clone())?;
-
-        // The new hash has no children yet
-        self.children_access.write(
-            BatchDbWriter::new(batch, &self.db),
-            hash,
-            BlockHashes::new(Vec::new()),
-        )?;
-
-        // Update `children` for each parent
-        for parent in parents.iter().cloned() {
-            let mut children = (*self.get_children(parent)?).clone();
-            children.push(hash);
-            self.children_access.write(
-                BatchDbWriter::new(batch, &self.db),
-                parent,
-                BlockHashes::new(children),
-            )?;
-        }
-
-        Ok(())
-    }
 }
 
 impl RelationsStoreReader for DbRelationsStore {
@@ -147,35 +113,53 @@ impl RelationsStoreReader for DbRelationsStore {
 }
 
 impl RelationsStore for DbRelationsStore {
-    /// See `insert_batch` as well
-    /// TODO: use one function with DbWriter for both this function and insert_batch
     fn insert(&self, hash: Hash, parents: BlockHashes) -> Result<(), StoreError> {
         if self.has(hash)? {
             return Err(StoreError::KeyAlreadyExists(hash.to_string()));
         }
 
-        // Insert a new entry for `hash`
-        self.parents_access
-            .write(DirectDbWriter::new(&self.db), hash, parents.clone())?;
-
-        // The new hash has no children yet
-        self.children_access.write(
-            DirectDbWriter::new(&self.db),
-            hash,
-            BlockHashes::new(Vec::new()),
-        )?;
-
-        // Update `children` for each parent
+        let mut parent_to_children = HashMap::new();
         for parent in parents.iter().cloned() {
-            let mut children = (*self.get_children(parent)?).clone();
+            let mut children = match self.get_children(parent) {
+                Ok(children) => (*children).clone(),
+                Err(e) => match e {
+                    StoreError::KeyNotFound(_) => vec![],
+                    _ => return std::result::Result::Err(e),
+                },
+            };
             children.push(hash);
-            self.children_access.write(
-                DirectDbWriter::new(&self.db),
-                parent,
-                BlockHashes::new(children),
-            )?;
+            parent_to_children.insert(
+                parent.to_vec(),
+                bcs_ext::to_bytes(&children).map_err(|e| StoreError::EncodeError(e.to_string()))?,
+            );
         }
 
+        let batch = WriteBatchWithColumn {
+            data: vec![
+                WriteBatchData {
+                    column: PARENTS_CF.to_string(),
+                    row_data: WriteBatch::new_with_rows(vec![(
+                        hash.to_vec(),
+                        WriteOp::Value(
+                            bcs_ext::to_bytes(&parents)
+                                .map_err(|e| StoreError::EncodeError(e.to_string()))?,
+                        ),
+                    )]),
+                },
+                WriteBatchData {
+                    column: CHILDREN_CF.to_string(),
+                    row_data: WriteBatch::new_with_rows(
+                        parent_to_children
+                            .into_iter()
+                            .map(|(key, value)| (key, WriteOp::Value(value)))
+                            .collect(),
+                    ),
+                },
+            ],
+        };
+        self.db
+            .write_batch_with_column(batch)
+            .map_err(|e| StoreError::DBIoError(e.to_string()))?;
         Ok(())
     }
 }

storage/src/batch/mod.rs

@@ -56,3 +56,14 @@ where
         Ok(Self::new_with_rows(rows?))
     }
 }
+
+#[derive(Debug, Default, Clone)]
+pub struct WriteBatchData {
+    pub column: String,
+    pub row_data: WriteBatch,
+}
+
+#[derive(Debug, Default, Clone)]
+pub struct WriteBatchWithColumn {
+    pub data: Vec<WriteBatchData>,
+}

storage/src/cache_storage/mod.rs

@@ -1,7 +1,7 @@
 // Copyright (c) The Starcoin Core Contributors
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::batch::GWriteBatch;
+use crate::batch::{GWriteBatch, WriteBatchWithColumn};
 use crate::{
     batch::WriteBatch,
     metrics::{record_metrics, StorageMetrics},
@@ -91,6 +91,32 @@ impl InnerStore for CacheStorage {
         })
     }
 
+    fn write_batch_with_column(&self, batch: WriteBatchWithColumn) -> Result<()> {
+        let rows = batch
+            .data
+            .into_iter()
+            .flat_map(|data| {
+                data.row_data
+                    .rows
+                    .iter()
+                    .cloned()
+                    .map(|(k, v)| (compose_key(Some(&data.column), k), v))
+                    .collect::<Vec<_>>()
+            })
+            .collect();
+        let batch = WriteBatch { rows };
+        record_metrics(
+            "cache",
+            "write_batch_column_prefix",
+            "write_batch",
+            self.metrics.as_ref(),
+        )
+        .call(|| {
+            self.write_batch_inner(batch);
+            Ok(())
+        })
+    }
+
     fn get_len(&self) -> Result<u64, Error> {
         Ok(self.cache.lock().len() as u64)
     }
@@ -111,6 +137,10 @@ impl InnerStore for CacheStorage {
         self.write_batch(prefix_name, batch)
     }
 
+    fn write_batch_with_column_sync(&self, batch: WriteBatchWithColumn) -> Result<()> {
+        self.write_batch_with_column(batch)
+    }
+
     fn multi_get(&self, prefix_name: &str, keys: Vec<Vec<u8>>) -> Result<Vec<Option<Vec<u8>>>> {
         let composed_keys = keys
             .into_iter()

storage/src/db_storage/mod.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    batch::WriteBatch,
+    batch::{WriteBatch, WriteBatchWithColumn},
     errors::StorageInitError,
     metrics::{record_metrics, StorageMetrics},
     storage::{ColumnFamilyName, InnerStore, KeyCodec, RawDBStorage, ValueCodec, WriteOp},
@@ -414,6 +414,53 @@ impl InnerStore for DBStorage {
         })
     }
 
+    fn write_batch_with_column(&self, batch: WriteBatchWithColumn) -> Result<()> {
+        let mut db_batch = DBWriteBatch::default();
+        batch.data.into_iter().for_each(|data| {
+            let cf_handle = self.get_cf_handle(&data.column);
+            for (key, write_op) in data.row_data.rows {
+                match write_op {
+                    WriteOp::Value(value) => db_batch.put_cf(cf_handle, key, value),
+                    WriteOp::Deletion => db_batch.delete_cf(cf_handle, key),
+                };
+            }
+        });
+        record_metrics(
+            "db",
+            "write_batch_column",
+            "write_batch",
+            self.metrics.as_ref(),
+        )
+        .call(|| {
+            self.db
+                .write_opt(db_batch, &Self::default_write_options())?;
+            Ok(())
+        })
+    }
+
+    fn write_batch_with_column_sync(&self, batch: WriteBatchWithColumn) -> Result<()> {
+        let mut db_batch = DBWriteBatch::default();
+        batch.data.into_iter().for_each(|data| {
+            let cf_handle = self.get_cf_handle(&data.column);
+            for (key, write_op) in data.row_data.rows {
+                match write_op {
+                    WriteOp::Value(value) => db_batch.put_cf(cf_handle, key, value),
+                    WriteOp::Deletion => db_batch.delete_cf(cf_handle, key),
+                };
+            }
+        });
+        record_metrics(
+            "db",
+            "write_batch_column",
+            "write_batch",
+            self.metrics.as_ref(),
+        )
+        .call(|| {
+            self.db.write_opt(db_batch, &Self::sync_write_options())?;
+            Ok(())
+        })
+    }
+
     fn get_len(&self) -> Result<u64> {
         unimplemented!()
     }

storage/src/storage.rs

@@ -3,6 +3,7 @@
 
 pub use crate::batch::WriteBatch;
 use crate::{
+    batch::WriteBatchWithColumn,
     cache_storage::CacheStorage,
     db_storage::{DBStorage, SchemaIterator},
     upgrade::DBUpgrade,
@@ -39,10 +40,12 @@ pub trait InnerStore: Send + Sync {
     fn contains_key(&self, prefix_name: &str, key: Vec<u8>) -> Result<bool>;
     fn remove(&self, prefix_name: &str, key: Vec<u8>) -> Result<()>;
     fn write_batch(&self, prefix_name: &str, batch: WriteBatch) -> Result<()>;
+    fn write_batch_with_column(&self, batch: WriteBatchWithColumn) -> Result<()>;
     fn get_len(&self) -> Result<u64>;
     fn keys(&self) -> Result<Vec<Vec<u8>>>;
     fn put_sync(&self, prefix_name: &str, key: Vec<u8>, value: Vec<u8>) -> Result<()>;
     fn write_batch_sync(&self, prefix_name: &str, batch: WriteBatch) -> Result<()>;
+    fn write_batch_with_column_sync(&self, batch: WriteBatchWithColumn) -> Result<()>;
     fn multi_get(&self, prefix_name: &str, keys: Vec<Vec<u8>>) -> Result<Vec<Option<Vec<u8>>>>;
 }
 
@@ -201,6 +204,18 @@ impl InnerStore for StorageInstance {
             },
         }
     }
+
+    fn write_batch_with_column(&self, batch: WriteBatchWithColumn) -> Result<()> {
+        match self {
+            Self::CACHE { cache } => cache.write_batch_with_column(batch),
+            Self::DB { db } => db.write_batch_with_column(batch),
+            Self::CacheAndDb { cache, db } => {
+                db.write_batch_with_column(batch.clone())?;
+                cache.write_batch_with_column(batch)
+            }
+        }
+    }
+
     fn get_len(&self) -> Result<u64> {
         match self {
             Self::CACHE { cache } => cache.get_len(),
@@ -240,6 +255,17 @@ impl InnerStore for StorageInstance {
         }
     }
 
+    fn write_batch_with_column_sync(&self, batch: WriteBatchWithColumn) -> Result<()> {
+        match self {
+            Self::CACHE { cache } => cache.write_batch_with_column_sync(batch),
+            Self::DB { db } => db.write_batch_with_column_sync(batch),
+            Self::CacheAndDb { cache, db } => {
+                db.write_batch_with_column_sync(batch.clone())?;
+                cache.write_batch_with_column_sync(batch)
+            }
+        }
+    }
+
     fn multi_get(&self, prefix_name: &str, keys: Vec<Vec<u8>>) -> Result<Vec<Option<Vec<u8>>>> {
         match self {
             Self::CACHE { cache } => cache.multi_get(prefix_name, keys),