use std::fmt;
use std::fs::create_dir_all;
use std::path::Path;
use std::str;
use std::sync::Arc;
use crate::conf::Conf;
use crate::error::{Error, Result};
use crate::key_composer::KeyComposer;
use crate::meta::KeyMeta;
use crate::string::conf::Set;
use crate::string::{StringSetType, decode_string_value, is_string_expired};
use fjall::config::{BlockSizePolicy, CompressionPolicy};
use fjall::{CompressionType, Database, Keyspace, KeyspaceCreateOptions, PersistMode};
#[derive(Clone)]
pub struct WeDb {
pub db: Arc<Database>,
pub data_ks: Keyspace,
pub meta_ks: Keyspace,
}
impl fmt::Debug for WeDb {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("WeDb").finish()
}
}
impl WeDb {
pub const DEFAULT_KEYSPACE: &'static str = "data";
pub const META_KEYSPACE: &'static str = "meta";
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
Self::open_with_conf(&Conf {
data_path: path.as_ref().to_string_lossy().to_string(),
..Default::default()
})
}
pub fn open_with_conf(conf: &Conf) -> Result<Self> {
let path = Path::new(&conf.data_path);
if let Some(parent) = path.parent() {
create_dir_all(parent)?;
}
let mut builder = Database::builder(path);
let mut comp_type = CompressionType::Lz4;
let mut data_comp_policy = CompressionPolicy::new([
CompressionType::None,
CompressionType::None,
CompressionType::Lz4,
]);
if let Some(cache_size) = conf.cache_size {
builder = builder.cache_size(cache_size as u64);
}
if let Some(ref comp) = conf.compression {
match comp.to_ascii_lowercase().as_str() {
"none" | "off" | "disabled" | "false" | "0" => {
comp_type = CompressionType::None;
data_comp_policy = CompressionPolicy::disabled();
}
"lz4" | "on" | "enabled" | "true" | "1" => {
comp_type = CompressionType::Lz4;
data_comp_policy = CompressionPolicy::new([
CompressionType::None,
CompressionType::None,
CompressionType::Lz4,
]);
}
_ => {}
}
}
builder = builder.journal_compression(comp_type);
let db = builder.open().map_err(|e| {
Error::internal_with_source(format!("Failed to open Fjall at {path:?}"), e)
})?;
let data_policy = data_comp_policy.clone();
let data_ks = db
.keyspace(Self::DEFAULT_KEYSPACE, move || {
KeyspaceCreateOptions::default()
.data_block_size_policy(BlockSizePolicy::all(16 * 1024))
.data_block_compression_policy(data_policy)
})
.map_err(|e| Error::internal_with_source("Failed to open default keyspace", e))?;
let meta_policy = data_comp_policy;
let meta_ks = db
.keyspace(Self::META_KEYSPACE, move || {
KeyspaceCreateOptions::default()
.data_block_size_policy(BlockSizePolicy::all(4 * 1024))
.data_block_compression_policy(meta_policy)
})
.map_err(|e| Error::internal_with_source("Failed to open meta keyspace", e))?;
Ok(Self {
db: Arc::new(db),
data_ks,
meta_ks,
})
}
#[inline]
pub fn database(&self) -> &Arc<Database> {
&self.db
}
#[inline]
pub fn data_keyspace(&self) -> &Keyspace {
&self.data_ks
}
#[inline]
pub fn meta_keyspace(&self) -> &Keyspace {
&self.meta_ks
}
#[inline]
pub fn keyspace(&self, name: &str) -> Result<Keyspace> {
self.db
.keyspace(name, KeyspaceCreateOptions::default)
.map_err(|e| Error::internal_with_source(format!("Keyspace '{name}' error"), e))
}
#[inline]
pub fn persist(&self, mode: PersistMode) -> Result<()> {
self.db
.persist(mode)
.map_err(|e| Error::internal_with_source("Persist error", e))
}
#[inline]
pub fn get(&self, key: impl AsRef<[u8]>) -> Result<Option<Vec<u8>>> {
let (val, _) = self.get_with_expire(key)?;
Ok(val)
}
#[inline]
pub fn get_with_expire(&self, key: impl AsRef<[u8]>) -> Result<(Option<Vec<u8>>, u64)> {
let key_bytes = key.as_ref();
let kc = KeyComposer::new("default");
let raw_k = kc.raw_key_bytes(key_bytes);
let now_ms = coarsetime::Clock::now_since_epoch().as_millis();
match self.data_ks.get(&raw_k)? {
Some(raw) => {
let (expire_at, payload) = decode_string_value(&raw);
if is_string_expired(expire_at, now_ms) {
Ok((None, 0))
} else {
Ok((Some(payload.to_vec()), expire_at))
}
}
None => Ok((None, 0)),
}
}
pub fn set<'a>(
&self,
key: impl AsRef<[u8]>,
val: impl AsRef<[u8]>,
conf_li: impl AsRef<[Set<'a>]>,
) -> Result<Option<Vec<u8>>> {
let options = conf_li.as_ref();
let mut set_type = StringSetType::None;
let mut get = false;
let mut keep_ttl = false;
let mut expire = 0u64;
let mut cmp_value: Option<&'a [u8]> = None;
let now_ms = coarsetime::Clock::now_since_epoch().as_millis();
for opt in options {
match opt {
Set::Ex(sec) => expire = now_ms + sec * 1000,
Set::Px(ms) => expire = now_ms + ms,
Set::ExAt(sec) => expire = sec * 1000,
Set::PxAt(ms) => expire = *ms,
Set::KeepTtl => keep_ttl = true,
Set::Nx => set_type = StringSetType::Nx,
Set::Xx => set_type = StringSetType::Xx,
Set::IfEq(expected) => {
set_type = StringSetType::IfEq;
cmp_value = Some(expected);
}
Set::IfNe(expected) => {
set_type = StringSetType::IfNe;
cmp_value = Some(expected);
}
Set::IfDeq(expected) => {
set_type = StringSetType::IfDeq;
cmp_value = Some(expected);
}
Set::IfDne(expected) => {
set_type = StringSetType::IfDne;
cmp_value = Some(expected);
}
Set::Get => get = true,
}
}
let args = crate::string::StringSetArgs {
expire,
set_type,
get,
keep_ttl,
cmp_value,
};
self.set_args(key, val, &args)
}
pub fn del(&self, keys: &[impl AsRef<[u8]>]) -> Result<usize> {
let mut deleted = 0;
let kc = KeyComposer::new("default");
let mut batch = self.db.batch();
for k in keys {
let k_bytes = k.as_ref();
let mut key_deleted = false;
let raw_k = kc.raw_key_bytes(k_bytes);
if self.data_ks.contains_key(&raw_k)? {
key_deleted = true;
batch.remove(&self.data_ks, raw_k.as_ref());
}
if self.meta_ks.contains_key(&raw_k)? {
key_deleted = true;
batch.remove(&self.meta_ks, raw_k.as_ref());
}
if let Ok(k_str) = str::from_utf8(k_bytes) {
for meta_k in kc.all_meta_keys(k_str) {
if self.meta_ks.contains_key(meta_k.as_bytes())? {
key_deleted = true;
batch.remove(&self.meta_ks, meta_k.as_bytes());
}
}
for data_prefix in kc.all_data_prefixes(k_str) {
for item in self.data_ks.prefix(&data_prefix) {
key_deleted = true;
batch.remove(&self.data_ks, item.key()?);
}
}
}
if key_deleted {
deleted += 1;
}
}
batch.commit()?;
Ok(deleted)
}
pub fn exists(&self, keys: &[impl AsRef<[u8]>]) -> Result<usize> {
let mut count = 0;
let kc = KeyComposer::new("default");
let now_ms = coarsetime::Clock::now_since_epoch().as_millis();
for k in keys {
let k_bytes = k.as_ref();
let raw_k = kc.raw_key_bytes(k_bytes);
if let Some(raw) = self.data_ks.get(&raw_k)? {
let (expire_at, _) = decode_string_value(&raw);
if !is_string_expired(expire_at, now_ms) {
count += 1;
continue;
}
}
if let Some(raw_meta) = self.meta_ks.get(&raw_k)?
&& let Some(meta) = KeyMeta::decode(&raw_meta)
&& !meta.is_expired(now_ms)
{
count += 1;
continue;
}
if let Ok(k_str) = str::from_utf8(k_bytes) {
let mut found = false;
for meta_k in kc.all_meta_keys(k_str) {
if let Some(raw_meta) = self.meta_ks.get(meta_k.as_bytes())?
&& let Some(meta) = KeyMeta::decode(&raw_meta)
&& !meta.is_expired(now_ms)
{
found = true;
break;
}
}
if found {
count += 1;
}
}
}
Ok(count)
}
pub fn flushall(&self) -> Result<()> {
let mut batch = self.db.batch();
for item in self.data_ks.iter() {
let k = item.key()?;
batch.remove(&self.data_ks, k);
}
for item in self.meta_ks.iter() {
let k = item.key()?;
batch.remove(&self.meta_ks, k);
}
batch.commit()?;
Ok(())
}
}