use std::{path::PathBuf, sync::Arc, time::Duration};
use object_store::ObjectStore;
use slatedb::{
Db, DbReader,
config::{DbReaderOptions, ObjectStoreCacheOptions, Settings},
};
use crate::{
error::{Error, Result},
store::segment::TagSegmentExtractor,
};
const DEFAULT_FLUSH_INTERVAL: Duration = Duration::from_millis(100);
pub(crate) struct StoreBuilder<'a> {
path: &'a str,
object_store: Arc<dyn ObjectStore>,
flush_interval: Duration,
cache_dir: Option<PathBuf>,
}
impl<'a> StoreBuilder<'a> {
pub(crate) fn new(path: &'a str, object_store: Arc<dyn ObjectStore>) -> Self {
Self {
path,
object_store,
flush_interval: DEFAULT_FLUSH_INTERVAL,
cache_dir: None,
}
}
pub(crate) fn flush_interval(mut self, flush_interval: Duration) -> Self {
self.flush_interval = flush_interval;
self
}
pub(crate) fn cache_dir(mut self, cache_dir: Option<PathBuf>) -> Self {
self.cache_dir = cache_dir;
self
}
pub(crate) async fn open_writer(self) -> Result<Db> {
let settings = self.settings();
Db::builder(self.path, self.object_store)
.with_settings(settings)
.with_segment_extractor(Arc::new(TagSegmentExtractor))
.build()
.await
.map_err(Error::from)
}
pub(crate) async fn open_reader(self) -> Result<DbReader> {
let options = DbReaderOptions {
object_store_cache_options: self.cache_options(),
..Default::default()
};
DbReader::builder(self.path, self.object_store)
.with_segment_extractor(Arc::new(TagSegmentExtractor))
.with_options(options)
.build()
.await
.map_err(Error::from)
}
fn settings(&self) -> Settings {
Settings {
flush_interval: Some(self.flush_interval),
object_store_cache_options: self.cache_options(),
..Default::default()
}
}
fn cache_options(&self) -> ObjectStoreCacheOptions {
ObjectStoreCacheOptions {
root_folder: self.cache_dir.clone(),
..Default::default()
}
}
}
#[cfg(test)]
mod tests {
use object_store::memory::InMemory;
use slatedb::{IsolationLevel, config::WriteOptions};
use super::*;
use crate::store::key::{self, Key, SysKey};
fn memory_store() -> Arc<dyn ObjectStore> {
Arc::new(InMemory::new())
}
#[tokio::test]
async fn multi_subspace_transaction_and_prefix_scans() {
let db = StoreBuilder::new("test/store", memory_store())
.open_writer()
.await
.unwrap();
let head = Key::Sys(SysKey::Head).encode();
let snapshot = Key::Snapshot { snapshot_id: 1 }.encode();
let table = Key::current(key::EntityKey::Table { table_id: 7 }).encode();
let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
tx.put(&head, b"head").unwrap();
tx.put(&snapshot, b"snap").unwrap();
tx.put(&table, b"table").unwrap();
tx.commit_with_options(&WriteOptions {
await_durable: true,
..Default::default()
})
.await
.unwrap();
assert_eq!(db.get(&head).await.unwrap().unwrap().as_ref(), b"head");
let mut iter = db
.scan_prefix(key::subspace_prefix(key::Subspace::Current), ..)
.await
.unwrap();
let entry = iter.next().await.unwrap().unwrap();
assert_eq!(entry.key.as_ref(), table.as_slice());
assert!(iter.next().await.unwrap().is_none());
let mut iter = db
.scan_prefix(key::subspace_prefix(key::Subspace::Snapshot), ..)
.await
.unwrap();
let entry = iter.next().await.unwrap().unwrap();
assert_eq!(entry.key.as_ref(), snapshot.as_slice());
assert!(iter.next().await.unwrap().is_none());
db.close().await.unwrap();
}
#[tokio::test]
async fn zero_flush_interval_opens_a_working_store() {
let db = StoreBuilder::new("test/store", memory_store())
.flush_interval(Duration::ZERO)
.open_writer()
.await
.unwrap();
let head = Key::Sys(SysKey::Head).encode();
db.put(&head, b"head").await.unwrap();
assert_eq!(db.get(&head).await.unwrap().unwrap().as_ref(), b"head");
db.close().await.unwrap();
}
#[tokio::test]
async fn explicit_flush_interval_opens_a_working_store() {
let db = StoreBuilder::new("test/store", memory_store())
.flush_interval(Duration::from_millis(1))
.open_writer()
.await
.unwrap();
let head = Key::Sys(SysKey::Head).encode();
db.put(&head, b"head").await.unwrap();
assert_eq!(db.get(&head).await.unwrap().unwrap().as_ref(), b"head");
db.close().await.unwrap();
}
#[tokio::test]
async fn reopen_without_extractor_is_refused() {
let object_store = memory_store();
let db = StoreBuilder::new("test/store", object_store.clone())
.open_writer()
.await
.unwrap();
db.put(&Key::Sys(SysKey::Head).encode(), b"head")
.await
.unwrap();
db.close().await.unwrap();
let bare = Db::builder("test/store", object_store).build().await;
assert!(bare.is_err(), "unsegmented reopen must be refused");
}
#[tokio::test]
async fn cache_dir_backs_the_reader_with_an_on_disk_cache() {
let object_store = memory_store();
let cache = std::env::temp_dir().join(format!("moraine-cache-test-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&cache);
let head = Key::Sys(SysKey::Head).encode();
let db = StoreBuilder::new("s", object_store.clone())
.flush_interval(Duration::from_millis(1))
.cache_dir(Some(cache.clone()))
.open_writer()
.await
.unwrap();
db.put(&head, b"head").await.unwrap();
db.close().await.unwrap();
let reader = StoreBuilder::new("s", object_store)
.cache_dir(Some(cache.clone()))
.open_reader()
.await
.unwrap();
assert_eq!(reader.get(head).await.unwrap().unwrap().as_ref(), b"head");
reader.close().await.unwrap();
let populated = std::fs::read_dir(&cache).is_ok_and(|mut entries| entries.next().is_some());
assert!(populated, "expected an on-disk cache under {cache:?}");
let _ = std::fs::remove_dir_all(&cache);
}
}