feat: RocksDB storage

src/core/src/index/revindex/disk_revindex.rs

@@ -1,5 +1,5 @@
 use std::hash::{BuildHasher, BuildHasherDefault, Hash, Hasher};
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 
@@ -19,7 +19,7 @@ use crate::manifest::Manifest;
 use crate::prelude::*;
 use crate::sketch::minhash::{KmerMinHash, KmerMinHashBTree};
 use crate::sketch::Sketch;
-use crate::storage::{InnerStorage, Storage};
+use crate::storage::{InnerStorage, Storage, STORAGE};
 use crate::Result;
 
 const DB_VERSION: u8 = 1;
@@ -35,6 +35,7 @@ fn compute_color(idxs: &Datasets) -> Color {
 pub struct RevIndex {
     db: Arc<DB>,
     collection: Arc<CollectionSet>,
+    path: PathBuf,
 }
 
 fn merge_datasets(
@@ -79,6 +80,7 @@ impl RevIndex {
         let index = Self {
             db,
             collection: Arc::new(collection),
+            path: path.into(),
         };
 
         index.collection.par_iter().for_each(|(dataset_id, _)| {
@@ -128,7 +130,11 @@ impl RevIndex {
             storage_spec,
         )?);
 
-        Ok(module::RevIndex::Plain(Self { db, collection }))
+        Ok(module::RevIndex::Plain(Self {
+            db,
+            collection,
+            path: path.as_ref().into(),
+        }))
     }
 
     fn load_collection_from_rocksdb(
@@ -498,7 +504,7 @@ impl RevIndexOps for RevIndex {
     }
 }
 
-fn cf_descriptors() -> Vec<ColumnFamilyDescriptor> {
+pub(crate) fn cf_descriptors() -> Vec<ColumnFamilyDescriptor> {
     let mut cfopts = Options::default();
     cfopts.set_max_write_buffer_number(16);
     cfopts.set_merge_operator_associative("datasets operator", merge_datasets);
@@ -514,15 +520,20 @@ fn cf_descriptors() -> Vec<ColumnFamilyDescriptor> {
     cfopts.set_max_write_buffer_number(16);
     // Updated default
     cfopts.set_level_compaction_dynamic_level_bytes(true);
-    //cfopts.set_merge_operator_associative("colors operator", merge_colors);
 
     let cf_metadata = ColumnFamilyDescriptor::new(METADATA, cfopts);
 
     let mut cfopts = Options::default();
     cfopts.set_max_write_buffer_number(16);
     // Updated default
     cfopts.set_level_compaction_dynamic_level_bytes(true);
-    //cfopts.set_merge_operator_associative("colors operator", merge_colors);
 
-    vec![cf_hashes, cf_metadata]
+    let cf_storage = ColumnFamilyDescriptor::new(STORAGE, cfopts);
+
+    let mut cfopts = Options::default();
+    cfopts.set_max_write_buffer_number(16);
+    // Updated default
+    cfopts.set_level_compaction_dynamic_level_bytes(true);
+
+    vec![cf_hashes, cf_metadata, cf_storage]
 }

src/core/src/index/revindex/mod.rs

@@ -23,7 +23,7 @@ use crate::sketch::Sketch;
 use crate::HashIntoType;
 use crate::Result;
 
-type DB = rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>;
+pub type DB = rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>;
 
 type QueryColors = HashMap<Color, Datasets>;
 type HashToColorT = HashMap<HashIntoType, Color, BuildNoHashHasher<HashIntoType>>;
@@ -198,7 +198,7 @@ impl RevIndex {
         }
     }
 
-    fn db_options() -> rocksdb::Options {
+    pub(crate) fn db_options() -> rocksdb::Options {
         let mut opts = rocksdb::Options::default();
         opts.set_max_open_files(500);
 

src/core/src/storage.rs

@@ -30,7 +30,14 @@ pub trait Storage {
     fn args(&self) -> StorageArgs;
 
     /// Load signature from internal path
-    fn load_sig(&self, path: &str) -> Result<SigStore>;
+    fn load_sig(&self, path: &str) -> Result<SigStore> {
+        let raw = self.load(path)?;
+        let sig = Signature::from_reader(&mut &raw[..])?
+            // TODO: select the right sig?
+            .swap_remove(0);
+
+        Ok(sig.into())
+    }
 
     /// Return a spec for creating/opening a storage
     fn spec(&self) -> String;
@@ -129,6 +136,17 @@ pub struct MemStorage {
     sigs: Arc<RwLock<HashMap<String, SigStore>>>,
 }
 
+#[cfg(feature = "branchwater")]
+// Column family for using rocksdb as a Storage
+pub(crate) const STORAGE: &str = "storage";
+
+#[cfg(feature = "branchwater")]
+/// Store data in RocksDB
+#[derive(Debug, Clone)]
+pub struct RocksDBStorage {
+    db: Arc<crate::index::revindex::DB>,
+}
+
 pub type Metadata<'a> = BTreeMap<&'a OsStr, &'a piz::read::FileMetadata<'a>>;
 
 // =========================================
@@ -145,6 +163,10 @@ impl InnerStorage {
                 InnerStorage::new(FSStorage::new("", path))
             }
             x if x.starts_with("memory") => InnerStorage::new(MemStorage::new()),
+            x if x.starts_with("rocksdb") => {
+                let path = x.split("://").last().expect("not a valid path");
+                InnerStorage::new(RocksDBStorage::from_path(path))
+            }
             x if x.starts_with("zip") => {
                 let path = x.split("://").last().expect("not a valid path");
                 InnerStorage::new(ZipStorage::from_file(path)?)
@@ -639,3 +661,49 @@ impl Storage for MemStorage {
         "memory://".into()
     }
 }
+
+#[cfg(feature = "branchwater")]
+impl RocksDBStorage {
+    pub fn from_path(path: &str) -> Self {
+        let mut opts = crate::index::revindex::RevIndex::db_options();
+        opts.create_if_missing(true);
+        opts.create_missing_column_families(true);
+        opts.prepare_for_bulk_load();
+
+        // prepare column family descriptors
+        let cfs = crate::index::revindex::disk_revindex::cf_descriptors();
+
+        let db =
+            Arc::new(crate::index::revindex::DB::open_cf_descriptors(&opts, path, cfs).unwrap());
+
+        Self { db }
+    }
+
+    pub fn from_db(db: Arc<crate::index::revindex::DB>) -> Self {
+        Self { db: db.clone() }
+    }
+}
+
+#[cfg(feature = "branchwater")]
+impl Storage for RocksDBStorage {
+    fn save(&self, path: &str, content: &[u8]) -> Result<String> {
+        let cf_storage = self.db.cf_handle(STORAGE).unwrap();
+        // TODO(lirber): deal with conflict for path?
+        self.db.put_cf(&cf_storage, path.as_bytes(), &content[..])?;
+        Ok(path.into())
+    }
+
+    fn load(&self, path: &str) -> Result<Vec<u8>> {
+        let cf_storage = self.db.cf_handle(STORAGE).unwrap();
+        let data = self.db.get_cf(&cf_storage, path.as_bytes())?;
+        data.ok_or_else(|| StorageError::DataReadError(path.into()).into())
+    }
+
+    fn args(&self) -> StorageArgs {
+        unimplemented!()
+    }
+
+    fn spec(&self) -> String {
+        "rocksdb://".into()
+    }
+}