1 files changed, 24 insertions, 8 deletions
diff --git a/src/cache.rs b/src/cache.rs
index a4d6dfb..ebf5e82 100644
--- a/src/cache.rs
+++ b/src/cache.rs
@@ -1,7 +1,6 @@
 use anyhow::Result;
 use futures_util::{StreamExt, stream::FuturesUnordered};
 use log::{debug, info};
-use redb::{Database, Durability, TableDefinition};
 use std::{
     fs::create_dir_all,
     path::{Path, PathBuf},
@@ -14,9 +13,15 @@ use tokio::{
 
 pub enum Cache {
     Directory(PathBuf),
-    Redb { db: Database, counter: AtomicUsize },
+    Redb {
+        db: redb::Database,
+        counter: AtomicUsize,
+    },
+    Rocksdb {
+        db: rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>,
+    },
 }
-const T_DOWNLOAD: TableDefinition<&str, &[u8]> = TableDefinition::new("dl");
+const T_DOWNLOAD: redb::TableDefinition<&str, &[u8]> = redb::TableDefinition::new("dl");
 impl Cache {
     pub fn new_directory() -> Result<Self> {
         let cachedir = xdg::BaseDirectories::with_prefix("weareearth")
@@ -27,9 +32,9 @@ impl Cache {
         create_dir_all(cachedir.join("NodeData"))?;
         Ok(Self::Directory(cachedir))
     }
-    pub fn new_db(path: &Path) -> Result<Self> {
-        info!("opening db...");
-        let db = Database::create(path)?;
+    pub fn new_redb(path: &Path) -> Result<Self> {
+        info!("opening redb...");
+        let db = redb::Database::create(path)?;
         info!("done");
         {
             let txn = db.begin_write()?;
@@ -41,6 +46,12 @@ impl Cache {
             counter: AtomicUsize::new(0),
         })
     }
+    pub fn new_rocksdb(path: &Path) -> Result<Self> {
+        info!("opening rocksdb...");
+        let db = rocksdb::DBWithThreadMode::open_default(path)?;
+        info!("done");
+        Ok(Self::Rocksdb { db })
+    }
 
     pub async fn transfer_entries(&self, other: &Cache) -> Result<()> {
         match self {
@@ -104,6 +115,7 @@ impl Cache {
                     Ok(None)
                 }
             }
+            Cache::Rocksdb { db } => Ok(db.get(path)?),
         }
     }
     pub async fn insert(&self, path: &str, data: &[u8]) -> Result<()> {
@@ -118,10 +130,10 @@ impl Cache {
                 let mut txn = db.begin_write()?;
                 if n % 100_000 == 0 {
                     info!("flushing db");
-                    txn.set_durability(Durability::Immediate);
+                    txn.set_durability(redb::Durability::Immediate);
                     txn.set_quick_repair(true);
                 } else {
-                    txn.set_durability(Durability::None)
+                    txn.set_durability(redb::Durability::None)
                 };
                 let mut table = txn.open_table(T_DOWNLOAD)?;
                 table.insert(path, data)?;
@@ -129,6 +141,10 @@ impl Cache {
                 txn.commit()?;
                 Ok(())
             }
+            Cache::Rocksdb { db } => {
+                db.put(path, data)?;
+                Ok(())
+            }
         }
     }
 }