133 lines
3.1 KiB
Rust
133 lines
3.1 KiB
Rust
|
use std::collections::HashMap;
|
||
|
use std::sync::{Arc, RwLock};
|
||
|
|
||
|
use arc_swap::ArcSwapOption;
|
||
|
|
||
|
use sled::transaction::{
|
||
|
ConflictableTransactionError, TransactionError, Transactional, TransactionalTree, UnabortableTransactionError
|
||
|
};
|
||
|
|
||
|
use crate::{Error, IDb, ITx, ITxFn, Result, TxError, TxFnResult, TxResult, Value, Db};
|
||
|
|
||
|
impl From<sled::Error> for Error {
|
||
|
fn from(e: sled::Error) -> Error {
|
||
|
Error(format!("{}", e).into())
|
||
|
}
|
||
|
}
|
||
|
|
||
|
pub struct SledDb {
|
||
|
db: sled::Db,
|
||
|
trees: RwLock<(Vec<sled::Tree>, HashMap<String, usize>)>,
|
||
|
}
|
||
|
|
||
|
impl SledDb {
|
||
|
pub fn new(db: sled::Db) -> Db {
|
||
|
let s = Self {
|
||
|
db,
|
||
|
trees: RwLock::new((Vec::new(), HashMap::new())),
|
||
|
};
|
||
|
Db(Arc::new(s))
|
||
|
}
|
||
|
|
||
|
fn get_tree(&self, i: usize) -> Result<sled::Tree> {
|
||
|
self.trees
|
||
|
.read()
|
||
|
.unwrap()
|
||
|
.0
|
||
|
.get(i)
|
||
|
.cloned()
|
||
|
.ok_or(Error("invalid tree id".into()))
|
||
|
}
|
||
|
}
|
||
|
|
||
|
impl IDb for SledDb {
|
||
|
fn tree(&self, name: &str) -> Result<usize> {
|
||
|
let mut trees = self.trees.write().unwrap();
|
||
|
if let Some(i) = trees.1.get(name) {
|
||
|
Ok(*i)
|
||
|
} else {
|
||
|
let tree = self.db.open_tree(name)?;
|
||
|
let i = trees.0.len();
|
||
|
trees.0.push(tree);
|
||
|
trees.1.insert(name.to_string(), i);
|
||
|
Ok(i)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
fn get<'a>(&'a self, tree: usize, key: &[u8]) -> Result<Option<Value<'a>>> {
|
||
|
let tree = self.get_tree(tree)?;
|
||
|
Ok(tree.get(key)?.map(|v| v.to_vec().into()))
|
||
|
}
|
||
|
fn put(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<()> {
|
||
|
let tree = self.get_tree(tree)?;
|
||
|
tree.insert(key, value)?;
|
||
|
Ok(())
|
||
|
}
|
||
|
|
||
|
fn transaction(&self, f: &dyn ITxFn) -> TxResult<(), ()> {
|
||
|
let trees = self.trees.read().unwrap();
|
||
|
let res = trees.0.transaction(|txtrees| {
|
||
|
let tx = SledTx {
|
||
|
trees: txtrees,
|
||
|
err: ArcSwapOption::new(None),
|
||
|
};
|
||
|
match f.try_on(&tx) {
|
||
|
TxFnResult::Ok => {
|
||
|
assert!(tx.err.into_inner().is_none());
|
||
|
Ok(())
|
||
|
}
|
||
|
TxFnResult::Abort => Err(ConflictableTransactionError::Abort(())),
|
||
|
TxFnResult::Err => {
|
||
|
let err_arc = tx
|
||
|
.err
|
||
|
.into_inner()
|
||
|
.expect("Transaction did not store error");
|
||
|
let err = Arc::try_unwrap(err_arc).ok().expect("Many refs");
|
||
|
Err(err.into())
|
||
|
}
|
||
|
}
|
||
|
});
|
||
|
match res {
|
||
|
Ok(()) => Ok(()),
|
||
|
Err(TransactionError::Abort(())) => Err(TxError::Abort(())),
|
||
|
Err(TransactionError::Storage(s)) => Err(TxError::Db(s.into())),
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// ----
|
||
|
|
||
|
struct SledTx<'a> {
|
||
|
trees: &'a [TransactionalTree],
|
||
|
err: ArcSwapOption<UnabortableTransactionError>,
|
||
|
}
|
||
|
|
||
|
impl<'a> SledTx<'a> {
|
||
|
fn save_error<R>(&self, v: std::result::Result<R, UnabortableTransactionError>) -> Result<R> {
|
||
|
match v {
|
||
|
Ok(x) => Ok(x),
|
||
|
Err(e) => {
|
||
|
let txt = format!("{}", e);
|
||
|
self.err.store(Some(Arc::new(e)));
|
||
|
Err(Error(txt.into()))
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
impl<'a> ITx<'a> for SledTx<'a> {
|
||
|
fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value<'a>>> {
|
||
|
let tree = self.trees.get(tree)
|
||
|
.ok_or(Error("invalid tree id".into()))?;
|
||
|
let tmp = self.save_error(tree.get(key))?;
|
||
|
Ok(tmp.map(|v| v.to_vec().into()))
|
||
|
}
|
||
|
|
||
|
fn put(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<()> {
|
||
|
let tree = self.trees.get(tree)
|
||
|
.ok_or(Error("invalid tree id".into()))?;
|
||
|
self.save_error(tree.insert(key, value))?;
|
||
|
Ok(())
|
||
|
}
|
||
|
}
|