use std::{fmt, ops::Bound, sync::Arc};
use parking_lot::Mutex;
use crate::{
Db,
api::key::cleanup_composite_data_raw,
engine::{Engine, KvEntry, Partition},
error::{Error, Result},
key_composer::{
CATALOG_PREFIX, KeyComposer, KeyTag, NS_NEXT_ID_KEY, decode_oppv_u64,
encode_catalog_db_key_fixed, encode_catalog_ns_prefix_fixed, encode_db_next_id_key_fixed,
encode_oppv_u64_fixed,
},
meta::{KeyMeta, current_now_ms, init_version_counter},
string::{decode_string_value, is_string_expired},
wedb::{DbBatch, IntoOptId, Namespace, Namespaces},
};
pub const DATA: &str = "data";
pub const META: &str = "meta";
#[inline]
pub fn clear_ks_prefix<E: Engine>(
partition: &E::Partition,
prefix: &[u8],
batch: &mut DbBatch<E>,
count: &mut u64,
) -> Result<()>
where
Error: From<E::Error>,
{
for item in partition.prefix(prefix) {
let entry = item?;
let k = entry.key();
if !k.starts_with(prefix) {
break;
}
batch.rm(partition, k);
*count += 1;
}
Ok(())
}
#[inline]
fn sweep_ks<P: Partition, F>(
ks: &P,
cursor: &mut Option<Vec<u8>>,
sample_limit: usize,
mut on_entry: F,
) -> Result<()>
where
Error: From<P::Error>,
F: FnMut(&[u8], &[u8]) -> Result<()>,
{
let mut count = 0;
let mut last_key = None;
let start_bound = cursor.as_deref().map_or(Bound::Unbounded, Bound::Excluded);
for guard in ks.range((start_bound, Bound::Unbounded)).take(sample_limit) {
let entry = guard?;
on_entry(entry.key(), entry.value())?;
count += 1;
if count == sample_limit {
last_key = Some(entry.key().to_vec());
}
}
*cursor = last_key;
Ok(())
}
#[derive(Default)]
pub struct ExpireCursors {
pub(crate) data_cursor: Option<Vec<u8>>,
pub(crate) meta_cursor: Option<Vec<u8>>,
}
pub(crate) struct WeDbInner<E: Engine> {
pub(crate) engine: Arc<E>,
pub(crate) data: E::Partition,
pub(crate) meta: E::Partition,
pub(crate) ns_lock: Mutex<()>,
pub(crate) expire_cursor: Mutex<ExpireCursors>,
}
#[inline]
pub(crate) fn activate_db_impl<E: Engine>(meta: &E::Partition, ns_id: u64, db_id: u64) -> Result<()>
where
Error: From<E::Error>,
{
let mut buf = [0u8; 20];
let len = encode_catalog_db_key_fixed(ns_id, db_id, &mut buf);
meta.insert(&buf[..len], b"")?;
Ok(())
}
#[inline]
pub(crate) fn next_namespace_id_impl<E: Engine>(
meta: &E::Partition,
lock: &Mutex<()>,
) -> Result<u64>
where
Error: From<E::Error>,
{
let _guard = lock.lock();
let current_id = meta
.get(NS_NEXT_ID_KEY)?
.and_then(|val| decode_oppv_u64(&val))
.map(|(v, _)| v)
.unwrap_or(1);
let new_id = current_id;
let next_val = current_id + 1;
let mut next_val_buf = [0u8; 9];
let next_len = encode_oppv_u64_fixed(next_val, &mut next_val_buf);
meta.insert(NS_NEXT_ID_KEY, &next_val_buf[..next_len])?;
Ok(new_id)
}
#[inline]
pub(crate) fn next_db_id_impl<E: Engine>(
meta: &E::Partition,
lock: &Mutex<()>,
ns_id: u64,
) -> Result<u64>
where
Error: From<E::Error>,
{
let _guard = lock.lock();
let mut key_buf = [0u8; 12];
let key_len = encode_db_next_id_key_fixed(ns_id, &mut key_buf);
let key = &key_buf[..key_len];
let current_id = meta
.get(key)?
.and_then(|val| decode_oppv_u64(&val))
.map(|(v, _)| v)
.unwrap_or(1);
let new_id = current_id;
let next_val = current_id + 1;
let mut next_val_buf = [0u8; 9];
let next_len = encode_oppv_u64_fixed(next_val, &mut next_val_buf);
meta.insert(key, &next_val_buf[..next_len])?;
Ok(new_id)
}
#[inline]
pub(crate) fn db_rm_impl<E: Engine>(
data: &E::Partition,
meta: &E::Partition,
engine: &E,
ns_id: u64,
db_id: u64,
) -> Result<u64>
where
Error: From<E::Error>,
{
let mut count = 0u64;
let mut batch = DbBatch::<E>::new(data.clone(), meta.clone(), engine.batch());
let kc = KeyComposer::new(ns_id, db_id);
let mut prefix_buf = [0u8; 19];
let prefix_len = kc.encode_scope_prefix_fixed(&mut prefix_buf);
let prefix = &prefix_buf[..prefix_len];
clear_ks_prefix(data, prefix, &mut batch, &mut count)?;
clear_ks_prefix(meta, prefix, &mut batch, &mut count)?;
let mut cat_key_buf = [0u8; 20];
let cat_len = encode_catalog_db_key_fixed(ns_id, db_id, &mut cat_key_buf);
batch.rm_meta(&cat_key_buf[..cat_len]);
batch.commit()?;
Ok(count)
}
#[inline]
pub(crate) fn namespace_rm_impl<E: Engine>(
data: &E::Partition,
meta: &E::Partition,
engine: &E,
ns_id: u64,
) -> Result<u64>
where
Error: From<E::Error>,
{
let mut count = 0u64;
let mut batch = DbBatch::<E>::new(data.clone(), meta.clone(), engine.batch());
let mut ns_prefix_buf = [0u8; 10];
let ns_len = KeyComposer::encode_ns_prefix_fixed(ns_id, &mut ns_prefix_buf);
let ns_prefix = &ns_prefix_buf[..ns_len];
clear_ks_prefix(data, ns_prefix, &mut batch, &mut count)?;
clear_ks_prefix(meta, ns_prefix, &mut batch, &mut count)?;
let mut cat_prefix_buf = [0u8; 11];
let cat_len = encode_catalog_ns_prefix_fixed(ns_id, &mut cat_prefix_buf);
let mut _dummy = 0;
clear_ks_prefix(meta, &cat_prefix_buf[..cat_len], &mut batch, &mut _dummy)?;
let mut db_id_key_buf = [0u8; 12];
let id_len = encode_db_next_id_key_fixed(ns_id, &mut db_id_key_buf);
batch.rm_meta(&db_id_key_buf[..id_len]);
batch.commit()?;
Ok(count)
}
pub struct WeDb<E: Engine> {
pub(crate) inner: Arc<WeDbInner<E>>,
}
impl<E: Engine> Clone for WeDb<E> {
#[inline(always)]
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
}
}
}
impl<E: Engine> fmt::Debug for WeDb<E> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("WeDb").finish()
}
}
impl<E: Engine> WeDb<E>
where
Error: From<E::Error>,
{
#[inline]
pub fn new(engine: E) -> Self {
init_version_counter();
let data = engine.partition(DATA).expect("open data partition");
let meta = engine.partition(META).expect("open meta partition");
let wedb = Self {
inner: Arc::new(WeDbInner {
engine: Arc::new(engine),
data,
meta,
ns_lock: Mutex::new(()),
expire_cursor: Mutex::new(ExpireCursors::default()),
}),
};
let _ = activate_db_impl::<E>(&wedb.inner.meta, 0, 0);
wedb
}
#[inline(always)]
pub fn data(&self) -> &E::Partition {
&self.inner.data
}
#[inline(always)]
pub fn meta(&self) -> &E::Partition {
&self.inner.meta
}
#[inline(always)]
pub fn engine(&self) -> &Arc<E> {
&self.inner.engine
}
#[inline(always)]
pub fn batch(&self) -> DbBatch<E> {
DbBatch::new(
self.inner.data.clone(),
self.inner.meta.clone(),
self.inner.engine.batch(),
)
}
#[inline(always)]
pub fn batch_with_capacity(&self, capacity: usize) -> DbBatch<E> {
DbBatch::new(
self.inner.data.clone(),
self.inner.meta.clone(),
self.inner.engine.batch_with_capacity(capacity),
)
}
#[inline]
pub fn iter(&self, begin: u64) -> Namespaces<'_, E> {
let mut buf = [0u8; 11];
let start_key: &[u8] = if begin == 0 {
CATALOG_PREFIX
} else {
let len = encode_catalog_ns_prefix_fixed(begin, &mut buf);
&buf[..len]
};
let iter = self
.inner
.meta
.range((Bound::Included(start_key), Bound::Unbounded));
Namespaces {
wedb: self.clone(),
iter,
last_emitted_ns: None,
}
}
#[inline]
pub fn ns(&self, id: impl IntoOptId) -> Result<Namespace<E>> {
let ns_id = match id.into_opt_id() {
Some(id) => id,
None => {
let new_id = next_namespace_id_impl::<E>(&self.inner.meta, &self.inner.ns_lock)?;
activate_db_impl::<E>(&self.inner.meta, new_id, 0)?;
new_id
}
};
Ok(Namespace {
id: ns_id,
inner: self.inner.clone(),
})
}
#[inline]
pub fn db(&self, id: impl IntoOptId) -> Result<Db<E>> {
self.ns(0)?.db(id)
}
#[inline]
pub fn rm(&self) -> Result<u64> {
let count = self.inner.data.approximate_len()? as u64;
self.inner.data.clear()?;
self.inner.meta.clear()?;
Ok(count)
}
#[inline]
pub fn persist(&self) -> Result<()> {
Ok(self.inner.engine.persist()?)
}
#[inline]
pub fn compact(&self) -> Result<()> {
self.inner.data.compact()?;
self.inner.meta.compact()?;
Ok(self.inner.engine.compact()?)
}
#[inline]
pub fn disk_space(&self) -> Result<u64> {
Ok(self.inner.engine.disk_space()?)
}
#[inline]
pub fn write_buffer_size(&self) -> u64 {
self.inner.engine.write_buffer_size()
}
#[inline]
pub fn cache_size(&self) -> u64 {
self.inner.engine.cache_size()
}
#[inline]
pub fn cache_capacity(&self) -> u64 {
self.inner.engine.cache_capacity()
}
#[inline]
pub fn outstanding_flushes(&self) -> usize {
self.inner.engine.outstanding_flushes()
}
#[inline]
pub fn is_kv_separated(&self) -> bool {
self.inner.data.is_kv_separated()
}
#[inline]
pub fn fragmented_blob_bytes(&self) -> u64 {
self.inner.data.fragmented_blob_bytes() + self.inner.meta.fragmented_blob_bytes()
}
#[inline]
pub fn table_count(&self) -> usize {
self.inner.data.table_count() + self.inner.meta.table_count()
}
#[inline]
pub fn blob_file_count(&self) -> usize {
self.inner.data.blob_file_count() + self.inner.meta.blob_file_count()
}
#[inline]
pub fn journal_count(&self) -> usize {
self.inner.engine.journal_count()
}
#[inline]
pub fn journal_disk_space(&self) -> Result<u64> {
Ok(self.inner.engine.journal_disk_space()?)
}
#[inline]
pub fn active_compactions(&self) -> usize {
self.inner.engine.active_compactions()
}
#[inline]
pub fn compactions_completed(&self) -> usize {
self.inner.engine.compactions_completed()
}
#[inline]
pub fn list_partitions(&self) -> Result<Vec<String>> {
Ok(self.inner.engine.list_partitions()?)
}
#[inline]
pub fn dbsize(&self) -> Result<usize> {
Ok(self.inner.data.approximate_len()?)
}
pub fn sweep_expired(&self, sample_limit: usize) -> Result<usize> {
let mut expired_count = 0;
let now_ms = current_now_ms();
let mut batch = self.batch();
let mut buf = Vec::with_capacity(64);
let (mut data_cur, mut meta_cur) = {
let guard = self.inner.expire_cursor.lock();
(guard.data_cursor.clone(), guard.meta_cursor.clone())
};
sweep_ks(self.data(), &mut data_cur, sample_limit, |k, v| {
let (expire_at, _) = decode_string_value(v);
if is_string_expired(expire_at, now_ms) {
batch.rm_data(k);
expired_count += 1;
}
Ok(())
})?;
sweep_ks(self.meta(), &mut meta_cur, sample_limit, |k, v| {
if let Some((kc, _, remain)) = KeyComposer::parse_scoped_prefix(k)
&& !remain.is_empty()
{
let meta_tag = remain[0];
let k_bytes = &remain[1..];
if let Some(tag) = KeyTag::from_u8(meta_tag)
&& tag.is_meta()
&& let Some(base_meta) = KeyMeta::decode(v)
&& base_meta.is_expired(now_ms)
{
batch.rm_meta(k);
cleanup_composite_data_raw(
self.data(),
self.meta(),
&kc,
meta_tag,
k_bytes,
&mut batch,
&mut buf,
)?;
expired_count += 1;
}
}
Ok(())
})?;
if expired_count > 0 {
batch.commit()?;
}
{
let mut guard = self.inner.expire_cursor.lock();
guard.data_cursor = data_cur;
guard.meta_cursor = meta_cur;
}
Ok(expired_count)
}
}