1use std::collections::HashSet;
16use std::fmt::Debug;
17use std::sync::{Arc, LazyLock};
18use std::time::Duration;
19
20use enum_as_inner::EnumAsInner;
21use foyer::{
22 BlockEngineConfig, CacheBuilder, DeviceBuilder, FifoPicker, FsDeviceBuilder, HybridCacheBuilder,
23};
24use futures::FutureExt;
25use futures::future::BoxFuture;
26use mixtrics::registry::prometheus::PrometheusMetricsRegistry;
27use risingwave_common::catalog::TableId;
28use risingwave_common::config::Role;
29use risingwave_common::config::storage::FileCacheRuntimeConfig;
30use risingwave_common::license::Feature;
31use risingwave_common::monitor::GLOBAL_METRICS_REGISTRY;
32use risingwave_common_service::RpcNotificationClient;
33use risingwave_hummock_sdk::{HummockEpoch, HummockSstableObjectId, SyncResult};
34use risingwave_object_store::object::build_remote_object_store;
35use thiserror_ext::AsReport;
36
37use crate::StateStore;
38use crate::compaction_catalog_manager::{CompactionCatalogManager, RemoteTableAccessor};
39use crate::error::StorageResult;
40use crate::hummock::all::AllRecentFilter;
41use crate::hummock::hummock_meta_client::MonitoredHummockMetaClient;
42use crate::hummock::none::NoneRecentFilter;
43use crate::hummock::sharded::ShardedRecentFilter;
44use crate::hummock::simple::SimpleRecentFilter;
45use crate::hummock::{
46 Block, BlockCacheEventListener, HummockError, HummockStorage, SstableBlockIndex,
47 SstableMetaCacheValue, SstableStore, SstableStoreConfig,
48};
49use crate::memory::MemoryStateStore;
50use crate::memory::sled::SledStateStore;
51use crate::monitor::{
52 CompactorMetrics, HummockStateStoreMetrics, MonitoredStateStore, MonitoredStorageMetrics,
53 ObjectStoreMetrics,
54};
55use crate::opts::StorageOpts;
56
57fn build_file_cache_spawner(
58 name: &str,
59 config: &FileCacheRuntimeConfig,
60) -> StorageResult<foyer::Spawner> {
61 let config = match config {
62 FileCacheRuntimeConfig::Disabled => return Ok(foyer::Spawner::current()),
63 FileCacheRuntimeConfig::Unified(config) => config.clone(),
64 FileCacheRuntimeConfig::Separated { .. } => {
65 return Err(HummockError::other(format!(
66 "{name} runtime_config.Separated is unsupported with Foyer 0.22; use runtime_config.Unified instead"
67 ))
68 .into());
69 }
70 };
71
72 let mut builder = tokio::runtime::Builder::new_multi_thread();
73 #[cfg(madsim)]
74 let _ = &config;
75 #[cfg(not(madsim))]
76 if config.worker_threads != 0 {
77 builder.worker_threads(config.worker_threads);
78 }
79 #[cfg(not(madsim))]
80 if config.max_blocking_threads != 0 {
81 builder.max_blocking_threads(config.max_blocking_threads);
82 }
83 builder.thread_name(format!("{name}-unified"));
84 let runtime = builder.enable_all().build().map_err(|error| {
85 HummockError::other(format!(
86 "failed to build {name} dedicated runtime: {}",
87 error.as_report()
88 ))
89 })?;
90 Ok(runtime.into())
91}
92
93#[cfg(test)]
94mod tests {
95 use risingwave_common::config::storage::FileCacheTokioRuntimeConfig;
96
97 use super::*;
98
99 #[test]
100 fn test_build_file_cache_spawner_rejects_separated_runtime() {
101 let runtime_options = FileCacheTokioRuntimeConfig {
102 worker_threads: 1,
103 max_blocking_threads: 1,
104 };
105 let error = build_file_cache_spawner(
106 "foyer.test",
107 &FileCacheRuntimeConfig::Separated {
108 read_runtime_options: runtime_options.clone(),
109 write_runtime_options: runtime_options,
110 },
111 )
112 .unwrap_err()
113 .to_string();
114
115 assert!(error.contains("foyer.test runtime_config.Separated"));
116 assert!(error.contains("use runtime_config.Unified instead"));
117 }
118
119 #[cfg(not(madsim))]
120 #[tokio::test]
121 async fn test_build_file_cache_spawner_runtime_modes() {
122 let current_runtime_id = tokio::runtime::Handle::current().id();
123
124 let disabled =
125 build_file_cache_spawner("foyer.test", &FileCacheRuntimeConfig::Disabled).unwrap();
126 let disabled_runtime_id = disabled
127 .spawn(async { tokio::runtime::Handle::current().id() })
128 .await
129 .unwrap();
130 assert_eq!(disabled_runtime_id, current_runtime_id);
131
132 let unified = build_file_cache_spawner(
133 "foyer.test",
134 &FileCacheRuntimeConfig::Unified(FileCacheTokioRuntimeConfig {
135 worker_threads: 1,
136 max_blocking_threads: 1,
137 }),
138 )
139 .unwrap();
140 let (unified_runtime_id, thread_name) = unified
141 .spawn(async {
142 (
143 tokio::runtime::Handle::current().id(),
144 std::thread::current().name().map(str::to_owned),
145 )
146 })
147 .await
148 .unwrap();
149
150 assert_ne!(unified_runtime_id, current_runtime_id);
151 assert_eq!(thread_name.as_deref(), Some("foyer.test-unified"));
152 }
153}
154
155static FOYER_METRICS_REGISTRY: LazyLock<Box<PrometheusMetricsRegistry>> = LazyLock::new(|| {
156 Box::new(PrometheusMetricsRegistry::new(
157 GLOBAL_METRICS_REGISTRY.clone(),
158 ))
159});
160
161mod opaque_type {
162 use super::*;
163
164 pub type HummockStorageType = impl StateStore + AsHummock;
165 pub type MemoryStateStoreType = impl StateStore + AsHummock;
166 pub type SledStateStoreType = impl StateStore + AsHummock;
167
168 #[define_opaque(MemoryStateStoreType)]
169 pub fn in_memory(state_store: MemoryStateStore) -> MemoryStateStoreType {
170 may_dynamic_dispatch(state_store)
171 }
172
173 #[define_opaque(HummockStorageType)]
174 pub fn hummock(state_store: HummockStorage) -> HummockStorageType {
175 may_dynamic_dispatch(may_verify(state_store))
176 }
177
178 #[define_opaque(SledStateStoreType)]
179 pub fn sled(state_store: SledStateStore) -> SledStateStoreType {
180 may_dynamic_dispatch(state_store)
181 }
182}
183pub use opaque_type::{HummockStorageType, MemoryStateStoreType, SledStateStoreType};
184use opaque_type::{hummock, in_memory, sled};
185
186#[cfg(feature = "hm-trace")]
187type Monitored<S> = MonitoredStateStore<crate::monitor::traced_store::TracedStateStore<S>>;
188
189#[cfg(not(feature = "hm-trace"))]
190type Monitored<S> = MonitoredStateStore<S>;
191
192fn monitored<S: StateStore>(
193 state_store: S,
194 storage_metrics: Arc<MonitoredStorageMetrics>,
195) -> Monitored<S> {
196 let inner = {
197 #[cfg(feature = "hm-trace")]
198 {
199 crate::monitor::traced_store::TracedStateStore::new_global(state_store)
200 }
201 #[cfg(not(feature = "hm-trace"))]
202 {
203 state_store
204 }
205 };
206 inner.monitored(storage_metrics)
207}
208
209fn inner<S>(state_store: &Monitored<S>) -> &S {
210 let inner = state_store.inner();
211 {
212 #[cfg(feature = "hm-trace")]
213 {
214 inner.inner()
215 }
216 #[cfg(not(feature = "hm-trace"))]
217 {
218 inner
219 }
220 }
221}
222
223#[derive(Clone, EnumAsInner)]
225#[expect(clippy::enum_variant_names)]
226pub enum StateStoreImpl {
227 HummockStateStore(Monitored<HummockStorageType>),
236 MemoryStateStore(Monitored<MemoryStateStoreType>),
241 SledStateStore(Monitored<SledStateStoreType>),
242}
243
244fn may_dynamic_dispatch(state_store: impl StateStore + AsHummock) -> impl StateStore + AsHummock {
245 #[cfg(not(debug_assertions))]
246 {
247 state_store
248 }
249 #[cfg(debug_assertions)]
250 {
251 use crate::store_impl::dyn_state_store::StateStorePointer;
252 StateStorePointer(Arc::new(state_store) as _)
253 }
254}
255
256fn may_verify(state_store: impl StateStore + AsHummock) -> impl StateStore + AsHummock {
257 #[cfg(not(debug_assertions))]
258 {
259 state_store
260 }
261 #[cfg(debug_assertions)]
262 {
263 use std::marker::PhantomData;
264
265 use risingwave_common::util::env_var::env_var_is_true;
266 use tracing::info;
267
268 use crate::store_impl::verify::VerifyStateStore;
269
270 let expected = if env_var_is_true("ENABLE_STATE_STORE_VERIFY") {
271 info!("enable verify state store");
272 Some(SledStateStore::new_temp())
273 } else {
274 info!("verify state store is not enabled");
275 None
276 };
277 VerifyStateStore {
278 actual: state_store,
279 expected,
280 _phantom: PhantomData::<()>,
281 }
282 }
283}
284
285impl StateStoreImpl {
286 fn in_memory(
287 state_store: MemoryStateStore,
288 storage_metrics: Arc<MonitoredStorageMetrics>,
289 ) -> Self {
290 Self::MemoryStateStore(monitored(in_memory(state_store), storage_metrics))
292 }
293
294 pub fn hummock(
295 state_store: HummockStorage,
296 storage_metrics: Arc<MonitoredStorageMetrics>,
297 ) -> Self {
298 Self::HummockStateStore(monitored(hummock(state_store), storage_metrics))
300 }
301
302 pub fn sled(
303 state_store: SledStateStore,
304 storage_metrics: Arc<MonitoredStorageMetrics>,
305 ) -> Self {
306 Self::SledStateStore(monitored(sled(state_store), storage_metrics))
307 }
308
309 pub fn shared_in_memory_store(storage_metrics: Arc<MonitoredStorageMetrics>) -> Self {
310 Self::in_memory(MemoryStateStore::shared(), storage_metrics)
311 }
312
313 pub fn for_test() -> Self {
314 Self::in_memory(
315 MemoryStateStore::new(),
316 Arc::new(MonitoredStorageMetrics::unused()),
317 )
318 }
319
320 pub fn as_hummock(&self) -> Option<&HummockStorage> {
321 match self {
322 StateStoreImpl::HummockStateStore(hummock) => {
323 Some(inner(hummock).as_hummock().expect("should be hummock"))
324 }
325 _ => None,
326 }
327 }
328}
329
330impl Debug for StateStoreImpl {
331 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
332 match self {
333 StateStoreImpl::HummockStateStore(_) => write!(f, "HummockStateStore"),
334 StateStoreImpl::MemoryStateStore(_) => write!(f, "MemoryStateStore"),
335 StateStoreImpl::SledStateStore(_) => write!(f, "SledStateStore"),
336 }
337 }
338}
339
340#[macro_export]
341macro_rules! dispatch_state_store {
342 ($impl:expr, $store:ident, $body:tt) => {{
343 use $crate::store_impl::StateStoreImpl;
344
345 match $impl {
346 StateStoreImpl::MemoryStateStore($store) => {
347 #[cfg(debug_assertions)]
350 {
351 $body
352 }
353 #[cfg(not(debug_assertions))]
354 {
355 let _store = $store;
356 unimplemented!("memory state store should never be used in release mode");
357 }
358 }
359
360 StateStoreImpl::SledStateStore($store) => {
361 #[cfg(debug_assertions)]
364 {
365 $body
366 }
367 #[cfg(not(debug_assertions))]
368 {
369 let _store = $store;
370 unimplemented!("sled state store should never be used in release mode");
371 }
372 }
373
374 StateStoreImpl::HummockStateStore($store) => $body,
375 }
376 }};
377}
378
379#[cfg(any(debug_assertions, test, feature = "test"))]
380pub mod verify {
381 use std::fmt::Debug;
382 use std::future::Future;
383 use std::marker::PhantomData;
384 use std::ops::Deref;
385 use std::sync::Arc;
386
387 use bytes::Bytes;
388 use risingwave_common::array::VectorRef;
389 use risingwave_common::bitmap::Bitmap;
390 use risingwave_common::hash::VirtualNode;
391 use risingwave_hummock_sdk::HummockReadEpoch;
392 use risingwave_hummock_sdk::key::{FullKey, TableKey, TableKeyRange};
393 use tracing::log::warn;
394
395 use crate::error::StorageResult;
396 use crate::hummock::HummockStorage;
397 use crate::store::*;
398 use crate::store_impl::AsHummock;
399
400 #[expect(dead_code)]
401 fn assert_result_eq<Item: PartialEq + Debug, E>(
402 first: &std::result::Result<Item, E>,
403 second: &std::result::Result<Item, E>,
404 ) {
405 match (first, second) {
406 (Ok(first), Ok(second)) => {
407 if first != second {
408 warn!("result different: {:?} {:?}", first, second);
409 }
410 assert_eq!(first, second);
411 }
412 (Err(_), Err(_)) => {}
413 _ => {
414 warn!("one success and one failed");
415 panic!("result not equal");
416 }
417 }
418 }
419
420 #[derive(Clone)]
421 pub struct VerifyStateStore<A, E, T = ()> {
422 pub actual: A,
423 pub expected: Option<E>,
424 pub _phantom: PhantomData<T>,
425 }
426
427 impl<A: AsHummock, E: AsHummock> AsHummock for VerifyStateStore<A, E> {
428 fn as_hummock(&self) -> Option<&HummockStorage> {
429 self.actual.as_hummock()
430 }
431 }
432
433 impl<A: StateStoreGet, E: StateStoreGet> StateStoreGet for VerifyStateStore<A, E> {
434 async fn on_key_value<'a, O: Send + 'a>(
435 &'a self,
436 key: TableKey<Bytes>,
437 read_options: ReadOptions,
438 on_key_value_fn: impl KeyValueFn<'a, O>,
439 ) -> StorageResult<Option<O>> {
440 let actual: Option<(FullKey<Bytes>, Bytes)> = self
441 .actual
442 .on_key_value(key.clone(), read_options.clone(), |key, value| {
443 Ok((key.copy_into(), Bytes::copy_from_slice(value)))
444 })
445 .await?;
446 if let Some(expected) = &self.expected {
447 let expected: Option<(FullKey<Bytes>, Bytes)> = expected
448 .on_key_value(key, read_options, |key, value| {
449 Ok((key.copy_into(), Bytes::copy_from_slice(value)))
450 })
451 .await?;
452 assert_eq!(
453 actual
454 .as_ref()
455 .map(|item| (item.0.epoch_with_gap.pure_epoch(), item)),
456 expected
457 .as_ref()
458 .map(|item| (item.0.epoch_with_gap.pure_epoch(), item))
459 );
460 }
461
462 actual
463 .map(|(key, value)| on_key_value_fn(key.to_ref(), value.as_ref()))
464 .transpose()
465 }
466 }
467
468 impl<A: StateStoreReadVector, E: StateStoreReadVector> StateStoreReadVector
469 for VerifyStateStore<A, E>
470 {
471 fn nearest<'a, O: Send + 'a>(
472 &'a self,
473 vec: VectorRef<'a>,
474 options: VectorNearestOptions,
475 on_nearest_item_fn: impl OnNearestItemFn<'a, O>,
476 ) -> impl StorageFuture<'a, Vec<O>> {
477 self.actual.nearest(vec, options, on_nearest_item_fn)
478 }
479 }
480
481 impl<A: StateStoreRead, E: StateStoreRead> StateStoreRead for VerifyStateStore<A, E> {
482 type Iter = impl StateStoreReadIter;
483 type RevIter = impl StateStoreReadIter;
484
485 fn read_epoch(&self) -> Option<u64> {
486 self.actual.read_epoch()
487 }
488
489 #[expect(clippy::manual_async_fn)]
492 fn iter(
493 &self,
494 key_range: TableKeyRange,
495
496 read_options: ReadOptions,
497 ) -> impl Future<Output = StorageResult<Self::Iter>> + '_ {
498 async move {
499 let actual = self
500 .actual
501 .iter(key_range.clone(), read_options.clone())
502 .await?;
503 let expected = if let Some(expected) = &self.expected {
504 Some(expected.iter(key_range, read_options).await?)
505 } else {
506 None
507 };
508
509 Ok(verify_iter::<StateStoreKeyedRow>(actual, expected))
510 }
511 }
512
513 #[expect(clippy::manual_async_fn)]
514 fn rev_iter(
515 &self,
516 key_range: TableKeyRange,
517
518 read_options: ReadOptions,
519 ) -> impl Future<Output = StorageResult<Self::RevIter>> + '_ {
520 async move {
521 let actual = self
522 .actual
523 .rev_iter(key_range.clone(), read_options.clone())
524 .await?;
525 let expected = if let Some(expected) = &self.expected {
526 Some(expected.rev_iter(key_range, read_options).await?)
527 } else {
528 None
529 };
530
531 Ok(verify_iter::<StateStoreKeyedRow>(actual, expected))
532 }
533 }
534 }
535
536 impl<A: StateStoreReadLog, E: StateStoreReadLog> StateStoreReadLog for VerifyStateStore<A, E> {
537 type ChangeLogIter = impl StateStoreReadChangeLogIter;
538
539 async fn next_epoch(&self, epoch: u64, options: NextEpochOptions) -> StorageResult<u64> {
540 let actual = self.actual.next_epoch(epoch, options.clone()).await?;
541 if let Some(expected) = &self.expected {
542 assert_eq!(actual, expected.next_epoch(epoch, options).await?);
543 }
544 Ok(actual)
545 }
546
547 async fn iter_log(
548 &self,
549 epoch_range: (u64, u64),
550 key_range: TableKeyRange,
551 options: ReadLogOptions,
552 ) -> StorageResult<Self::ChangeLogIter> {
553 let actual = self
554 .actual
555 .iter_log(epoch_range, key_range.clone(), options.clone())
556 .await?;
557 let expected = if let Some(expected) = &self.expected {
558 Some(expected.iter_log(epoch_range, key_range, options).await?)
559 } else {
560 None
561 };
562
563 Ok(verify_iter::<StateStoreReadLogItem>(actual, expected))
564 }
565 }
566
567 impl<A: StateStoreIter<T>, E: StateStoreIter<T>, T: IterItem> StateStoreIter<T>
568 for VerifyStateStore<A, E, T>
569 where
570 for<'a> T::ItemRef<'a>: PartialEq + Debug,
571 {
572 async fn try_next(&mut self) -> StorageResult<Option<T::ItemRef<'_>>> {
573 let actual = self.actual.try_next().await?;
574 if let Some(expected) = self.expected.as_mut() {
575 let expected = expected.try_next().await?;
576 assert_eq!(actual, expected);
577 }
578 Ok(actual)
579 }
580 }
581
582 fn verify_iter<T: IterItem>(
583 actual: impl StateStoreIter<T>,
584 expected: Option<impl StateStoreIter<T>>,
585 ) -> impl StateStoreIter<T>
586 where
587 for<'a> T::ItemRef<'a>: PartialEq + Debug,
588 {
589 VerifyStateStore {
590 actual,
591 expected,
592 _phantom: PhantomData::<T>,
593 }
594 }
595
596 impl<A: LocalStateStore, E: LocalStateStore> LocalStateStore for VerifyStateStore<A, E> {
597 type FlushedSnapshotReader =
598 VerifyStateStore<A::FlushedSnapshotReader, E::FlushedSnapshotReader>;
599
600 type Iter<'a> = impl StateStoreIter + 'a;
601 type RevIter<'a> = impl StateStoreIter + 'a;
602
603 #[expect(clippy::manual_async_fn)]
604 fn iter(
605 &self,
606 key_range: TableKeyRange,
607 read_options: ReadOptions,
608 ) -> impl Future<Output = StorageResult<Self::Iter<'_>>> + Send + '_ {
609 async move {
610 let actual = self
611 .actual
612 .iter(key_range.clone(), read_options.clone())
613 .await?;
614 let expected = if let Some(expected) = &self.expected {
615 Some(expected.iter(key_range, read_options).await?)
616 } else {
617 None
618 };
619
620 Ok(verify_iter::<StateStoreKeyedRow>(actual, expected))
621 }
622 }
623
624 #[expect(clippy::manual_async_fn)]
625 fn rev_iter(
626 &self,
627 key_range: TableKeyRange,
628 read_options: ReadOptions,
629 ) -> impl Future<Output = StorageResult<Self::RevIter<'_>>> + Send + '_ {
630 async move {
631 let actual = self
632 .actual
633 .rev_iter(key_range.clone(), read_options.clone())
634 .await?;
635 let expected = if let Some(expected) = &self.expected {
636 Some(expected.rev_iter(key_range, read_options).await?)
637 } else {
638 None
639 };
640
641 Ok(verify_iter::<StateStoreKeyedRow>(actual, expected))
642 }
643 }
644
645 fn insert(
646 &mut self,
647 key: TableKey<Bytes>,
648 new_val: Bytes,
649 old_val: Option<Bytes>,
650 ) -> StorageResult<()> {
651 if let Some(expected) = &mut self.expected {
652 expected.insert(key.clone(), new_val.clone(), old_val.clone())?;
653 }
654 self.actual.insert(key, new_val, old_val)?;
655
656 Ok(())
657 }
658
659 fn delete(&mut self, key: TableKey<Bytes>, old_val: Bytes) -> StorageResult<()> {
660 if let Some(expected) = &mut self.expected {
661 expected.delete(key.clone(), old_val.clone())?;
662 }
663 self.actual.delete(key, old_val)?;
664 Ok(())
665 }
666
667 async fn update_vnode_bitmap(&mut self, vnodes: Arc<Bitmap>) -> StorageResult<Arc<Bitmap>> {
668 let ret = self.actual.update_vnode_bitmap(vnodes.clone()).await?;
669 if let Some(expected) = &mut self.expected {
670 assert_eq!(ret, expected.update_vnode_bitmap(vnodes).await?);
671 }
672 Ok(ret)
673 }
674
675 fn get_table_watermark(&self, vnode: VirtualNode) -> Option<Bytes> {
676 let ret = self.actual.get_table_watermark(vnode);
677 if let Some(expected) = &self.expected {
678 assert_eq!(ret, expected.get_table_watermark(vnode));
679 }
680 ret
681 }
682
683 fn new_flushed_snapshot_reader(&self) -> Self::FlushedSnapshotReader {
684 VerifyStateStore {
685 actual: self.actual.new_flushed_snapshot_reader(),
686 expected: self.expected.as_ref().map(E::new_flushed_snapshot_reader),
687 _phantom: Default::default(),
688 }
689 }
690 }
691
692 impl<A: StateStoreWriteEpochControl, E: StateStoreWriteEpochControl> StateStoreWriteEpochControl
693 for VerifyStateStore<A, E>
694 {
695 async fn flush(&mut self) -> StorageResult<usize> {
696 if let Some(expected) = &mut self.expected {
697 expected.flush().await?;
698 }
699 self.actual.flush().await
700 }
701
702 async fn try_flush(&mut self) -> StorageResult<()> {
703 if let Some(expected) = &mut self.expected {
704 expected.try_flush().await?;
705 }
706 self.actual.try_flush().await
707 }
708
709 async fn init(&mut self, options: InitOptions) -> StorageResult<()> {
710 self.actual.init(options.clone()).await?;
711 if let Some(expected) = &mut self.expected {
712 expected.init(options).await?;
713 }
714 Ok(())
715 }
716
717 fn seal_current_epoch(&mut self, next_epoch: u64, opts: SealCurrentEpochOptions) {
718 if let Some(expected) = &mut self.expected {
719 expected.seal_current_epoch(next_epoch, opts.clone());
720 }
721 self.actual.seal_current_epoch(next_epoch, opts);
722 }
723 }
724
725 impl<A: StateStore, E: StateStore> StateStore for VerifyStateStore<A, E> {
726 type Local = VerifyStateStore<A::Local, E::Local>;
727 type ReadSnapshot = VerifyStateStore<A::ReadSnapshot, E::ReadSnapshot>;
728 type VectorWriter = A::VectorWriter;
729
730 fn try_wait_epoch(
731 &self,
732 epoch: HummockReadEpoch,
733 options: TryWaitEpochOptions,
734 ) -> impl Future<Output = StorageResult<()>> + Send + '_ {
735 self.actual.try_wait_epoch(epoch, options)
736 }
737
738 async fn new_local(&self, option: NewLocalOptions) -> Self::Local {
739 let expected = if let Some(expected) = &self.expected {
740 Some(expected.new_local(option.clone()).await)
741 } else {
742 None
743 };
744 VerifyStateStore {
745 actual: self.actual.new_local(option).await,
746 expected,
747 _phantom: PhantomData::<()>,
748 }
749 }
750
751 async fn new_read_snapshot(
752 &self,
753 epoch: HummockReadEpoch,
754 options: NewReadSnapshotOptions,
755 ) -> StorageResult<Self::ReadSnapshot> {
756 let expected = if let Some(expected) = &self.expected {
757 Some(expected.new_read_snapshot(epoch, options).await?)
758 } else {
759 None
760 };
761 Ok(VerifyStateStore {
762 actual: self.actual.new_read_snapshot(epoch, options).await?,
763 expected,
764 _phantom: PhantomData::<()>,
765 })
766 }
767
768 fn new_vector_writer(
769 &self,
770 options: NewVectorWriterOptions,
771 ) -> impl Future<Output = Self::VectorWriter> + Send + '_ {
772 self.actual.new_vector_writer(options)
773 }
774 }
775
776 impl<A, E> Deref for VerifyStateStore<A, E> {
777 type Target = A;
778
779 fn deref(&self) -> &Self::Target {
780 &self.actual
781 }
782 }
783}
784
785impl StateStoreImpl {
786 #[cfg_attr(not(target_os = "linux"), allow(unused_variables))]
787 #[expect(clippy::borrowed_box)]
788 pub async fn new(
789 s: &str,
790 role: Role,
791 opts: Arc<StorageOpts>,
792 hummock_meta_client: Arc<MonitoredHummockMetaClient>,
793 state_store_metrics: Arc<HummockStateStoreMetrics>,
794 object_store_metrics: Arc<ObjectStoreMetrics>,
795 storage_metrics: Arc<MonitoredStorageMetrics>,
796 compactor_metrics: Arc<CompactorMetrics>,
797 await_tree_config: Option<await_tree::Config>,
798 use_new_object_prefix_strategy: bool,
799 ) -> StorageResult<Self> {
800 const KB: usize = 1 << 10;
801 const MB: usize = 1 << 20;
802
803 let meta_cache = {
804 let mut builder = HybridCacheBuilder::new()
805 .with_name("foyer.meta")
806 .with_metrics_registry(FOYER_METRICS_REGISTRY.clone())
807 .memory(opts.meta_cache_capacity_mb * MB)
808 .with_shards(opts.meta_cache_shard_num)
809 .with_eviction_config(opts.meta_cache_eviction_config.clone())
810 .with_weighter(
811 |_: &HummockSstableObjectId, value: &SstableMetaCacheValue| {
812 std::mem::size_of::<HummockSstableObjectId>()
813 + value.estimated_meta_cache_memory_weight()
814 },
815 )
816 .storage();
817
818 if !opts.meta_file_cache_dir.is_empty() {
819 if let Err(e) = Feature::ElasticDiskCache.check_available() {
820 tracing::warn!(error = %e.as_report(), "ElasticDiskCache is not available.");
821 } else {
822 let device_builder = FsDeviceBuilder::new(&opts.meta_file_cache_dir)
823 .with_capacity(opts.meta_file_cache_capacity_mb * MB)
824 .with_throttle(opts.meta_file_cache_throttle.clone());
825 #[cfg(target_os = "linux")]
826 let device_builder = device_builder.with_direct(opts.meta_file_cache_direct_io);
827 let device = device_builder.build().map_err(HummockError::foyer_error)?;
828 let engine_builder = BlockEngineConfig::new(device)
829 .with_block_size(opts.meta_file_cache_file_capacity_mb * MB)
830 .with_indexer_shards(opts.meta_file_cache_indexer_shards)
831 .with_flushers(opts.meta_file_cache_flushers)
832 .with_reclaimers(opts.meta_file_cache_reclaimers)
833 .with_buffer_pool_size(opts.meta_file_cache_flush_buffer_threshold_mb * MB)
834 .with_submit_queue_size_threshold(
835 opts.meta_file_cache_submit_queue_size_threshold_mb * MB,
836 )
837 .with_clean_block_threshold(
838 opts.meta_file_cache_reclaimers + opts.meta_file_cache_reclaimers / 2,
839 )
840 .with_recover_concurrency(opts.meta_file_cache_recover_concurrency)
841 .with_blob_index_size(opts.meta_file_cache_blob_index_size_kb * KB)
842 .with_eviction_pickers(vec![Box::new(FifoPicker::new(
843 opts.meta_file_cache_fifo_probation_ratio,
844 ))]);
845 builder = builder
846 .with_engine_config(engine_builder)
847 .with_recover_mode(opts.meta_file_cache_recover_mode)
848 .with_compression(opts.meta_file_cache_compression);
849 builder = builder.with_spawner(build_file_cache_spawner(
850 "foyer.meta",
851 &opts.meta_file_cache_runtime_config,
852 )?);
853 }
854 }
855
856 builder.build().await.map_err(HummockError::foyer_error)?
857 };
858
859 let block_cache = {
860 let mut builder = HybridCacheBuilder::new()
861 .with_name("foyer.data")
862 .with_metrics_registry(FOYER_METRICS_REGISTRY.clone())
863 .with_event_listener(Arc::new(BlockCacheEventListener::new(
864 state_store_metrics.clone(),
865 )))
866 .memory(opts.block_cache_capacity_mb * MB)
867 .with_shards(opts.block_cache_shard_num)
868 .with_eviction_config(opts.block_cache_eviction_config.clone())
869 .with_weighter(|_: &SstableBlockIndex, value: &Box<Block>| {
870 std::mem::size_of::<SstableBlockIndex>() + value.estimated_memory_weight()
871 })
872 .storage();
873
874 if !opts.data_file_cache_dir.is_empty() {
875 if let Err(e) = Feature::ElasticDiskCache.check_available() {
876 tracing::warn!(error = %e.as_report(), "ElasticDiskCache is not available.");
877 } else {
878 let device_builder = FsDeviceBuilder::new(&opts.data_file_cache_dir)
879 .with_capacity(opts.data_file_cache_capacity_mb * MB)
880 .with_throttle(opts.data_file_cache_throttle.clone());
881 #[cfg(target_os = "linux")]
882 let device_builder = device_builder.with_direct(opts.data_file_cache_direct_io);
883 let device = device_builder.build().map_err(HummockError::foyer_error)?;
884 let engine_builder = BlockEngineConfig::new(device)
885 .with_block_size(opts.data_file_cache_file_capacity_mb * MB)
886 .with_indexer_shards(opts.data_file_cache_indexer_shards)
887 .with_flushers(opts.data_file_cache_flushers)
888 .with_reclaimers(opts.data_file_cache_reclaimers)
889 .with_buffer_pool_size(opts.data_file_cache_flush_buffer_threshold_mb * MB)
890 .with_submit_queue_size_threshold(
891 opts.data_file_cache_submit_queue_size_threshold_mb * MB,
892 )
893 .with_clean_block_threshold(
894 opts.data_file_cache_reclaimers + opts.data_file_cache_reclaimers / 2,
895 )
896 .with_recover_concurrency(opts.data_file_cache_recover_concurrency)
897 .with_blob_index_size(opts.data_file_cache_blob_index_size_kb * KB)
898 .with_eviction_pickers(vec![Box::new(FifoPicker::new(
899 opts.data_file_cache_fifo_probation_ratio,
900 ))]);
901 builder = builder
902 .with_engine_config(engine_builder)
903 .with_recover_mode(opts.data_file_cache_recover_mode)
904 .with_compression(opts.data_file_cache_compression);
905 builder = builder.with_spawner(build_file_cache_spawner(
906 "foyer.data",
907 &opts.data_file_cache_runtime_config,
908 )?);
909 }
910 }
911
912 builder.build().await.map_err(HummockError::foyer_error)?
913 };
914
915 let vector_meta_cache = CacheBuilder::new(opts.vector_meta_cache_capacity_mb * MB)
916 .with_shards(opts.vector_meta_cache_shard_num)
917 .with_eviction_config(opts.vector_meta_cache_eviction_config.clone())
918 .build();
919
920 let vector_block_cache = CacheBuilder::new(opts.vector_block_cache_capacity_mb * MB)
921 .with_shards(opts.vector_block_cache_shard_num)
922 .with_eviction_config(opts.vector_block_cache_eviction_config.clone())
923 .build();
924
925 let recent_filter = if opts.data_file_cache_dir.is_empty() {
926 Arc::new(NoneRecentFilter::default().into())
927 } else if opts.cache_refill_recent_filter_shards == 1 {
928 Arc::new(
929 SimpleRecentFilter::new(
930 opts.cache_refill_recent_filter_layers,
931 Duration::from_millis(
932 opts.cache_refill_recent_filter_rotate_interval_ms as u64,
933 ),
934 )
935 .into(),
936 )
937 } else if opts.cache_refill_skip_recent_filter {
938 Arc::new(AllRecentFilter::default().into())
939 } else {
940 Arc::new(
941 ShardedRecentFilter::new(
942 opts.cache_refill_recent_filter_layers,
943 Duration::from_millis(
944 opts.cache_refill_recent_filter_rotate_interval_ms as u64,
945 ),
946 opts.cache_refill_recent_filter_shards,
947 )
948 .into(),
949 )
950 };
951
952 let store = match s {
953 hummock if hummock.starts_with("hummock+") => {
954 let object_store = build_remote_object_store(
955 hummock.strip_prefix("hummock+").unwrap(),
956 object_store_metrics.clone(),
957 "Hummock",
958 Arc::new(opts.object_store_config.clone()),
959 )
960 .await;
961
962 let sstable_store = Arc::new(SstableStore::new(SstableStoreConfig {
963 store: Arc::new(object_store),
964 path: opts.data_directory.clone(),
965 prefetch_buffer_capacity: opts.prefetch_buffer_capacity_mb * (1 << 20),
966 max_prefetch_block_number: opts.max_prefetch_block_number,
967 recent_filter,
968 state_store_metrics: state_store_metrics.clone(),
969 use_new_object_prefix_strategy,
970 skip_bloom_filter_in_serde: opts.sst_skip_bloom_filter_in_serde,
971
972 meta_cache,
973 block_cache,
974 vector_meta_cache,
975 vector_block_cache,
976 }));
977 let notification_client =
978 RpcNotificationClient::new(hummock_meta_client.get_inner().clone());
979 let compaction_catalog_manager_ref =
980 Arc::new(CompactionCatalogManager::new(Box::new(
981 RemoteTableAccessor::new(hummock_meta_client.get_inner().clone()),
982 )));
983
984 let inner = HummockStorage::new(
985 role,
986 opts.clone(),
987 sstable_store,
988 hummock_meta_client.clone(),
989 notification_client,
990 compaction_catalog_manager_ref,
991 state_store_metrics.clone(),
992 compactor_metrics.clone(),
993 await_tree_config,
994 )
995 .await?;
996
997 StateStoreImpl::hummock(inner, storage_metrics)
998 }
999
1000 "in_memory" | "in-memory" => {
1001 tracing::warn!(
1002 "In-memory state store should never be used in end-to-end benchmarks or production environment. Scaling and recovery are not supported."
1003 );
1004 StateStoreImpl::shared_in_memory_store(storage_metrics.clone())
1005 }
1006
1007 sled if sled.starts_with("sled://") => {
1008 tracing::warn!(
1009 "sled state store should never be used in end-to-end benchmarks or production environment. Scaling and recovery are not supported."
1010 );
1011 let path = sled.strip_prefix("sled://").unwrap();
1012 StateStoreImpl::sled(SledStateStore::new(path), storage_metrics.clone())
1013 }
1014
1015 other => unimplemented!("{} state store is not supported", other),
1016 };
1017
1018 Ok(store)
1019 }
1020}
1021
1022pub trait AsHummock: Send + Sync {
1023 fn as_hummock(&self) -> Option<&HummockStorage>;
1024
1025 fn sync(
1026 &self,
1027 sync_table_epochs: Vec<(HummockEpoch, HashSet<TableId>)>,
1028 ) -> BoxFuture<'_, StorageResult<SyncResult>> {
1029 async move {
1030 if let Some(hummock) = self.as_hummock() {
1031 hummock.sync(sync_table_epochs).await
1032 } else {
1033 Ok(SyncResult::default())
1034 }
1035 }
1036 .boxed()
1037 }
1038}
1039
1040impl AsHummock for HummockStorage {
1041 fn as_hummock(&self) -> Option<&HummockStorage> {
1042 Some(self)
1043 }
1044}
1045
1046impl AsHummock for MemoryStateStore {
1047 fn as_hummock(&self) -> Option<&HummockStorage> {
1048 None
1049 }
1050}
1051
1052impl AsHummock for SledStateStore {
1053 fn as_hummock(&self) -> Option<&HummockStorage> {
1054 None
1055 }
1056}
1057
1058#[cfg(debug_assertions)]
1059mod dyn_state_store {
1060 use std::future::Future;
1061 use std::ops::DerefMut;
1062 use std::sync::Arc;
1063
1064 use bytes::Bytes;
1065 use risingwave_common::array::VectorRef;
1066 use risingwave_common::bitmap::Bitmap;
1067 use risingwave_common::hash::VirtualNode;
1068 use risingwave_hummock_sdk::HummockReadEpoch;
1069 use risingwave_hummock_sdk::key::{TableKey, TableKeyRange};
1070
1071 use crate::error::StorageResult;
1072 use crate::hummock::HummockStorage;
1073 use crate::store::*;
1074 use crate::store_impl::AsHummock;
1075 use crate::vector::VectorDistance;
1076
1077 #[async_trait::async_trait]
1078 pub trait DynStateStoreIter<T: IterItem>: Send {
1079 async fn try_next(&mut self) -> StorageResult<Option<T::ItemRef<'_>>>;
1080 }
1081
1082 #[async_trait::async_trait]
1083 impl<T: IterItem, I: StateStoreIter<T>> DynStateStoreIter<T> for I {
1084 async fn try_next(&mut self) -> StorageResult<Option<T::ItemRef<'_>>> {
1085 self.try_next().await
1086 }
1087 }
1088
1089 pub type BoxStateStoreIter<'a, T> = Box<dyn DynStateStoreIter<T> + 'a>;
1090 impl<T: IterItem> StateStoreIter<T> for BoxStateStoreIter<'_, T> {
1091 fn try_next(
1092 &mut self,
1093 ) -> impl Future<Output = StorageResult<Option<T::ItemRef<'_>>>> + Send + '_ {
1094 self.deref_mut().try_next()
1095 }
1096 }
1097
1098 pub type BoxStateStoreReadIter = BoxStateStoreIter<'static, StateStoreKeyedRow>;
1101 pub type BoxStateStoreReadChangeLogIter = BoxStateStoreIter<'static, StateStoreReadLogItem>;
1102
1103 #[async_trait::async_trait]
1104 pub trait DynStateStoreGet: StaticSendSync {
1105 async fn get_keyed_row(
1106 &self,
1107 key: TableKey<Bytes>,
1108 read_options: ReadOptions,
1109 ) -> StorageResult<Option<StateStoreKeyedRow>>;
1110 }
1111
1112 #[async_trait::async_trait]
1113 pub trait DynStateStoreRead: DynStateStoreGet + StaticSendSync {
1114 fn read_epoch(&self) -> Option<u64>;
1115
1116 async fn iter(
1117 &self,
1118 key_range: TableKeyRange,
1119
1120 read_options: ReadOptions,
1121 ) -> StorageResult<BoxStateStoreReadIter>;
1122
1123 async fn rev_iter(
1124 &self,
1125 key_range: TableKeyRange,
1126
1127 read_options: ReadOptions,
1128 ) -> StorageResult<BoxStateStoreReadIter>;
1129 }
1130
1131 #[async_trait::async_trait]
1132 pub trait DynStateStoreReadLog: StaticSendSync {
1133 async fn next_epoch(&self, epoch: u64, options: NextEpochOptions) -> StorageResult<u64>;
1134 async fn iter_log(
1135 &self,
1136 epoch_range: (u64, u64),
1137 key_range: TableKeyRange,
1138 options: ReadLogOptions,
1139 ) -> StorageResult<BoxStateStoreReadChangeLogIter>;
1140 }
1141
1142 pub type StateStoreReadDynRef = StateStorePointer<Arc<dyn DynStateStoreRead>>;
1143
1144 #[async_trait::async_trait]
1145 impl<S: StateStoreGet> DynStateStoreGet for S {
1146 async fn get_keyed_row(
1147 &self,
1148 key: TableKey<Bytes>,
1149 read_options: ReadOptions,
1150 ) -> StorageResult<Option<StateStoreKeyedRow>> {
1151 self.on_key_value(key, read_options, move |key, value| {
1152 Ok((key.copy_into(), Bytes::copy_from_slice(value)))
1153 })
1154 .await
1155 }
1156 }
1157
1158 #[async_trait::async_trait]
1159 impl<S: StateStoreRead> DynStateStoreRead for S {
1160 fn read_epoch(&self) -> Option<u64> {
1161 StateStoreRead::read_epoch(self)
1162 }
1163
1164 async fn iter(
1165 &self,
1166 key_range: TableKeyRange,
1167
1168 read_options: ReadOptions,
1169 ) -> StorageResult<BoxStateStoreReadIter> {
1170 Ok(Box::new(self.iter(key_range, read_options).await?))
1171 }
1172
1173 async fn rev_iter(
1174 &self,
1175 key_range: TableKeyRange,
1176
1177 read_options: ReadOptions,
1178 ) -> StorageResult<BoxStateStoreReadIter> {
1179 Ok(Box::new(self.rev_iter(key_range, read_options).await?))
1180 }
1181 }
1182
1183 #[async_trait::async_trait]
1184 impl<S: StateStoreReadLog> DynStateStoreReadLog for S {
1185 async fn next_epoch(&self, epoch: u64, options: NextEpochOptions) -> StorageResult<u64> {
1186 self.next_epoch(epoch, options).await
1187 }
1188
1189 async fn iter_log(
1190 &self,
1191 epoch_range: (u64, u64),
1192 key_range: TableKeyRange,
1193 options: ReadLogOptions,
1194 ) -> StorageResult<BoxStateStoreReadChangeLogIter> {
1195 Ok(Box::new(
1196 self.iter_log(epoch_range, key_range, options).await?,
1197 ))
1198 }
1199 }
1200
1201 pub type BoxLocalStateStoreIterStream<'a> = BoxStateStoreIter<'a, StateStoreKeyedRow>;
1203 #[async_trait::async_trait]
1204 pub trait DynLocalStateStore:
1205 DynStateStoreGet + DynStateStoreWriteEpochControl + StaticSendSync
1206 {
1207 async fn iter(
1208 &self,
1209 key_range: TableKeyRange,
1210 read_options: ReadOptions,
1211 ) -> StorageResult<BoxLocalStateStoreIterStream<'_>>;
1212
1213 async fn rev_iter(
1214 &self,
1215 key_range: TableKeyRange,
1216 read_options: ReadOptions,
1217 ) -> StorageResult<BoxLocalStateStoreIterStream<'_>>;
1218
1219 fn new_flushed_snapshot_reader(&self) -> StateStoreReadDynRef;
1220
1221 fn insert(
1222 &mut self,
1223 key: TableKey<Bytes>,
1224 new_val: Bytes,
1225 old_val: Option<Bytes>,
1226 ) -> StorageResult<()>;
1227
1228 fn delete(&mut self, key: TableKey<Bytes>, old_val: Bytes) -> StorageResult<()>;
1229
1230 async fn update_vnode_bitmap(&mut self, vnodes: Arc<Bitmap>) -> StorageResult<Arc<Bitmap>>;
1231
1232 fn get_table_watermark(&self, vnode: VirtualNode) -> Option<Bytes>;
1233 }
1234
1235 #[async_trait::async_trait]
1236 pub trait DynStateStoreWriteEpochControl: StaticSendSync {
1237 async fn flush(&mut self) -> StorageResult<usize>;
1238
1239 async fn try_flush(&mut self) -> StorageResult<()>;
1240
1241 async fn init(&mut self, epoch: InitOptions) -> StorageResult<()>;
1242
1243 fn seal_current_epoch(&mut self, next_epoch: u64, opts: SealCurrentEpochOptions);
1244 }
1245
1246 #[async_trait::async_trait]
1247 impl<S: LocalStateStore> DynLocalStateStore for S {
1248 async fn iter(
1249 &self,
1250 key_range: TableKeyRange,
1251 read_options: ReadOptions,
1252 ) -> StorageResult<BoxLocalStateStoreIterStream<'_>> {
1253 Ok(Box::new(self.iter(key_range, read_options).await?))
1254 }
1255
1256 async fn rev_iter(
1257 &self,
1258 key_range: TableKeyRange,
1259 read_options: ReadOptions,
1260 ) -> StorageResult<BoxLocalStateStoreIterStream<'_>> {
1261 Ok(Box::new(self.rev_iter(key_range, read_options).await?))
1262 }
1263
1264 fn new_flushed_snapshot_reader(&self) -> StateStoreReadDynRef {
1265 StateStorePointer(Arc::new(self.new_flushed_snapshot_reader()) as _)
1266 }
1267
1268 fn insert(
1269 &mut self,
1270 key: TableKey<Bytes>,
1271 new_val: Bytes,
1272 old_val: Option<Bytes>,
1273 ) -> StorageResult<()> {
1274 self.insert(key, new_val, old_val)
1275 }
1276
1277 fn delete(&mut self, key: TableKey<Bytes>, old_val: Bytes) -> StorageResult<()> {
1278 self.delete(key, old_val)
1279 }
1280
1281 async fn update_vnode_bitmap(&mut self, vnodes: Arc<Bitmap>) -> StorageResult<Arc<Bitmap>> {
1282 self.update_vnode_bitmap(vnodes).await
1283 }
1284
1285 fn get_table_watermark(&self, vnode: VirtualNode) -> Option<Bytes> {
1286 self.get_table_watermark(vnode)
1287 }
1288 }
1289
1290 #[async_trait::async_trait]
1291 impl<S: StateStoreWriteEpochControl> DynStateStoreWriteEpochControl for S {
1292 async fn flush(&mut self) -> StorageResult<usize> {
1293 self.flush().await
1294 }
1295
1296 async fn try_flush(&mut self) -> StorageResult<()> {
1297 self.try_flush().await
1298 }
1299
1300 async fn init(&mut self, options: InitOptions) -> StorageResult<()> {
1301 self.init(options).await
1302 }
1303
1304 fn seal_current_epoch(&mut self, next_epoch: u64, opts: SealCurrentEpochOptions) {
1305 self.seal_current_epoch(next_epoch, opts)
1306 }
1307 }
1308
1309 pub type BoxDynLocalStateStore = StateStorePointer<Box<dyn DynLocalStateStore>>;
1310
1311 impl LocalStateStore for BoxDynLocalStateStore {
1312 type FlushedSnapshotReader = StateStoreReadDynRef;
1313 type Iter<'a> = BoxLocalStateStoreIterStream<'a>;
1314 type RevIter<'a> = BoxLocalStateStoreIterStream<'a>;
1315
1316 fn iter(
1317 &self,
1318 key_range: TableKeyRange,
1319 read_options: ReadOptions,
1320 ) -> impl Future<Output = StorageResult<Self::Iter<'_>>> + Send + '_ {
1321 (*self.0).iter(key_range, read_options)
1322 }
1323
1324 fn rev_iter(
1325 &self,
1326 key_range: TableKeyRange,
1327 read_options: ReadOptions,
1328 ) -> impl Future<Output = StorageResult<Self::RevIter<'_>>> + Send + '_ {
1329 (*self.0).rev_iter(key_range, read_options)
1330 }
1331
1332 fn new_flushed_snapshot_reader(&self) -> Self::FlushedSnapshotReader {
1333 (*self.0).new_flushed_snapshot_reader()
1334 }
1335
1336 fn get_table_watermark(&self, vnode: VirtualNode) -> Option<Bytes> {
1337 (*self.0).get_table_watermark(vnode)
1338 }
1339
1340 fn insert(
1341 &mut self,
1342 key: TableKey<Bytes>,
1343 new_val: Bytes,
1344 old_val: Option<Bytes>,
1345 ) -> StorageResult<()> {
1346 (*self.0).insert(key, new_val, old_val)
1347 }
1348
1349 fn delete(&mut self, key: TableKey<Bytes>, old_val: Bytes) -> StorageResult<()> {
1350 (*self.0).delete(key, old_val)
1351 }
1352
1353 async fn update_vnode_bitmap(&mut self, vnodes: Arc<Bitmap>) -> StorageResult<Arc<Bitmap>> {
1354 (*self.0).update_vnode_bitmap(vnodes).await
1355 }
1356 }
1357
1358 impl<P> StateStoreWriteEpochControl for StateStorePointer<P>
1359 where
1360 StateStorePointer<P>: AsMut<dyn DynStateStoreWriteEpochControl> + StaticSendSync,
1361 {
1362 fn flush(&mut self) -> impl Future<Output = StorageResult<usize>> + Send + '_ {
1363 self.as_mut().flush()
1364 }
1365
1366 fn try_flush(&mut self) -> impl Future<Output = StorageResult<()>> + Send + '_ {
1367 self.as_mut().try_flush()
1368 }
1369
1370 fn init(
1371 &mut self,
1372 options: InitOptions,
1373 ) -> impl Future<Output = StorageResult<()>> + Send + '_ {
1374 self.as_mut().init(options)
1375 }
1376
1377 fn seal_current_epoch(&mut self, next_epoch: u64, opts: SealCurrentEpochOptions) {
1378 self.as_mut().seal_current_epoch(next_epoch, opts)
1379 }
1380 }
1381
1382 #[async_trait::async_trait]
1383 pub trait DynStateStoreWriteVector: DynStateStoreWriteEpochControl + StaticSendSync {
1384 fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()>;
1385 }
1386
1387 #[async_trait::async_trait]
1388 impl<S: StateStoreWriteVector> DynStateStoreWriteVector for S {
1389 fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()> {
1390 self.insert(vec, info)
1391 }
1392 }
1393
1394 pub type BoxDynStateStoreWriteVector = StateStorePointer<Box<dyn DynStateStoreWriteVector>>;
1395
1396 impl StateStoreWriteVector for BoxDynStateStoreWriteVector {
1397 fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()> {
1398 self.0.insert(vec, info)
1399 }
1400 }
1401
1402 #[async_trait::async_trait]
1405 pub trait DynStateStoreReadVector: StaticSendSync {
1406 async fn nearest(
1407 &self,
1408 vec: VectorRef<'_>,
1409 options: VectorNearestOptions,
1410 ) -> StorageResult<Vec<(Vector, VectorDistance, Bytes)>>;
1411 }
1412
1413 #[async_trait::async_trait]
1414 impl<S: StateStoreReadVector> DynStateStoreReadVector for S {
1415 async fn nearest(
1416 &self,
1417 vec: VectorRef<'_>,
1418 options: VectorNearestOptions,
1419 ) -> StorageResult<Vec<(Vector, VectorDistance, Bytes)>> {
1420 use risingwave_common::types::ScalarRef;
1421 self.nearest(vec, options, |vec, distance, info| {
1422 (
1423 vec.to_owned_scalar(),
1424 distance,
1425 Bytes::copy_from_slice(info),
1426 )
1427 })
1428 .await
1429 }
1430 }
1431
1432 impl<P> StateStoreReadVector for StateStorePointer<P>
1433 where
1434 StateStorePointer<P>: AsRef<dyn DynStateStoreReadVector> + StaticSendSync,
1435 {
1436 async fn nearest<'a, O: Send + 'a>(
1437 &'a self,
1438 vec: VectorRef<'a>,
1439 options: VectorNearestOptions,
1440 on_nearest_item_fn: impl OnNearestItemFn<'a, O>,
1441 ) -> StorageResult<Vec<O>> {
1442 let output = self.as_ref().nearest(vec, options).await?;
1443 Ok(output
1444 .into_iter()
1445 .map(|(vec, distance, info)| {
1446 on_nearest_item_fn(vec.to_ref(), distance, info.as_ref())
1447 })
1448 .collect())
1449 }
1450 }
1451
1452 pub trait DynStateStoreReadSnapshot:
1453 DynStateStoreRead + DynStateStoreReadVector + StaticSendSync
1454 {
1455 }
1456
1457 impl<S: DynStateStoreRead + DynStateStoreReadVector + StaticSendSync> DynStateStoreReadSnapshot
1458 for S
1459 {
1460 }
1461
1462 pub type StateStoreReadSnapshotDynRef = StateStorePointer<Arc<dyn DynStateStoreReadSnapshot>>;
1463 #[async_trait::async_trait]
1464 pub trait DynStateStoreExt: StaticSendSync {
1465 async fn try_wait_epoch(
1466 &self,
1467 epoch: HummockReadEpoch,
1468 options: TryWaitEpochOptions,
1469 ) -> StorageResult<()>;
1470
1471 async fn new_local(&self, option: NewLocalOptions) -> BoxDynLocalStateStore;
1472 async fn new_read_snapshot(
1473 &self,
1474 epoch: HummockReadEpoch,
1475 options: NewReadSnapshotOptions,
1476 ) -> StorageResult<StateStoreReadSnapshotDynRef>;
1477 async fn new_vector_writer(
1478 &self,
1479 options: NewVectorWriterOptions,
1480 ) -> BoxDynStateStoreWriteVector;
1481 }
1482
1483 #[async_trait::async_trait]
1484 impl<S: StateStore> DynStateStoreExt for S {
1485 async fn try_wait_epoch(
1486 &self,
1487 epoch: HummockReadEpoch,
1488 options: TryWaitEpochOptions,
1489 ) -> StorageResult<()> {
1490 self.try_wait_epoch(epoch, options).await
1491 }
1492
1493 async fn new_local(&self, option: NewLocalOptions) -> BoxDynLocalStateStore {
1494 StateStorePointer(Box::new(self.new_local(option).await))
1495 }
1496
1497 async fn new_read_snapshot(
1498 &self,
1499 epoch: HummockReadEpoch,
1500 options: NewReadSnapshotOptions,
1501 ) -> StorageResult<StateStoreReadSnapshotDynRef> {
1502 Ok(StateStorePointer(Arc::new(
1503 self.new_read_snapshot(epoch, options).await?,
1504 )))
1505 }
1506
1507 async fn new_vector_writer(
1508 &self,
1509 options: NewVectorWriterOptions,
1510 ) -> BoxDynStateStoreWriteVector {
1511 StateStorePointer(Box::new(self.new_vector_writer(options).await))
1512 }
1513 }
1514
1515 pub type StateStoreDynRef = StateStorePointer<Arc<dyn DynStateStore>>;
1516
1517 macro_rules! state_store_pointer_dyn_as_ref {
1518 ($pointer:ident < dyn $source_dyn_trait:ident > , $target_dyn_trait:ident) => {
1519 impl AsRef<dyn $target_dyn_trait>
1520 for StateStorePointer<$pointer<dyn $source_dyn_trait>>
1521 {
1522 fn as_ref(&self) -> &dyn $target_dyn_trait {
1523 (&*self.0) as _
1524 }
1525 }
1526 };
1527 }
1528
1529 state_store_pointer_dyn_as_ref!(Arc<dyn DynStateStoreReadSnapshot>, DynStateStoreRead);
1530 state_store_pointer_dyn_as_ref!(Arc<dyn DynStateStoreReadSnapshot>, DynStateStoreGet);
1531 state_store_pointer_dyn_as_ref!(Arc<dyn DynStateStoreReadSnapshot>, DynStateStoreReadVector);
1532 state_store_pointer_dyn_as_ref!(Arc<dyn DynStateStoreRead>, DynStateStoreRead);
1533 state_store_pointer_dyn_as_ref!(Arc<dyn DynStateStoreRead>, DynStateStoreGet);
1534 state_store_pointer_dyn_as_ref!(Box<dyn DynLocalStateStore>, DynStateStoreGet);
1535
1536 macro_rules! state_store_pointer_dyn_as_mut {
1537 ($pointer:ident < dyn $source_dyn_trait:ident > , $target_dyn_trait:ident) => {
1538 impl AsMut<dyn $target_dyn_trait>
1539 for StateStorePointer<$pointer<dyn $source_dyn_trait>>
1540 {
1541 fn as_mut(&mut self) -> &mut dyn $target_dyn_trait {
1542 (&mut *self.0) as _
1543 }
1544 }
1545 };
1546 }
1547
1548 state_store_pointer_dyn_as_mut!(Box<dyn DynLocalStateStore>, DynStateStoreWriteEpochControl);
1549 state_store_pointer_dyn_as_mut!(
1550 Box<dyn DynStateStoreWriteVector>,
1551 DynStateStoreWriteEpochControl
1552 );
1553
1554 #[derive(Clone)]
1555 pub struct StateStorePointer<P>(pub(crate) P);
1556
1557 impl<P> StateStoreGet for StateStorePointer<P>
1558 where
1559 StateStorePointer<P>: AsRef<dyn DynStateStoreGet> + StaticSendSync,
1560 {
1561 async fn on_key_value<'a, O: Send + 'a>(
1562 &'a self,
1563 key: TableKey<Bytes>,
1564 read_options: ReadOptions,
1565 on_key_value_fn: impl KeyValueFn<'a, O>,
1566 ) -> StorageResult<Option<O>> {
1567 let option = self.as_ref().get_keyed_row(key, read_options).await?;
1568 option
1569 .map(|(key, value)| on_key_value_fn(key.to_ref(), value.as_ref()))
1570 .transpose()
1571 }
1572 }
1573
1574 impl<P> StateStoreRead for StateStorePointer<P>
1575 where
1576 StateStorePointer<P>: AsRef<dyn DynStateStoreRead> + StateStoreGet + StaticSendSync,
1577 {
1578 type Iter = BoxStateStoreReadIter;
1579 type RevIter = BoxStateStoreReadIter;
1580
1581 fn read_epoch(&self) -> Option<u64> {
1582 self.as_ref().read_epoch()
1583 }
1584
1585 fn iter(
1586 &self,
1587 key_range: TableKeyRange,
1588
1589 read_options: ReadOptions,
1590 ) -> impl Future<Output = StorageResult<Self::Iter>> + '_ {
1591 self.as_ref().iter(key_range, read_options)
1592 }
1593
1594 fn rev_iter(
1595 &self,
1596 key_range: TableKeyRange,
1597
1598 read_options: ReadOptions,
1599 ) -> impl Future<Output = StorageResult<Self::RevIter>> + '_ {
1600 self.as_ref().rev_iter(key_range, read_options)
1601 }
1602 }
1603
1604 impl StateStoreReadLog for StateStoreDynRef {
1605 type ChangeLogIter = BoxStateStoreReadChangeLogIter;
1606
1607 async fn next_epoch(&self, epoch: u64, options: NextEpochOptions) -> StorageResult<u64> {
1608 (*self.0).next_epoch(epoch, options).await
1609 }
1610
1611 fn iter_log(
1612 &self,
1613 epoch_range: (u64, u64),
1614 key_range: TableKeyRange,
1615 options: ReadLogOptions,
1616 ) -> impl Future<Output = StorageResult<Self::ChangeLogIter>> + Send + '_ {
1617 (*self.0).iter_log(epoch_range, key_range, options)
1618 }
1619 }
1620
1621 pub trait DynStateStore: DynStateStoreReadLog + DynStateStoreExt + AsHummock {}
1622
1623 impl AsHummock for StateStoreDynRef {
1624 fn as_hummock(&self) -> Option<&HummockStorage> {
1625 (*self.0).as_hummock()
1626 }
1627 }
1628
1629 impl<S: DynStateStoreReadLog + DynStateStoreExt + AsHummock> DynStateStore for S {}
1630
1631 impl StateStore for StateStoreDynRef {
1632 type Local = BoxDynLocalStateStore;
1633 type ReadSnapshot = StateStoreReadSnapshotDynRef;
1634 type VectorWriter = BoxDynStateStoreWriteVector;
1635
1636 fn try_wait_epoch(
1637 &self,
1638 epoch: HummockReadEpoch,
1639 options: TryWaitEpochOptions,
1640 ) -> impl Future<Output = StorageResult<()>> + Send + '_ {
1641 (*self.0).try_wait_epoch(epoch, options)
1642 }
1643
1644 fn new_local(
1645 &self,
1646 option: NewLocalOptions,
1647 ) -> impl Future<Output = Self::Local> + Send + '_ {
1648 (*self.0).new_local(option)
1649 }
1650
1651 async fn new_read_snapshot(
1652 &self,
1653 epoch: HummockReadEpoch,
1654 options: NewReadSnapshotOptions,
1655 ) -> StorageResult<Self::ReadSnapshot> {
1656 (*self.0).new_read_snapshot(epoch, options).await
1657 }
1658
1659 fn new_vector_writer(
1660 &self,
1661 options: NewVectorWriterOptions,
1662 ) -> impl Future<Output = Self::VectorWriter> + Send + '_ {
1663 (*self.0).new_vector_writer(options)
1664 }
1665 }
1666}