reifydb_store_multi/store/

mod.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

use std::{
	ops::Deref,
	sync::{
		Arc, OnceLock,
		atomic::{AtomicU64, Ordering},
	},
};

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,
	interface::store::{EntryKind, storage_key},
	lifecycle::watermark::EvictionWatermark,
	metrics::collect::MetricsCollector,
};
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
use reifydb_filter::{actor::FilterActor, config::FilterConfig};
use reifydb_filter::{actor::FilterMessage, adaptive::FilterMetrics};
use reifydb_runtime::{
	actor::{mailbox::ActorRef, system::ActorSystem},
	context::clock::Clock,
	shutdown::Shutdown,
	sync::rwlock::RwLock,
};
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
use reifydb_sqlite::SqliteTempPathGuard;
use reifydb_store::metrics::PageCacheMetrics;
use reifydb_store_commit::store::{CommitStore, MultiCommitMetrics};
use reifydb_value::{count::Count, util::cowvec::CowVec};
use tracing::instrument;

#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
use crate::filter::source::MultiCurrentKeySource;
use crate::{
	CommitStoreConfig,
	config::MultiStoreConfig,
	flush::ObjectPersistence,
	tier::{
		persistent::MultiPersistentTier,
		point::{MultiPointShardMetrics, MultiPointTier},
		range::{MultiRangeShardMetrics, MultiRangeTier},
	},
};
#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
use crate::{
	config::PersistentConfig,
	flush::engine::FlushEngine,
	tier::{point::MultiPointConfig, range::MultiRangeConfig},
};

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, Copy, Debug, Default, PartialEq, Eq)]
pub struct MultiPersistentProbeMetrics {
	pub persistent_probes: Count,
	pub persistent_absent: Count,
}

#[derive(Clone)]
pub struct StandardMultiStore(Arc<StandardMultiStoreInner>);

pub struct StandardMultiStoreInner {
	pub(crate) commit: CommitStore,
	pub(crate) persistent: Option<MultiPersistentTier>,
	pub(crate) point: Option<MultiPointTier>,
	pub(crate) range: Option<MultiRangeTier>,

	#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
	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,

	pub(crate) persistent_probes: AtomicU64,
	pub(crate) persistent_absent: AtomicU64,

	pub(crate) filter: Option<ActorRef<FilterMessage>>,
}

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 point = config.persistent.is_some().then(|| config.point.and_then(MultiPointTier::new)).flatten();

		let range = config.persistent.is_some().then(|| config.range.and_then(MultiRangeTier::new)).flatten();

		#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
		let (persistent, flush_engine, filter) = {
			let persistent_config = config.persistent.clone();
			let persistent = persistent_config.as_ref().map(|c| c.storage.clone());
			let filter = persistent.as_ref().map(|tier| {
				let storage = tier.sqlite_storage().clone();
				let actor = FilterActor::spawn(&config.spawner);
				actor.send(FilterMessage::Register {
					filter: storage.filter().handle(),
					source: Box::new(MultiCurrentKeySource::new(storage)),
					config: FilterConfig::default(),
				})
				.expect("multi current filter source could not be registered");
				actor
			});
			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(),
						config.clock.clone(),
						config.event_bus.clone(),
					)
					.with_point(point.clone())
					.with_range(range.clone()),
				)),
				_ => None,
			};
			(persistent, flush_engine, filter)
		};

		#[cfg(not(all(feature = "sqlite", not(target_arch = "wasm32"))))]
		let (persistent, filter): (Option<MultiPersistentTier>, Option<ActorRef<FilterMessage>>) = {
			let _ = config.persistent;
			(None, None)
		};

		let point = persistent.as_ref().and(point);
		let range = persistent.as_ref().and(range);

		Ok(Self(Arc::new(StandardMultiStoreInner {
			commit,
			persistent,
			point,
			range,
			#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
			flush_engine,
			row_settings_provider,
			eviction_watermark,
			event_bus: config.event_bus,
			persistent_probes: AtomicU64::new(0),
			persistent_absent: AtomicU64::new(0),
			filter,
		})))
	}

	#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
	pub fn flush_engine(&self) -> Option<Arc<FlushEngine>> {
		self.flush_engine.clone()
	}

	pub fn insert_read_key(
		&self,
		table: EntryKind,
		key: EncodedKey,
		version: CommitVersion,
		value: Option<CowVec<u8>>,
	) {
		if let Some(range) = &self.range {
			range.insert(table, key.clone(), version, value.clone());
		}
		if let Some(point) = &self.point {
			let storage_key = storage_key(&key).1;
			point.insert(table, storage_key, key, version, value);
		}
	}

	pub fn invalidate_read_key(&self, table: EntryKind, key: &EncodedKey) {
		if let Some(range) = &self.range {
			range.invalidate(table, key);
		}
		if let Some(point) = &self.point {
			point.invalidate(table, storage_key(key).1, key);
		}
	}

	pub fn clear_read(&self) {
		if let Some(range) = &self.range {
			range.clear();
		}
		if let Some(point) = &self.point {
			point.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) -> &CommitStore {
		&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();
		collectors.push(Arc::new(self.commit.clone()));
		if let Some(point) = &self.point {
			collectors.push(Arc::new(point.clone()));
		}
		if let Some(range) = &self.range {
			collectors.push(Arc::new(range.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 point_shard_metrics(&self) -> Vec<MultiPointShardMetrics> {
		self.point.as_ref().map(|point| point.shard_metrics()).unwrap_or_default()
	}

	pub fn range_shard_metrics(&self) -> Vec<MultiRangeShardMetrics> {
		self.range.as_ref().map(|range| range.full_shard_metrics()).unwrap_or_default()
	}

	pub fn commit_metrics(&self) -> MultiCommitMetrics {
		self.commit.metrics()
	}

	pub fn persistent_page_cache_metrics(&self) -> Option<PageCacheMetrics> {
		self.persistent.as_ref().map(MultiPersistentTier::page_cache_metrics)
	}

	pub fn persistent_filter_metrics(&self) -> Option<FilterMetrics> {
		self.persistent.as_ref().map(|tier| tier.filter().metrics())
	}

	pub fn persistent_probe_metrics(&self) -> Option<MultiPersistentProbeMetrics> {
		self.persistent.as_ref().map(|_| MultiPersistentProbeMetrics {
			persistent_probes: Count::new(self.persistent_probes.load(Ordering::Relaxed)),
			persistent_absent: Count::new(self.persistent_absent.load(Ordering::Relaxed)),
		})
	}

	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(filter) = self.filter.as_ref() {
			let _ = filter.send(FilterMessage::Shutdown);
		}
		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: CommitStoreConfig {
				storage: CommitStore::new(),
			},
			persistent: None,
			point: None,
			range: 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: CommitStoreConfig {
				storage: CommitStore::new(),
			},
			persistent: None,
			point: None,
			range: 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_tiers(
			MultiPointConfig::testing(),
			MultiRangeConfig::testing(),
		)
	}

	#[cfg(all(feature = "sqlite", not(target_arch = "wasm32")))]
	pub fn testing_memory_with_persistent_sqlite_tiers(
		point: MultiPointConfig,
		range: MultiRangeConfig,
	) -> (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: CommitStoreConfig {
				storage: CommitStore::new(),
			},
			persistent: Some(persistent),
			point: Some(point),
			range: Some(range),
			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: CommitStoreConfig {
				storage: CommitStore::new(),
			},
			persistent: Some(persistent),
			point: Some(MultiPointConfig::testing()),
			range: Some(MultiRangeConfig::testing()),
			retention: Default::default(),
			merge_config: Default::default(),
			event_bus,
			spawner,
			clock,
		})
		.unwrap();
		(store, guard)
	}
}