rocksdb cache backend

5 files changed, 81 insertions, 3 deletions

diff --git a/cache/Cargo.toml b/cache/Cargo.toml
index d7fcd81..c545fd2 100644
--- a/cache/Cargo.toml
+++ b/cache/Cargo.toml
@@ -13,3 +13,4 @@ rand = "0.9.2"
 serde = "1.0.228"
 serde_json = "1.0.145"
 percent-encoding = "2.3.2"
+rocksdb = { version = "0.24.0", features = ["multi-threaded-cf"] }
diff --git a/cache/src/backends/mod.rs b/cache/src/backends/mod.rs
index 6b7dac3..5872255 100644
--- a/cache/src/backends/mod.rs
+++ b/cache/src/backends/mod.rs
@@ -4,10 +4,23 @@
     Copyright (C) 2025 metamuffin <metamuffin.org>
 */
 pub mod filesystem;
+pub mod rocksdb;
 
-use anyhow::Result;
+use crate::{
+    CONF,
+    backends::{filesystem::Filesystem, rocksdb::Rocksdb},
+};
+use anyhow::{Result, bail};
 
 pub(crate) trait CacheStorage: Send + Sync + 'static {
     fn store(&self, key: String, value: &[u8]) -> Result<()>;
     fn read(&self, key: &str) -> Result<Option<Vec<u8>>>;
 }
+
+pub fn init_backend() -> Result<Box<dyn CacheStorage>> {
+    Ok(match CONF.driver.as_str() {
+        "filesystem" => Box::new(Filesystem::new(&CONF)),
+        "rocksdb" => Box::new(Rocksdb::new(&CONF)?),
+        _ => bail!("unknown driver"),
+    })
+}
diff --git a/cache/src/backends/rocksdb.rs b/cache/src/backends/rocksdb.rs
new file mode 100644
index 0000000..183664f
--- /dev/null
+++ b/cache/src/backends/rocksdb.rs
@@ -0,0 +1,26 @@
+/*
+    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
+    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
+    Copyright (C) 2025 metamuffin <metamuffin.org>
+*/
+
+use crate::{Config, backends::CacheStorage};
+use anyhow::Result;
+use rocksdb::DB;
+
+pub struct Rocksdb(DB);
+
+impl Rocksdb {
+    pub fn new(config: &Config) -> Result<Self> {
+        Ok(Self(rocksdb::DB::open_default(config.path.clone())?))
+    }
+}
+
+impl CacheStorage for Rocksdb {
+    fn store(&self, key: String, value: &[u8]) -> Result<()> {
+        Ok(self.0.put(key, value)?)
+    }
+    fn read(&self, key: &str) -> Result<Option<Vec<u8>>> {
+        Ok(self.0.get(key)?)
+    }
+}
diff --git a/cache/src/bin/cache_fs_to_rocksdb.rs b/cache/src/bin/cache_fs_to_rocksdb.rs
new file mode 100644
index 0000000..0a3b43d
--- /dev/null
+++ b/cache/src/bin/cache_fs_to_rocksdb.rs
@@ -0,0 +1,37 @@
+/*
+    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
+    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
+    Copyright (C) 2025 metamuffin <metamuffin.org>
+*/
+
+use anyhow::{Result, bail};
+use rocksdb::DB;
+use std::{env::args, fs::File, io::Read, path::Path};
+
+fn main() -> Result<()> {
+    let in_path = args().nth(1).unwrap();
+    let out_path = args().nth(2).unwrap();
+    let db = DB::open_default(out_path)?;
+    if !in_path.ends_with("/") {
+        bail!("path needs to end with /")
+    }
+    traverse(&db, &in_path, in_path.as_ref())?;
+    Ok(())
+}
+
+fn traverse(db: &DB, prefix: &str, path: &Path) -> Result<()> {
+    if path.is_dir() {
+        for e in path.read_dir()? {
+            traverse(db, prefix, &e?.path())?;
+        }
+    }
+    if path.is_file() {
+        let key = path.to_string_lossy();
+        let key = key.strip_prefix(prefix).unwrap();
+        let mut value = Vec::new();
+        File::open(path)?.read_to_end(&mut value)?;
+        println!("{key}");
+        db.put(key, value)?;
+    }
+    Ok(())
+}
diff --git a/cache/src/lib.rs b/cache/src/lib.rs
index 20e1424..d68266f 100644
--- a/cache/src/lib.rs
+++ b/cache/src/lib.rs
@@ -6,7 +6,7 @@
 mod backends;
 mod helper;
 
-use crate::backends::{CacheStorage, filesystem::Filesystem};
+use crate::backends::{CacheStorage, init_backend};
 use anyhow::{Context, Result, anyhow};
 use log::{info, warn};
 use serde::{Deserialize, Serialize};
@@ -26,6 +26,7 @@ pub use helper::{EscapeKey, HashKey};
 
 #[derive(Debug, Deserialize)]
 pub struct Config {
+    driver: String,
     path: PathBuf,
     max_in_memory_cache_size: usize,
 }
@@ -49,7 +50,7 @@ static CACHE_STORE: OnceLock<Box<dyn CacheStorage>> = OnceLock::new();
 
 pub fn init_cache() -> Result<()> {
     CACHE_STORE
-        .set(Box::new(Filesystem::new(&CONF)))
+        .set(init_backend().context("cache backend")?)
         .map_err(|_| ())
         .unwrap();
     Ok(())