Skip to main content

risingwave_storage/
store_impl.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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/// The type erased [`StateStore`].
224#[derive(Clone, EnumAsInner)]
225#[expect(clippy::enum_variant_names)]
226pub enum StateStoreImpl {
227    /// The Hummock state store, which operates on an S3-like service. URLs beginning with
228    /// `hummock` will be automatically recognized as Hummock state store.
229    ///
230    /// Example URLs:
231    ///
232    /// * `hummock+s3://bucket`
233    /// * `hummock+minio://KEY:SECRET@minio-ip:port`
234    /// * `hummock+memory` (should only be used in 1 compute node mode)
235    HummockStateStore(Monitored<HummockStorageType>),
236    /// In-memory B-Tree state store. Should only be used in unit and integration tests. If you
237    /// want speed up e2e test, you should use Hummock in-memory mode instead. Also, this state
238    /// store misses some critical implementation to ensure the correctness of persisting streaming
239    /// state. (e.g., no `read_epoch` support, no async checkpoint)
240    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        // The specific type of MemoryStateStoreType in deducted here.
291        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        // The specific type of HummockStateStoreType in deducted here.
299        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                // WARNING: don't change this. Enabling memory backend will cause monomorphization
348                // explosion and thus slow compile time in release mode.
349                #[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                // WARNING: don't change this. Enabling memory backend will cause monomorphization
362                // explosion and thus slow compile time in release mode.
363                #[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        // TODO: may avoid manual async fn when the bug of rust compiler is fixed. Currently it will
490        // fail to compile.
491        #[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    // For StateStoreRead
1099
1100    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    // For LocalStateStore
1202    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    // For global StateStore
1403
1404    #[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}