transaction interface

4 files changed, 151 insertions, 51 deletions

diff --git a/database/src/backends/memory.rs b/database/src/backends/memory.rs
index f1952a3..3c2fdea 100644
--- a/database/src/backends/memory.rs
+++ b/database/src/backends/memory.rs
@@ -4,10 +4,14 @@
     Copyright (C) 2025 metamuffin <metamuffin.org>
 */
 
-use crate::backends::KV;
+use crate::backends::{Db, ReadTransaction, ReadTxnFunction, WriteTransaction, WriteTxnFunction};
 use anyhow::Result;
-use std::{collections::BTreeMap, sync::RwLock};
+use std::{
+    collections::BTreeMap,
+    sync::{RwLock, RwLockReadGuard, RwLockWriteGuard},
+};
 
+type Inner = BTreeMap<Vec<u8>, Vec<u8>>;
 pub struct Memory(RwLock<BTreeMap<Vec<u8>, Vec<u8>>>);
 
 impl Memory {
@@ -15,25 +19,54 @@ impl Memory {
         Self(RwLock::new(BTreeMap::new()))
     }
 }
-impl KV for Memory {
-    fn set(&self, key: &[u8], value: &[u8]) -> Result<()> {
-        self.0.write().unwrap().insert(key.to_vec(), value.to_vec());
+impl Db for Memory {
+    fn write_transaction(&self, f: &mut WriteTxnFunction) -> Result<()> {
+        f(&mut self.0.write().unwrap())
+    }
+    fn read_transaction(&self, f: &mut ReadTxnFunction) -> Result<()> {
+        f(&self.0.read().unwrap())
+    }
+}
+impl WriteTransaction for RwLockWriteGuard<'_, Inner> {
+    fn set(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
+        (**self).insert(key.to_vec(), value.to_vec());
+        Ok(())
+    }
+    fn del(&mut self, key: &[u8]) -> Result<()> {
+        (**self).remove(key);
         Ok(())
     }
+}
+impl ReadTransaction for RwLockWriteGuard<'_, Inner> {
     fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        Ok(self.0.read().unwrap().get(key).cloned())
+        Ok((**self).get(key).cloned())
     }
-    fn del(&self, key: &[u8]) -> Result<()> {
-        self.0.write().unwrap().remove(key);
-        Ok(())
+    fn next(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
+        Ok((**self)
+            .range(key.to_vec()..)
+            .next()
+            .map(|(k, _)| k.to_owned()))
+    }
+    fn prev(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
+        Ok((**self)
+            .range(..key.to_vec())
+            .next_back()
+            .map(|(k, _)| k.to_owned()))
+    }
+}
+impl ReadTransaction for RwLockReadGuard<'_, Inner> {
+    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
+        Ok((**self).get(key).cloned())
     }
     fn next(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        let r = self.0.read().unwrap();
-        Ok(r.range(key.to_vec()..).next().map(|(k, _)| k.to_owned()))
+        Ok((**self)
+            .range(key.to_vec()..)
+            .next()
+            .map(|(k, _)| k.to_owned()))
     }
     fn prev(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        let r = self.0.read().unwrap();
-        Ok(r.range(..key.to_vec())
+        Ok((**self)
+            .range(..key.to_vec())
             .next_back()
             .map(|(k, _)| k.to_owned()))
     }
diff --git a/database/src/backends/mod.rs b/database/src/backends/mod.rs
index 814cc50..a95b00a 100644
--- a/database/src/backends/mod.rs
+++ b/database/src/backends/mod.rs
@@ -12,15 +12,24 @@ use crate::backends::{memory::Memory, redb::Redb, rocksdb::Rocksdb};
 use anyhow::{Result, bail};
 use std::{path::Path, sync::Arc};
 
-pub trait KV {
-    fn set(&self, key: &[u8], value: &[u8]) -> Result<()>;
-    fn get<'a>(&'a self, key: &[u8]) -> Result<Option<Vec<u8>>>;
-    fn del(&self, key: &[u8]) -> Result<()>;
+pub type WriteTxnFunction = dyn FnMut(&mut dyn WriteTransaction) -> Result<()>;
+pub type ReadTxnFunction = dyn FnMut(&dyn ReadTransaction) -> Result<()>;
+
+pub trait Db {
+    fn write_transaction(&self, f: &mut WriteTxnFunction) -> Result<()>;
+    fn read_transaction(&self, f: &mut ReadTxnFunction) -> Result<()>;
+}
+pub trait WriteTransaction: ReadTransaction {
+    fn set(&mut self, key: &[u8], value: &[u8]) -> Result<()>;
+    fn del(&mut self, key: &[u8]) -> Result<()>;
+}
+pub trait ReadTransaction {
+    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>>;
     fn next(&self, key: &[u8]) -> Result<Option<Vec<u8>>>;
     fn prev(&self, key: &[u8]) -> Result<Option<Vec<u8>>>;
 }
 
-pub fn create_backend(driver: &str, path: &Path) -> Result<Arc<dyn KV>> {
+pub fn create_backend(driver: &str, path: &Path) -> Result<Arc<dyn Db>> {
     Ok(match driver {
         "rocksdb" => Arc::new(Rocksdb::new(path)?),
         "redb" => Arc::new(Redb::new(path)?),
diff --git a/database/src/backends/redb.rs b/database/src/backends/redb.rs
index 1b672b6..d2849de 100644
--- a/database/src/backends/redb.rs
+++ b/database/src/backends/redb.rs
@@ -4,11 +4,10 @@
     Copyright (C) 2025 metamuffin <metamuffin.org>
 */
 
-use std::path::Path;
-
-use crate::backends::KV;
+use crate::backends::{Db, ReadTransaction, ReadTxnFunction, WriteTransaction, WriteTxnFunction};
 use anyhow::Result;
-use redb::{Database, ReadableDatabase, TableDefinition};
+use redb::{Database, ReadableDatabase, ReadableTable, TableDefinition};
+use std::path::Path;
 
 pub struct Redb {
     db: Database,
@@ -23,29 +22,62 @@ impl Redb {
         })
     }
 }
-impl KV for Redb {
-    fn set(&self, key: &[u8], value: &[u8]) -> Result<()> {
-        let txn = self.db.begin_write()?;
-        txn.open_table(TABLE)?.insert(key, value)?;
+impl Db for Redb {
+    fn write_transaction(&self, f: &mut WriteTxnFunction) -> Result<()> {
+        let mut txn = self.db.begin_write()?;
+        f(&mut txn)?;
         txn.commit()?;
         Ok(())
     }
-    fn del(&self, key: &[u8]) -> Result<()> {
-        let txn = self.db.begin_write()?;
-        txn.open_table(TABLE)?.remove(key)?;
-        txn.commit()?;
+    fn read_transaction(&self, f: &mut ReadTxnFunction) -> Result<()> {
+        let mut txn = self.db.begin_read()?;
+        f(&mut txn)?;
         Ok(())
     }
+}
+impl WriteTransaction for redb::WriteTransaction {
+    fn set(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
+        self.open_table(TABLE)?.insert(key, value)?;
+        Ok(())
+    }
+    fn del(&mut self, key: &[u8]) -> Result<()> {
+        self.open_table(TABLE)?.remove(key)?;
+        Ok(())
+    }
+}
+impl ReadTransaction for redb::WriteTransaction {
+    fn get<'a>(&'a self, key: &[u8]) -> Result<Option<Vec<u8>>> {
+        match self.open_table(TABLE)?.get(key)? {
+            Some(v) => Ok(Some(v.value().to_vec())),
+            None => Ok(None),
+        }
+    }
+    fn next(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
+        let table = self.open_table(TABLE)?;
+        let mut iter = table.range(key..)?;
+        match iter.next() {
+            Some(k) => Ok(Some(k?.0.value().to_vec())),
+            None => Ok(None),
+        }
+    }
+    fn prev(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
+        let table = self.open_table(TABLE)?;
+        let mut iter = table.range(..key)?;
+        match iter.next_back() {
+            Some(k) => Ok(Some(k?.0.value().to_vec())),
+            None => Ok(None),
+        }
+    }
+}
+impl ReadTransaction for redb::ReadTransaction {
     fn get<'a>(&'a self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        let txn = self.db.begin_read()?;
-        match txn.open_table(TABLE)?.get(key)? {
+        match self.open_table(TABLE)?.get(key)? {
             Some(v) => Ok(Some(v.value().to_vec())),
             None => Ok(None),
         }
     }
     fn next(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        let txn = self.db.begin_read()?;
-        let table = txn.open_table(TABLE)?;
+        let table = self.open_table(TABLE)?;
         let mut iter = table.range(key..)?;
         match iter.next() {
             Some(k) => Ok(Some(k?.0.value().to_vec())),
@@ -53,8 +85,7 @@ impl KV for Redb {
         }
     }
     fn prev(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        let txn = self.db.begin_read()?;
-        let table = txn.open_table(TABLE)?;
+        let table = self.open_table(TABLE)?;
         let mut iter = table.range(..key)?;
         match iter.next_back() {
             Some(k) => Ok(Some(k?.0.value().to_vec())),
diff --git a/database/src/backends/rocksdb.rs b/database/src/backends/rocksdb.rs
index 31229a9..cdcb60a 100644
--- a/database/src/backends/rocksdb.rs
+++ b/database/src/backends/rocksdb.rs
@@ -4,41 +4,68 @@
     Copyright (C) 2025 metamuffin <metamuffin.org>
 */
 
-use std::path::Path;
-
-use crate::backends::KV;
+use crate::backends::{Db, ReadTransaction, WriteTransaction, WriteTxnFunction};
 use anyhow::Result;
-use rocksdb::DB;
+use rocksdb::{ErrorKind, OptimisticTransactionDB};
+use std::path::Path;
 
 pub struct Rocksdb {
-    db: DB,
+    db: OptimisticTransactionDB,
 }
 
 impl Rocksdb {
     pub fn new(path: &Path) -> Result<Self> {
         Ok(Self {
-            db: DB::open_default(path)?,
+            db: OptimisticTransactionDB::open_default(path)?,
         })
     }
 }
-impl KV for Rocksdb {
-    fn set(&self, key: &[u8], value: &[u8]) -> Result<()> {
-        Ok(self.db.put(key, value)?)
+impl Db for Rocksdb {
+    fn write_transaction(&self, f: &mut WriteTxnFunction) -> Result<()> {
+        loop {
+            let mut txn = self.db.transaction();
+            f(&mut txn)?;
+            match txn.commit() {
+                Ok(()) => break Ok(()),
+                Err(e) if e.kind() == ErrorKind::Busy => continue,
+                Err(e) => return Err(e.into()),
+            }
+        }
     }
-    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        Ok(self.db.get(key)?)
+    fn read_transaction(&self, f: &mut super::ReadTxnFunction) -> Result<()> {
+        loop {
+            let txn = self.db.transaction();
+            f(&txn)?;
+            match txn.commit() {
+                Ok(()) => break Ok(()),
+                Err(e) if e.kind() == ErrorKind::Busy => continue,
+                Err(e) => return Err(e.into()),
+            }
+        }
     }
-    fn del(&self, key: &[u8]) -> Result<()> {
-        Ok(self.db.delete(key)?)
+}
+impl WriteTransaction for rocksdb::Transaction<'_, OptimisticTransactionDB> {
+    fn set(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
+        Ok(self.put(key, value)?)
+    }
+
+    fn del(&mut self, key: &[u8]) -> Result<()> {
+        Ok(self.delete(key)?)
+    }
+}
+impl ReadTransaction for rocksdb::Transaction<'_, OptimisticTransactionDB> {
+    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
+        Ok(self.get(key)?)
     }
     fn next(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        let mut it = self.db.raw_iterator();
+        let mut it = self.raw_iterator();
         it.seek_for_prev(key);
         it.next();
         Ok(it.key().map(Vec::from))
     }
+
     fn prev(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        let mut it = self.db.raw_iterator();
+        let mut it = self.raw_iterator();
         it.seek(key);
         it.prev();
         Ok(it.key().map(Vec::from))