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