use std::{
ops::Deref,
sync::{Arc, OnceLock},
};
use reifydb_codec::key::encoded::EncodedKey;
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
use reifydb_core::metrics::sample::MetricsSample;
use reifydb_core::{
common::CommitVersion, event::EventBus, lifecycle::watermark::EvictionWatermark,
metrics::collect::MetricsCollector,
};
use reifydb_runtime::{actor::system::ActorSystem, context::clock::Clock, shutdown::Shutdown, sync::rwlock::RwLock};
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
use reifydb_sqlite::SqliteTempPathGuard;
use reifydb_value::util::cowvec::CowVec;
use tracing::instrument;
use crate::{
CommitBufferConfig,
config::MultiStoreConfig,
flush::{ObjectPersistence, engine::FlushEngine},
tier::{
commit::buffer::MultiCommitBufferTier,
persistent::MultiPersistentTier,
read::{MultiReadBufferTier, ReadBufferShardMetrics},
},
};
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
use crate::{config::PersistentConfig, tier::read::ReadBufferConfig};
pub mod multi;
pub mod router;
use crate::Result;
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
struct SqlitePageCacheCollector {
persistent: MultiPersistentTier,
}
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
impl MetricsCollector for SqlitePageCacheCollector {
fn collect(&self, out: &mut Vec<MetricsSample>) {
let metrics = self.persistent.page_cache_metrics();
out.push(MetricsSample::bytes("sqlite::multi", "page_cache_used_bytes", metrics.used));
out.push(MetricsSample::counter("sqlite::multi", "page_cache_hit_count", metrics.hits.as_u64()));
out.push(MetricsSample::counter("sqlite::multi", "page_cache_miss_count", metrics.misses.as_u64()));
out.push(MetricsSample::count(
"sqlite::multi",
"page_cache_sampled_connections",
metrics.connections_sampled.as_u64(),
));
}
}
#[derive(Clone)]
pub struct StandardMultiStore(Arc<StandardMultiStoreInner>);
pub struct StandardMultiStoreInner {
pub(crate) commit: MultiCommitBufferTier,
pub(crate) persistent: Option<MultiPersistentTier>,
pub(crate) read: Option<MultiReadBufferTier>,
#[allow(dead_code)]
pub(crate) flush_engine: Option<Arc<FlushEngine>>,
#[allow(dead_code)]
pub(crate) row_settings_provider: Arc<OnceLock<Arc<dyn ObjectPersistence>>>,
#[allow(dead_code)]
pub(crate) eviction_watermark: Arc<RwLock<Option<Arc<dyn EvictionWatermark>>>>,
pub(crate) event_bus: EventBus,
}
impl StandardMultiStore {
#[instrument(name = "store::multi::new", level = "debug", skip(config), fields(
has_persistent = config.persistent.is_some(),
))]
pub fn new(config: MultiStoreConfig) -> Result<Self> {
let commit = config.commit.storage;
let row_settings_provider: Arc<OnceLock<Arc<dyn ObjectPersistence>>> = Arc::new(OnceLock::new());
let eviction_watermark: Arc<RwLock<Option<Arc<dyn EvictionWatermark>>>> = Arc::new(RwLock::new(None));
let read =
config.persistent.is_some().then(|| config.read.and_then(MultiReadBufferTier::new)).flatten();
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
let (persistent, flush_engine) = {
let persistent_config = config.persistent.clone();
let persistent = persistent_config.as_ref().map(|c| c.storage.clone());
let flush_engine = match (persistent.as_ref(), persistent_config.as_ref()) {
(Some(persistent_storage), Some(_)) => Some(Arc::new(FlushEngine::new(
commit.clone(),
persistent_storage.clone(),
row_settings_provider.clone(),
eviction_watermark.clone(),
read.clone(),
config.clock.clone(),
config.event_bus.clone(),
))),
_ => None,
};
(persistent, flush_engine)
};
#[cfg(not(all(feature = "sqlite", not(target_arch = "wasm32"))))]
let (persistent, flush_engine): (Option<MultiPersistentTier>, Option<Arc<FlushEngine>>) = {
let _ = config.persistent;
(None, None)
};
let read = persistent.as_ref().and(read);
Ok(Self(Arc::new(StandardMultiStoreInner {
commit,
persistent,
read,
flush_engine,
row_settings_provider,
eviction_watermark,
event_bus: config.event_bus,
})))
}
pub fn flush_engine(&self) -> Option<Arc<FlushEngine>> {
self.flush_engine.clone()
}
pub fn insert_read_key(&self, key: EncodedKey, version: CommitVersion, value: Option<CowVec<u8>>) {
if let Some(read) = &self.read {
read.insert(key, version, value);
}
}
pub fn invalidate_read_key(&self, key: &EncodedKey) {
if let Some(read) = &self.read {
read.invalidate(key);
}
}
pub fn clear_read(&self) {
if let Some(read) = &self.read {
read.clear();
}
}
pub fn set_row_settings_provider(&self, provider: Arc<dyn ObjectPersistence>) {
let _ = self.row_settings_provider.set(provider);
}
pub fn set_eviction_watermark(&self, watermark: Arc<dyn EvictionWatermark>) {
*self.eviction_watermark.write() = Some(watermark);
}
pub fn clear_eviction_watermark(&self) {
*self.eviction_watermark.write() = None;
}
pub fn commit(&self) -> &MultiCommitBufferTier {
&self.commit
}
pub fn event_bus(&self) -> &EventBus {
&self.event_bus
}
pub fn metrics_collectors(&self) -> Vec<Arc<dyn MetricsCollector>> {
let mut collectors: Vec<Arc<dyn MetricsCollector>> = Vec::new();
if let Some(read) = &self.read {
collectors.push(Arc::new(read.clone()));
}
collectors.push(Arc::new(self.commit.clone()));
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
if let Some(persistent) = &self.persistent {
collectors.push(Arc::new(SqlitePageCacheCollector {
persistent: persistent.clone(),
}));
}
collectors
}
pub fn persistent(&self) -> Option<&MultiPersistentTier> {
self.persistent.as_ref()
}
pub fn read_buffer_shard_metrics(&self) -> Vec<ReadBufferShardMetrics> {
self.read.as_ref().map(|read| read.shard_metrics()).unwrap_or_default()
}
pub fn flush_pending_blocking(&self) {
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
if let Some(engine) = self.flush_engine.as_ref() {
self.event_bus.wait_for_completion();
engine.flush_pending();
}
}
pub fn flush_all_blocking(&self) {
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
if let Some(engine) = self.flush_engine.as_ref() {
self.event_bus.wait_for_completion();
engine.flush_all();
}
}
}
impl Deref for StandardMultiStore {
type Target = StandardMultiStoreInner;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl Shutdown for StandardMultiStore {
fn shutdown(&self) {
if let Some(persistent) = self.persistent.as_ref() {
persistent.shutdown();
}
}
}
impl StandardMultiStore {
pub fn testing_memory() -> Self {
let clock = Clock::testing();
let actor_system = ActorSystem::testing(clock.clone());
let spawner = actor_system.spawner();
let event_bus = EventBus::new(&spawner);
Self::new(MultiStoreConfig {
commit: CommitBufferConfig {
storage: MultiCommitBufferTier::memory(),
},
persistent: None,
read: None,
retention: Default::default(),
merge_config: Default::default(),
event_bus,
spawner,
clock,
})
.unwrap()
}
pub fn testing_memory_with_eventbus(event_bus: EventBus) -> Self {
let clock = Clock::testing();
let actor_system = ActorSystem::testing(clock.clone());
let spawner = actor_system.spawner();
Self::new(MultiStoreConfig {
commit: CommitBufferConfig {
storage: MultiCommitBufferTier::memory(),
},
persistent: None,
read: None,
retention: Default::default(),
merge_config: Default::default(),
event_bus,
spawner,
clock,
})
.unwrap()
}
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
pub fn testing_memory_with_persistent_sqlite() -> (Self, SqliteTempPathGuard) {
Self::testing_memory_with_persistent_sqlite_read(ReadBufferConfig::default())
}
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
pub fn testing_memory_with_persistent_sqlite_read(read: ReadBufferConfig) -> (Self, SqliteTempPathGuard) {
let clock = Clock::testing();
let actor_system = ActorSystem::testing(clock.clone());
let spawner = actor_system.spawner();
let event_bus = EventBus::new(&spawner);
let (persistent, guard) = PersistentConfig::sqlite_in_memory();
let store = Self::new(MultiStoreConfig {
commit: CommitBufferConfig {
storage: MultiCommitBufferTier::memory(),
},
persistent: Some(persistent),
read: Some(read),
retention: Default::default(),
merge_config: Default::default(),
event_bus,
spawner,
clock,
})
.unwrap();
(store, guard)
}
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
pub fn testing_memory_with_persistent_sqlite_with_eventbus(event_bus: EventBus) -> (Self, SqliteTempPathGuard) {
let clock = Clock::testing();
let actor_system = ActorSystem::testing(clock.clone());
let spawner = actor_system.spawner();
let (persistent, guard) = PersistentConfig::sqlite_in_memory();
let store = Self::new(MultiStoreConfig {
commit: CommitBufferConfig {
storage: MultiCommitBufferTier::memory(),
},
persistent: Some(persistent),
read: Some(ReadBufferConfig::default()),
retention: Default::default(),
merge_config: Default::default(),
event_bus,
spawner,
clock,
})
.unwrap();
(store, guard)
}
}