use std::{path::PathBuf, sync::Arc, time::Duration};
use object_store::ObjectStore;
use slatedb::{
Db, DbReader,
admin::AdminBuilder,
config::{
CheckpointOptions, CheckpointScope, DbReaderOptions, ObjectStoreCacheOptions, Settings,
},
};
use uuid::Uuid;
use crate::{
error::{Error, Result},
store::segment::TagSegmentExtractor,
};
const DEFAULT_FLUSH_INTERVAL: Duration = Duration::from_millis(100);
pub(crate) async fn create_checkpoint(db: &Db, lifetime: Option<Duration>) -> Result<Uuid> {
let created = db
.create_checkpoint(
CheckpointScope::All,
&CheckpointOptions {
lifetime,
..CheckpointOptions::default()
},
)
.await
.map_err(Error::from)?;
Ok(created.id)
}
pub(crate) struct StoreBuilder<'a> {
path: &'a str,
object_store: Arc<dyn ObjectStore>,
flush_interval: Duration,
cache_dir: Option<PathBuf>,
checkpoint: Option<Uuid>,
}
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,
checkpoint: 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) fn checkpoint(mut self, checkpoint: Option<Uuid>) -> Self {
self.checkpoint = checkpoint;
self
}
pub(crate) async fn open_writer(&self) -> Result<Db> {
let settings = self.settings();
Db::builder(self.path, Arc::clone(&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()
};
let mut builder = DbReader::builder(self.path, Arc::clone(&self.object_store))
.with_segment_extractor(Arc::new(TagSegmentExtractor))
.with_options(options);
if let Some(checkpoint) = self.checkpoint {
builder = builder.with_checkpoint_id(checkpoint);
}
builder.build().await.map_err(Error::from)
}
pub(crate) async fn delete_checkpoint(&self, checkpoint: Uuid) -> Result<()> {
AdminBuilder::new(self.path, Arc::clone(&self.object_store))
.build()
.delete_checkpoint(checkpoint)
.await
.map_err(Error::from)
}
pub(crate) async fn list_checkpoints(&self) -> Result<Vec<Uuid>> {
let checkpoints = AdminBuilder::new(self.path, Arc::clone(&self.object_store))
.build()
.list_checkpoints(None)
.await
.map_err(Error::from)?;
Ok(checkpoints.into_iter().map(|c| c.id).collect())
}
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);
}
}