use std::{path::PathBuf, sync::Arc, time::Duration};
use futures::{StreamExt, stream};
use object_store::ObjectStore;
use slatedb::{
BlockCachePolicy, CacheTarget, Db, DbReader, DbReaderMode,
admin::AdminBuilder,
config::{CheckpointOptions, CheckpointScope, DbReaderOptions, Settings, SstBlockSize},
};
use tracing::info;
use uuid::Uuid;
use crate::{
catalog::CachePreload,
error::{Error, Result},
store::{
cache, census,
handle::{ReadHandle, SCAN_READ_AHEAD_BYTES, ScanShape},
key, retry,
segment::TagSegmentExtractor,
},
};
const DEFAULT_FLUSH_INTERVAL: Duration = Duration::from_millis(100);
const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(10);
const WRITER_MANIFEST_POLL_INTERVAL: Duration = Duration::from_secs(5);
const MIN_FILTER_KEYS: u32 = 0;
const SST_BLOCK_SIZE: SstBlockSize = SstBlockSize::Block4Kib;
#[cfg(test)]
mod probe_tests;
#[cfg(test)]
mod probe_bench;
#[cfg(test)]
mod compaction_cache_tests;
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) fn preload_shortfall(store_bytes: u64, cache_size: Option<u64>) -> Option<u64> {
let cap = cache_size.unwrap_or(cache::DEFAULT_CACHE_DISK);
store_bytes
.checked_sub(cap)
.filter(|shortfall| *shortfall > 0)
}
pub(crate) struct StoreBuilder<'a> {
path: &'a str,
object_store: Arc<dyn ObjectStore>,
flush_interval: Option<Duration>,
poll_interval: Duration,
cache_dir: Option<PathBuf>,
cache_identity: crate::CacheIdentity,
cache_size: Option<u64>,
cache_memory: Option<u64>,
cache_preload: Option<CachePreload>,
cache_puts: bool,
cache_compaction_puts: bool,
checkpoint: Option<Uuid>,
warm_segments: Vec<census::SegmentSize>,
}
impl<'a> StoreBuilder<'a> {
pub(crate) fn new(path: &'a str, object_store: Arc<dyn ObjectStore>) -> Self {
Self {
path,
object_store: retry::Store::wrap(object_store),
flush_interval: Some(DEFAULT_FLUSH_INTERVAL),
poll_interval: DEFAULT_POLL_INTERVAL,
cache_dir: None,
cache_identity: crate::CacheIdentity::default(),
cache_size: None,
cache_memory: None,
cache_preload: None,
cache_puts: true,
cache_compaction_puts: true,
checkpoint: None,
warm_segments: Vec::new(),
}
}
pub(crate) fn poll_interval(mut self, poll_interval: Duration) -> Self {
self.poll_interval = poll_interval;
self
}
pub(crate) fn flush_interval(mut self, flush_interval: Duration) -> Self {
self.flush_interval = Some(flush_interval);
self
}
pub(crate) fn without_flush_timer(mut self) -> Self {
self.flush_interval = None;
self
}
pub(crate) fn cache_dir(mut self, cache_dir: Option<PathBuf>) -> Self {
self.cache_dir = cache_dir;
self
}
pub(crate) fn cache_identity(mut self, identity: Option<crate::CacheIdentity>) -> Self {
if let Some(identity) = identity {
self.cache_identity = identity;
}
self
}
pub(crate) fn cache_size(mut self, cache_size: Option<u64>) -> Self {
self.cache_size = cache_size;
self
}
pub(crate) fn cache_memory(mut self, cache_memory: Option<u64>) -> Self {
self.cache_memory = cache_memory;
self
}
pub(crate) fn cache_preload(mut self, cache_preload: Option<CachePreload>) -> Self {
self.cache_preload = cache_preload;
self
}
pub(crate) fn warm_segments(mut self, segments: Vec<census::SegmentSize>) -> Self {
self.warm_segments = segments;
self
}
pub(crate) fn cache_puts(mut self, cache_puts: bool) -> Self {
self.cache_puts = cache_puts;
self
}
pub(crate) fn cache_compaction_puts(mut self, enabled: bool) -> Self {
self.cache_compaction_puts = enabled;
self
}
pub(crate) fn checkpoint(mut self, checkpoint: Option<Uuid>) -> Self {
self.checkpoint = checkpoint;
self
}
pub(crate) async fn open_writer(&self) -> Result<(Db, Arc<cache::CacheCounters>)> {
let settings = self.settings();
let counters = cache::store_counters();
let mut builder = Db::builder(self.path, Arc::clone(&self.object_store))
.with_settings(settings)
.with_sst_block_size(SST_BLOCK_SIZE)
.with_segment_extractor(Arc::new(TagSegmentExtractor))
.with_filter_policies(super::index_filter::policies())
.with_block_cache_policy(self.block_cache_policy())
.with_metrics_recorder(cache::recorder(Arc::clone(&counters)));
if let Some(cache) =
cache::shared(&self.cache_config(), self.location(), Arc::clone(&counters)).await
{
builder = builder.with_db_cache(cache);
}
let db = builder.build().await.map_err(Error::from)?;
if let Some(preload) = self.cache_preload {
let tx = db
.begin(slatedb::IsolationLevel::Snapshot)
.await
.map_err(Error::from)?;
warm(preload, ReadHandle::Tx(&tx), &counters, &self.warm_segments).await;
tx.rollback();
}
Ok((db, counters))
}
pub(crate) async fn open_reader(&self) -> Result<(DbReader, Arc<cache::CacheCounters>)> {
let options = DbReaderOptions {
manifest_poll_interval: self.poll_interval,
..Default::default()
};
let counters = cache::store_counters();
let mut builder = DbReader::builder(self.path, Arc::clone(&self.object_store))
.with_segment_extractor(Arc::new(TagSegmentExtractor))
.with_filter_policies(super::index_filter::policies())
.with_metrics_recorder(cache::recorder(Arc::clone(&counters)))
.with_options(options);
if let Some(cache) =
cache::shared(&self.cache_config(), self.location(), Arc::clone(&counters)).await
{
builder = builder.with_db_cache(cache);
}
if let Some(checkpoint) = self.checkpoint {
builder = builder.with_reader_mode(DbReaderMode::Checkpoint(checkpoint));
}
let reader = builder.build().await.map_err(Error::from)?;
if let Some(preload) = self.cache_preload {
warm(
preload,
ReadHandle::Reader(&reader),
&counters,
&self.warm_segments,
)
.await;
}
Ok((reader, counters))
}
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: self.flush_interval,
manifest_poll_interval: WRITER_MANIFEST_POLL_INTERVAL,
min_filter_keys: MIN_FILTER_KEYS,
..Default::default()
}
}
fn location(&self) -> cache::StoreLocation {
cache::StoreLocation {
identity: self.cache_identity,
path: self.path.to_owned(),
}
}
fn cache_config(&self) -> cache::CacheConfig {
cache::CacheConfig {
memory: self.cache_memory,
dir: self.cache_dir.clone(),
disk_size: self.cache_size,
}
}
fn block_cache_policy(&self) -> BlockCachePolicy {
let targets: &[CacheTarget] = &[
CacheTarget::Index,
CacheTarget::Filters,
CacheTarget::Stats,
CacheTarget::data::<&[u8], _>(..),
];
BlockCachePolicy::default()
.with_flush_targets(if self.cache_puts { targets } else { &[] })
.with_compaction_output_targets(if self.cache_compaction_puts {
targets
} else {
&[]
})
}
}
pub(crate) fn warm_shape(bytes: Option<u64>) -> ScanShape {
let read_ahead = u64::try_from(SCAN_READ_AHEAD_BYTES).unwrap_or(u64::MAX);
match bytes {
Some(bytes) if bytes <= read_ahead => ScanShape::Probe,
_ => ScanShape::Seek,
}
}
async fn warm(
preload: CachePreload,
handle: ReadHandle<'_>,
counters: &cache::CacheCounters,
segments: &[census::SegmentSize],
) {
let metadata_only = [
key::Subspace::System,
key::Subspace::Current,
key::Subspace::History,
key::Subspace::Snapshot,
key::Subspace::Changelog,
key::Subspace::Index,
key::Subspace::Inline,
];
let whole = [
key::Subspace::System,
key::Subspace::Current,
key::Subspace::History,
key::Subspace::Snapshot,
key::Subspace::Changelog,
];
let before = counters.tally();
let (warmed, failed) = stream::iter(metadata_only.into_iter().map(|subspace| {
let deep = matches!(preload, CachePreload::All) && whole.contains(&subspace);
let prefix = key::subspace_prefix(subspace);
let shape = warm_shape(
segments
.iter()
.find(|segment| segment.prefix == prefix)
.map(|segment| segment.bytes),
);
warm_subspace(handle, subspace, deep, shape)
}))
.buffer_unordered(metadata_only.len())
.fold((0_usize, 0_usize), |(warmed, failed), result| async move {
match result {
Ok(()) => (warmed + 1, failed),
Err(_) => (warmed, failed + 1),
}
})
.await;
counters.record_preload(before, counters.tally(), failed);
info!(
warmed,
failed,
level = match preload {
CachePreload::L0 => "l0",
CachePreload::All => "all",
},
"warmed the cache"
);
}
async fn warm_subspace(
handle: ReadHandle<'_>,
subspace: key::Subspace,
deep: bool,
shape: ScanShape,
) -> Result<()> {
let prefix = key::subspace_prefix(subspace);
if !deep {
let mut iterator = handle.scan_prefix(&prefix, .., shape).await?;
iterator.next().await?;
return Ok(());
}
let split_kinds = match subspace {
key::Subspace::Current => key::SPLIT_SCAN_KINDS
.map(key::current_entity_kind_prefix)
.to_vec(),
key::Subspace::History => key::SPLIT_SCAN_KINDS
.map(key::history_entity_kind_prefix)
.to_vec(),
_ => Vec::new(),
};
handle
.scan_prefix_split(&prefix, &split_kinds, ScanShape::Probe)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use object_store::memory::InMemory;
use slatedb::IsolationLevel;
use super::*;
use crate::store::key::{self, Key, SysKey};
fn memory_store() -> Arc<dyn ObjectStore> {
Arc::new(InMemory::new())
}
#[test]
fn the_warm_shape_follows_the_segments_size() {
assert_eq!(warm_shape(Some(4 * 1024)), ScanShape::Probe);
assert_eq!(
warm_shape(Some(u64::try_from(SCAN_READ_AHEAD_BYTES).unwrap())),
ScanShape::Probe
);
assert_eq!(
warm_shape(Some(u64::try_from(SCAN_READ_AHEAD_BYTES).unwrap() + 1)),
ScanShape::Seek
);
assert_eq!(warm_shape(None), ScanShape::Seek);
}
#[test]
fn identical_display_names_do_not_share_catalog_cache_locations() {
let first = StoreBuilder::new("catalog", memory_store());
let second = StoreBuilder::new("catalog", memory_store());
assert_ne!(first.location(), second.location());
}
#[test]
fn explicit_catalog_cache_identities_share_only_the_same_namespace_and_path() {
let identity = crate::CacheIdentity::new("https://endpoint/bucket");
let location = |path, identity| {
StoreBuilder::new(path, memory_store())
.cache_identity(Some(identity))
.location()
};
assert_eq!(location("catalog", identity), location("catalog", identity));
assert_ne!(location("catalog", identity), location("other", identity));
assert_ne!(
location("catalog", identity),
location("catalog", crate::CacheIdentity::new("https://other/bucket"))
);
}
#[test]
fn the_sst_block_size_is_fixed_at_four_kibibytes() {
assert_eq!(SST_BLOCK_SIZE.as_bytes(), 4 * 1024);
}
#[test]
fn every_sst_carries_a_bloom_filter() {
let settings = StoreBuilder::new("s", memory_store()).settings();
assert_eq!(settings.min_filter_keys, 0);
}
#[test]
fn writers_poll_the_manifest_every_five_seconds() {
let settings = StoreBuilder::new("s", memory_store()).settings();
assert_eq!(settings.manifest_poll_interval, Duration::from_secs(5));
}
#[test]
fn writer_settings_pass_slatedb_validation() {
let settings = StoreBuilder::new("s", memory_store()).settings();
settings.validate().unwrap();
}
#[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()
.await
.unwrap()
.unwrap()
.await_durable()
.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()
.await_durable()
.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()
.await_durable()
.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_block_slot_with_a_device() {
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 _ = std::fs::remove_dir_all(&cache);
}
#[test]
fn the_cache_config_carries_what_was_asked_for() {
let object_store = memory_store();
let unset = StoreBuilder::new("s", Arc::clone(&object_store));
assert_eq!(unset.cache_config(), cache::CacheConfig::default());
let dir = std::path::PathBuf::from("/tmp/moraine-config-test");
let configured = StoreBuilder::new("s", object_store)
.cache_memory(Some(1 << 30))
.cache_dir(Some(dir.clone()))
.cache_size(Some(64 * 1024 * 1024));
assert_eq!(
configured.cache_config(),
cache::CacheConfig {
memory: Some(1 << 30),
dir: Some(dir),
disk_size: Some(64 * 1024 * 1024),
}
);
}
#[test]
fn a_preload_larger_than_the_cap_reports_its_shortfall() {
assert_eq!(preload_shortfall(3_400_000_000, Some(8_000_000_000)), None);
assert_eq!(preload_shortfall(64, Some(64)), None);
assert_eq!(preload_shortfall(100, Some(64)), Some(36));
assert_eq!(preload_shortfall(cache::DEFAULT_CACHE_DISK, None), None);
assert_eq!(
preload_shortfall(cache::DEFAULT_CACHE_DISK + 1, None),
Some(1)
);
}
#[test]
fn both_admissions_default_on_and_stay_independent() {
let object_store = memory_store();
let admitted = [
CacheTarget::Index,
CacheTarget::Filters,
CacheTarget::Stats,
CacheTarget::data::<&[u8], _>(..),
];
let unset = StoreBuilder::new("s", Arc::clone(&object_store));
assert_eq!(
unset.block_cache_policy(),
BlockCachePolicy::default()
.with_flush_targets(&admitted)
.with_compaction_output_targets(&admitted)
);
let no_flush = StoreBuilder::new("s", Arc::clone(&object_store)).cache_puts(false);
assert_eq!(
no_flush.block_cache_policy(),
BlockCachePolicy::default()
.with_flush_targets(&[])
.with_compaction_output_targets(&admitted)
);
let no_compaction =
StoreBuilder::new("s", Arc::clone(&object_store)).cache_compaction_puts(false);
assert_eq!(
no_compaction.block_cache_policy(),
BlockCachePolicy::default()
.with_flush_targets(&admitted)
.with_compaction_output_targets(&[])
);
let neither = StoreBuilder::new("s", object_store)
.cache_puts(false)
.cache_compaction_puts(false);
assert_eq!(
neither.block_cache_policy(),
BlockCachePolicy::default()
.with_flush_targets(&[])
.with_compaction_output_targets(&[])
);
}
#[tokio::test]
async fn a_preload_opens_a_readable_store_at_every_level() {
let object_store = memory_store();
let head = Key::Sys(SysKey::Head).encode();
let (db, _) = StoreBuilder::new("s", Arc::clone(&object_store))
.flush_interval(Duration::from_millis(1))
.open_writer()
.await
.unwrap();
db.put(&head, b"head").await.unwrap();
db.flush().await.unwrap();
db.close().await.unwrap();
for preload in [None, Some(CachePreload::L0), Some(CachePreload::All)] {
let (reader, _) = StoreBuilder::new("s", Arc::clone(&object_store))
.cache_preload(preload)
.open_reader()
.await
.unwrap();
assert_eq!(
reader.get(head.clone()).await.unwrap().unwrap().as_ref(),
b"head",
"preload {preload:?} left the store unreadable"
);
reader.close().await.unwrap();
}
}
}