header_map: change persist backend to rocksdb

Author: eval-exec

db/src/db.rs

@@ -8,7 +8,7 @@ use ckb_db_schema::Col;
 use ckb_logger::info;
 use rocksdb::ops::{
     CompactRangeCF, CreateCF, DropCF, GetColumnFamilys, GetPinned, GetPinnedCF, IterateCF, OpenCF,
-    Put, SetOptions, WriteOps,
+    Put, PutCF, SetOptions, WriteOps,
 };
 use rocksdb::{
     BlockBasedIndexType, BlockBasedOptions, Cache, ColumnFamily, ColumnFamilyDescriptor,
@@ -179,6 +179,16 @@ impl RocksDB {
         self.inner.get_pinned(key).map_err(internal_error)
     }
 
+    /// Put a valueinto the given column under the given key.
+    pub fn put<K, V>(&self, col: Col, key: K, value: V) -> Result<()>
+    where
+        K: AsRef<[u8]>,
+        V: AsRef<[u8]>,
+    {
+        let cf = cf_handle(&self.inner, col)?;
+        self.inner.put_cf(cf, key, value).map_err(internal_error)
+    }
+
     /// Insert a value into the database under the given key.
     pub fn put_default<K, V>(&self, key: K, value: V) -> Result<()>
     where

shared/src/shared_builder.rs

@@ -364,6 +364,7 @@ impl SharedBuilder {
         let ibd_finished = Arc::new(AtomicBool::new(false));
 
         let header_map = Arc::new(HeaderMap::new(
+            db.clone(),
             header_map_memory_limit,
             &async_handle,
             Arc::clone(&ibd_finished),

shared/src/types/header_map/kernel_lru.rs

@@ -1,20 +1,24 @@
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 
+use ckb_db::RocksDB;
+use ckb_db_schema::COLUMN_BLOCK_HEADER;
 #[cfg(feature = "stats")]
 use ckb_logger::info;
 use ckb_metrics::HistogramTimer;
+use ckb_types::packed;
 #[cfg(feature = "stats")]
 use ckb_util::{Mutex, MutexGuard};
 use ckb_util::{RwLock, RwLockReadGuard};
 
-use ckb_types::{U256, core::EpochNumberWithFraction, packed::Byte32};
+use ckb_types::{U256, core::EpochNumberWithFraction, packed::Byte32, prelude::*};
 
 use super::MemoryMap;
 use crate::types::HeaderIndexView;
 
 pub(crate) struct HeaderMapKernel {
     pub(crate) memory: MemoryMap,
+    db: RocksDB,
     // Configuration
     memory_limit: usize,
     // if ckb is in IBD mode, don't shrink memory map
@@ -39,14 +43,15 @@ struct HeaderMapKernelStats {
 }
 
 impl HeaderMapKernel {
-    pub(crate) fn new(memory_limit: usize, ibd_finished: Arc<AtomicBool>) -> Self {
+    pub(crate) fn new(db: RocksDB, memory_limit: usize, ibd_finished: Arc<AtomicBool>) -> Self {
         let memory = Default::default();
         let shared_best_header_value = Self::default_shared_best_header();
         let shared_best_header = RwLock::new(shared_best_header_value);
 
         #[cfg(not(feature = "stats"))]
         {
             Self {
+                db,
                 memory,
                 memory_limit,
                 ibd_finished,
@@ -57,6 +62,7 @@ impl HeaderMapKernel {
         #[cfg(feature = "stats")]
         {
             Self {
+                db,
                 memory,
                 memory_limit,
                 ibd_finished,
@@ -80,7 +86,13 @@ impl HeaderMapKernel {
         if let Some(metrics) = ckb_metrics::handle() {
             metrics.ckb_header_map_memory_hit_miss_count.miss.inc();
         }
-        false
+        let contains = self
+            .db
+            .get_pinned(COLUMN_BLOCK_HEADER, hash.as_slice())
+            .unwrap()
+            .is_some();
+
+        contains
     }
 
     pub(crate) fn get(&self, hash: &Byte32) -> Option<HeaderIndexView> {
@@ -99,27 +111,46 @@ impl HeaderMapKernel {
             metrics.ckb_header_map_memory_hit_miss_count.miss.inc();
         }
 
+        self.db
+            .get_pinned(COLUMN_BLOCK_HEADER, hash.as_slice())
+            .unwrap()
+            .and_then(|slice| {
+                let reader = packed::HeaderViewReader::from_slice_should_be_ok(slice.as_ref());
+
+                let header_view = Unpack::<ckb_types::core::HeaderView>::unpack(&reader);
+                Some(header_view)
+            });
+
         None
     }
 
-    pub(crate) fn insert(&self, view: HeaderIndexView) -> Option<()> {
+    pub(crate) fn insert(
+        &self,
+        view: ckb_types::core::HeaderView,
+        total_difficulty: U256,
+    ) -> Option<()> {
         #[cfg(feature = "stats")]
         {
             self.trace();
             self.stats().tick_primary_insert();
         }
-        self.memory.insert(view)
+        let packed_header: packed::HeaderView = view.clone().into();
+        let view: HeaderIndexView = (view, total_difficulty).into();
+        let hash = view.hash();
+        self.memory.insert(view.into());
+
+        self.db
+            .put(
+                COLUMN_BLOCK_HEADER,
+                hash.as_slice(),
+                packed_header.as_slice(),
+            )
+            .ok()
     }
 
-    pub(crate) fn remove(&self, hash: &Byte32) {
-        #[cfg(feature = "stats")]
-        {
-            self.trace();
-            self.stats().tick_primary_delete();
-        }
-        // If IBD is not finished, don't shrink memory map
-        let allow_shrink_to_fit = self.ibd_finished.load(Ordering::Acquire);
-        self.memory.remove(hash, allow_shrink_to_fit);
+    pub(crate) fn remove(&self, _hash: &Byte32) {
+        // TODO
+        // no need to remove
     }
 
     pub(crate) fn limit_memory(&self) {

shared/src/types/header_map/memory.rs

@@ -104,22 +104,6 @@ impl MemoryMap {
         ret.map(|_| ())
     }
 
-    pub(crate) fn remove(&self, key: &Byte32, shrink_to_fit: bool) -> Option<HeaderIndexView> {
-        let mut guard = self.0.write();
-        let ret = guard.remove(key);
-
-        if shrink_to_fit {
-            shrink_to_fit!(guard, SHRINK_THRESHOLD);
-        }
-        ret.map(|inner| {
-            if let Some(metrics) = ckb_metrics::handle() {
-                metrics.ckb_header_map_memory_count.dec();
-            }
-
-            (key.clone(), inner).into()
-        })
-    }
-
     pub(crate) fn excess_items(&self, size_limit: usize) -> Option<Vec<HeaderIndexView>> {
         let guard = self.0.read();
         let size = guard.len();

shared/src/types/header_map/mod.rs

@@ -1,7 +1,8 @@
 use ckb_async_runtime::Handle;
+use ckb_db::RocksDB;
 use ckb_logger::info;
 use ckb_stop_handler::{CancellationToken, new_tokio_exit_rx};
-use ckb_types::packed::Byte32;
+use ckb_types::{U256, packed::Byte32};
 use std::mem::size_of;
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
@@ -27,7 +28,12 @@ const ITEM_BYTES_SIZE: usize = size_of::<HeaderIndexView>();
 const WARN_THRESHOLD: usize = ITEM_BYTES_SIZE * 100_000;
 
 impl HeaderMap {
-    pub fn new(memory_limit: usize, async_handle: &Handle, ibd_finished: Arc<AtomicBool>) -> Self {
+    pub fn new(
+        db: RocksDB,
+        memory_limit: usize,
+        async_handle: &Handle,
+        ibd_finished: Arc<AtomicBool>,
+    ) -> Self {
         if memory_limit < ITEM_BYTES_SIZE {
             panic!("The limit setting is too low");
         }
@@ -38,7 +44,7 @@ impl HeaderMap {
             );
         }
         let size_limit = memory_limit / ITEM_BYTES_SIZE;
-        let inner = Arc::new(HeaderMapKernel::new(size_limit, ibd_finished));
+        let inner = Arc::new(HeaderMapKernel::new(db, size_limit, ibd_finished));
         let map = Arc::clone(&inner);
         let stop_rx: CancellationToken = new_tokio_exit_rx();
 
@@ -82,15 +88,15 @@ impl HeaderMap {
         self.inner.get(hash)
     }
 
-    pub fn insert(&self, view: HeaderIndexView) -> Option<()> {
+    pub fn insert(&self, view: ckb_types::core::HeaderView, total_difficulty: U256) -> Option<()> {
         let _trace_timer: Option<HistogramTimer> = ckb_metrics::handle().map(|metric| {
             metric
                 .ckb_header_map_ops_duration
                 .with_label_values(&["insert"])
                 .start_timer()
         });
 
-        self.inner.insert(view)
+        self.inner.insert(view, total_difficulty)
     }
 
     pub fn remove(&self, hash: &Byte32) {

sync/src/types/mod.rs

@@ -1131,13 +1131,14 @@ impl SyncShared {
         let parent_header_index = self
             .get_header_index_view(&parent_hash, store_first)
             .expect("parent should be verified");
+        let total_difficulty = parent_header_index.total_difficulty() + header.difficulty();
         let mut header_view = HeaderIndexView::new(
             header.hash(),
             header.number(),
             header.epoch(),
             header.timestamp(),
             parent_hash,
-            parent_header_index.total_difficulty() + header.difficulty(),
+            total_difficulty.clone(),
         );
 
         let snapshot = Arc::clone(&self.shared.snapshot());
@@ -1156,7 +1157,9 @@ impl SyncShared {
                 }
             },
         );
-        self.shared.header_map().insert(header_view.clone());
+        self.shared
+            .header_map()
+            .insert(header.to_owned(), total_difficulty);
         self.state
             .peers()
             .may_set_best_known_header(peer, header_view.as_header_index());