risingwave_storage/
store.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::default::Default;
16use std::fmt::{Debug, Formatter};
17use std::future::Future;
18use std::marker::PhantomData;
19use std::sync::{Arc, LazyLock};
20
21use bytes::Bytes;
22use futures::{Stream, TryStreamExt};
23use futures_async_stream::try_stream;
24use prost::Message;
25use risingwave_common::array::{Op, VectorRef};
26use risingwave_common::bitmap::Bitmap;
27use risingwave_common::catalog::{TableId, TableOption};
28use risingwave_common::hash::VirtualNode;
29use risingwave_common::id::FragmentId;
30use risingwave_common::util::epoch::{Epoch, EpochPair, MAX_EPOCH};
31use risingwave_common::vector::distance::DistanceMeasurement;
32use risingwave_hummock_sdk::HummockReadEpoch;
33use risingwave_hummock_sdk::key::{FullKey, TableKey, TableKeyRange};
34use risingwave_hummock_sdk::table_watermark::{
35    VnodeWatermark, WatermarkDirection, WatermarkSerdeType,
36};
37use risingwave_hummock_trace::{
38    TracedInitOptions, TracedNewLocalOptions, TracedOpConsistencyLevel, TracedPrefetchOptions,
39    TracedReadOptions, TracedSealCurrentEpochOptions, TracedTryWaitEpochOptions,
40};
41use risingwave_pb::hummock::{PbVnodeWatermark, PbWatermarkSerdeType};
42
43use crate::error::{StorageError, StorageResult};
44use crate::hummock::CachePolicy;
45use crate::monitor::{MonitoredStateStore, MonitoredStorageMetrics};
46pub(crate) use crate::vector::{OnNearestItem, Vector};
47
48mod auto_rebuild;
49pub use auto_rebuild::{AutoRebuildStateStoreReadIter, timeout_auto_rebuild};
50
51pub trait StaticSendSync = Send + Sync + 'static;
52
53pub trait IterItem: Send + 'static {
54    type ItemRef<'a>: Send + Copy + 'a;
55}
56
57impl IterItem for StateStoreKeyedRow {
58    type ItemRef<'a> = StateStoreKeyedRowRef<'a>;
59}
60
61impl IterItem for StateStoreReadLogItem {
62    type ItemRef<'a> = StateStoreReadLogItemRef<'a>;
63}
64
65pub trait StateStoreIter<T: IterItem = StateStoreKeyedRow>: Send {
66    fn try_next(&mut self) -> impl StorageFuture<'_, Option<T::ItemRef<'_>>>;
67}
68
69pub fn to_owned_item((key, value): StateStoreKeyedRowRef<'_>) -> StorageResult<StateStoreKeyedRow> {
70    Ok((key.copy_into(), Bytes::copy_from_slice(value)))
71}
72
73pub trait StateStoreIterExt<T: IterItem = StateStoreKeyedRow>: StateStoreIter<T> + Sized {
74    type ItemStream<O: Send, F: Send + for<'a> Fn(T::ItemRef<'a>) -> StorageResult<O>>: Stream<Item = StorageResult<O>>
75        + Send;
76
77    fn into_stream<O: Send, F: for<'a> Fn(T::ItemRef<'a>) -> StorageResult<O> + Send>(
78        self,
79        f: F,
80    ) -> Self::ItemStream<O, F>;
81
82    fn fused(self) -> FusedStateStoreIter<Self, T> {
83        FusedStateStoreIter::new(self)
84    }
85}
86
87#[try_stream(ok = O, error = StorageError)]
88async fn into_stream_inner<
89    T: IterItem,
90    I: StateStoreIter<T>,
91    O: Send,
92    F: for<'a> Fn(T::ItemRef<'a>) -> StorageResult<O> + Send,
93>(
94    iter: I,
95    f: F,
96) {
97    let mut iter = iter.fused();
98    while let Some(item) = iter.try_next().await? {
99        yield f(item)?;
100    }
101}
102
103pub struct FromStreamStateStoreIter<S> {
104    inner: S,
105    item_buffer: Option<StateStoreKeyedRow>,
106}
107
108impl<S> FromStreamStateStoreIter<S> {
109    pub fn new(inner: S) -> Self {
110        Self {
111            inner,
112            item_buffer: None,
113        }
114    }
115}
116
117impl<S: Stream<Item = StorageResult<StateStoreKeyedRow>> + Unpin + Send> StateStoreIter
118    for FromStreamStateStoreIter<S>
119{
120    async fn try_next(&mut self) -> StorageResult<Option<StateStoreKeyedRowRef<'_>>> {
121        self.item_buffer = self.inner.try_next().await?;
122        Ok(self
123            .item_buffer
124            .as_ref()
125            .map(|(key, value)| (key.to_ref(), value.as_ref())))
126    }
127}
128
129pub struct FusedStateStoreIter<I, T> {
130    inner: I,
131    finished: bool,
132    _phantom: PhantomData<T>,
133}
134
135impl<I, T> FusedStateStoreIter<I, T> {
136    fn new(inner: I) -> Self {
137        Self {
138            inner,
139            finished: false,
140            _phantom: PhantomData,
141        }
142    }
143}
144
145impl<T: IterItem, I: StateStoreIter<T>> FusedStateStoreIter<I, T> {
146    async fn try_next(&mut self) -> StorageResult<Option<T::ItemRef<'_>>> {
147        assert!(!self.finished, "call try_next after finish");
148        let result = self.inner.try_next().await;
149        match &result {
150            Ok(Some(_)) => {}
151            Ok(None) | Err(_) => {
152                self.finished = true;
153            }
154        }
155        result
156    }
157}
158
159impl<T: IterItem, I: StateStoreIter<T>> StateStoreIterExt<T> for I {
160    type ItemStream<O: Send, F: Send + for<'a> Fn(T::ItemRef<'a>) -> StorageResult<O>> =
161        impl Stream<Item = StorageResult<O>> + Send;
162
163    fn into_stream<O: Send, F: for<'a> Fn(T::ItemRef<'a>) -> StorageResult<O> + Send>(
164        self,
165        f: F,
166    ) -> Self::ItemStream<O, F> {
167        into_stream_inner(self, f)
168    }
169}
170
171pub type StateStoreKeyedRowRef<'a> = (FullKey<&'a [u8]>, &'a [u8]);
172pub type StateStoreKeyedRow = (FullKey<Bytes>, Bytes);
173pub trait StateStoreReadIter = StateStoreIter + 'static;
174
175#[derive(Clone, Copy, Eq, PartialEq, Debug)]
176pub enum ChangeLogValue<T> {
177    Insert(T),
178    Update { new_value: T, old_value: T },
179    Delete(T),
180}
181
182impl<T> ChangeLogValue<T> {
183    pub fn try_map<O>(self, f: impl Fn(T) -> StorageResult<O>) -> StorageResult<ChangeLogValue<O>> {
184        Ok(match self {
185            ChangeLogValue::Insert(value) => ChangeLogValue::Insert(f(value)?),
186            ChangeLogValue::Update {
187                new_value,
188                old_value,
189            } => ChangeLogValue::Update {
190                new_value: f(new_value)?,
191                old_value: f(old_value)?,
192            },
193            ChangeLogValue::Delete(value) => ChangeLogValue::Delete(f(value)?),
194        })
195    }
196
197    pub fn into_op_value_iter(self) -> impl Iterator<Item = (Op, T)> {
198        std::iter::from_coroutine(
199            #[coroutine]
200            move || match self {
201                Self::Insert(row) => {
202                    yield (Op::Insert, row);
203                }
204                Self::Delete(row) => {
205                    yield (Op::Delete, row);
206                }
207                Self::Update {
208                    old_value,
209                    new_value,
210                } => {
211                    yield (Op::UpdateDelete, old_value);
212                    yield (Op::UpdateInsert, new_value);
213                }
214            },
215        )
216    }
217}
218
219impl<T: AsRef<[u8]>> ChangeLogValue<T> {
220    pub fn to_ref(&self) -> ChangeLogValue<&[u8]> {
221        match self {
222            ChangeLogValue::Insert(val) => ChangeLogValue::Insert(val.as_ref()),
223            ChangeLogValue::Update {
224                new_value,
225                old_value,
226            } => ChangeLogValue::Update {
227                new_value: new_value.as_ref(),
228                old_value: old_value.as_ref(),
229            },
230            ChangeLogValue::Delete(val) => ChangeLogValue::Delete(val.as_ref()),
231        }
232    }
233}
234
235pub type StateStoreReadLogItem = (TableKey<Bytes>, ChangeLogValue<Bytes>);
236pub type StateStoreReadLogItemRef<'a> = (TableKey<&'a [u8]>, ChangeLogValue<&'a [u8]>);
237
238#[derive(Clone)]
239pub struct NextEpochOptions {
240    pub table_id: TableId,
241}
242
243#[derive(Clone)]
244pub struct ReadLogOptions {
245    pub table_id: TableId,
246}
247
248pub trait StateStoreReadChangeLogIter = StateStoreIter<StateStoreReadLogItem> + Send + 'static;
249pub trait StorageFuture<'a, T> = Future<Output = StorageResult<T>> + Send + 'a;
250
251pub trait StateStoreReadLog: StaticSendSync {
252    type ChangeLogIter: StateStoreReadChangeLogIter;
253
254    fn next_epoch(&self, epoch: u64, options: NextEpochOptions) -> impl StorageFuture<'_, u64>;
255
256    fn iter_log(
257        &self,
258        epoch_range: (u64, u64),
259        key_range: TableKeyRange,
260        options: ReadLogOptions,
261    ) -> impl StorageFuture<'_, Self::ChangeLogIter>;
262}
263
264pub trait KeyValueFn<'a, O> =
265    for<'kv> FnOnce(FullKey<&'kv [u8]>, &'kv [u8]) -> StorageResult<O> + Send + 'a;
266
267pub trait StateStoreGet: StaticSendSync {
268    fn on_key_value<'a, O: Send + 'a>(
269        &'a self,
270        key: TableKey<Bytes>,
271        read_options: ReadOptions,
272        on_key_value_fn: impl KeyValueFn<'a, O>,
273    ) -> impl StorageFuture<'a, Option<O>>;
274}
275
276pub trait StateStoreRead: StateStoreGet + StaticSendSync {
277    type Iter: StateStoreReadIter;
278    type RevIter: StateStoreReadIter;
279
280    /// The resolved epoch of a read snapshot, if this reader has one. Tracing uses this to
281    /// record a numeric epoch for requests resolved by storage, such as latest committed.
282    fn read_epoch(&self) -> Option<u64> {
283        None
284    }
285
286    /// Opens and returns an iterator for given `prefix_hint` and `full_key_range`
287    /// Internally, `prefix_hint` will be used for checking the SST filter and
288    /// `full_key_range` used for iter. (if the `prefix_hint` not None, it should be be included
289    /// in `key_range`) The returned iterator will iterate data based on a snapshot
290    /// corresponding to the given `epoch`.
291    fn iter(
292        &self,
293        key_range: TableKeyRange,
294        read_options: ReadOptions,
295    ) -> impl StorageFuture<'_, Self::Iter>;
296
297    fn rev_iter(
298        &self,
299        key_range: TableKeyRange,
300        read_options: ReadOptions,
301    ) -> impl StorageFuture<'_, Self::RevIter>;
302}
303
304#[derive(Clone)]
305pub struct TryWaitEpochOptions {
306    pub table_id: TableId,
307}
308
309impl TryWaitEpochOptions {
310    #[cfg(any(test, feature = "test"))]
311    pub fn for_test(table_id: TableId) -> Self {
312        Self { table_id }
313    }
314}
315
316impl From<TracedTryWaitEpochOptions> for TryWaitEpochOptions {
317    fn from(value: TracedTryWaitEpochOptions) -> Self {
318        Self {
319            table_id: value.table_id.into(),
320        }
321    }
322}
323
324impl From<TryWaitEpochOptions> for TracedTryWaitEpochOptions {
325    fn from(value: TryWaitEpochOptions) -> Self {
326        Self {
327            table_id: value.table_id.into(),
328        }
329    }
330}
331
332#[derive(Clone, Copy)]
333pub struct NewReadSnapshotOptions {
334    pub table_id: TableId,
335    pub table_option: TableOption,
336}
337
338#[derive(Clone)]
339pub struct NewVectorWriterOptions {
340    pub table_id: TableId,
341}
342
343pub trait StateStore: StateStoreReadLog + StaticSendSync + Clone {
344    type Local: LocalStateStore;
345    type ReadSnapshot: StateStoreRead + StateStoreReadVector + Clone;
346    type VectorWriter: StateStoreWriteVector;
347
348    /// If epoch is `Committed`, we will wait until the epoch is committed and its data is ready to
349    /// read. If epoch is `Current`, we will only check if the data can be read with this epoch.
350    fn try_wait_epoch(
351        &self,
352        epoch: HummockReadEpoch,
353        options: TryWaitEpochOptions,
354    ) -> impl StorageFuture<'_, ()>;
355
356    /// Creates a [`MonitoredStateStore`] from this state store, with given `stats`.
357    fn monitored(self, storage_metrics: Arc<MonitoredStorageMetrics>) -> MonitoredStateStore<Self> {
358        MonitoredStateStore::new(self, storage_metrics)
359    }
360
361    fn new_local(&self, option: NewLocalOptions) -> impl Future<Output = Self::Local> + Send + '_;
362
363    fn new_read_snapshot(
364        &self,
365        epoch: HummockReadEpoch,
366        options: NewReadSnapshotOptions,
367    ) -> impl StorageFuture<'_, Self::ReadSnapshot>;
368
369    fn new_vector_writer(
370        &self,
371        options: NewVectorWriterOptions,
372    ) -> impl Future<Output = Self::VectorWriter> + Send + '_;
373}
374
375/// A state store that is dedicated for streaming operator, which only reads the uncommitted data
376/// written by itself. Each local state store is not `Clone`, and is owned by a streaming state
377/// table.
378pub trait LocalStateStore: StateStoreGet + StateStoreWriteEpochControl + StaticSendSync {
379    type FlushedSnapshotReader: StateStoreRead;
380    type Iter<'a>: StateStoreIter + 'a;
381    type RevIter<'a>: StateStoreIter + 'a;
382
383    /// Opens and returns an iterator for given `prefix_hint` and `full_key_range`
384    /// Internally, `prefix_hint` will be used for checking the SST filter and
385    /// `full_key_range` used for iter. (if the `prefix_hint` not None, it should be be included
386    /// in `key_range`) The returned iterator will iterate data based on the latest written
387    /// snapshot.
388    fn iter(
389        &self,
390        key_range: TableKeyRange,
391        read_options: ReadOptions,
392    ) -> impl StorageFuture<'_, Self::Iter<'_>>;
393
394    fn rev_iter(
395        &self,
396        key_range: TableKeyRange,
397        read_options: ReadOptions,
398    ) -> impl StorageFuture<'_, Self::RevIter<'_>>;
399
400    fn new_flushed_snapshot_reader(&self) -> Self::FlushedSnapshotReader;
401
402    /// Get last persisted watermark for a given vnode.
403    fn get_table_watermark(&self, vnode: VirtualNode) -> Option<Bytes>;
404
405    /// Inserts a key-value entry associated with a given `epoch` into the state store.
406    fn insert(
407        &mut self,
408        key: TableKey<Bytes>,
409        new_val: Bytes,
410        old_val: Option<Bytes>,
411    ) -> StorageResult<()>;
412
413    /// Deletes a key-value entry from the state store. Only the key-value entry with epoch smaller
414    /// than the given `epoch` will be deleted.
415    fn delete(&mut self, key: TableKey<Bytes>, old_val: Bytes) -> StorageResult<()>;
416
417    // Updates the vnode bitmap corresponding to the local state store
418    // Returns the previous vnode bitmap
419    fn update_vnode_bitmap(&mut self, vnodes: Arc<Bitmap>) -> impl StorageFuture<'_, Arc<Bitmap>>;
420}
421
422pub trait StateStoreWriteEpochControl: StaticSendSync {
423    fn flush(&mut self) -> impl StorageFuture<'_, usize>;
424
425    fn try_flush(&mut self) -> impl StorageFuture<'_, ()>;
426
427    /// Initializes the state store with given `epoch` pair.
428    /// Typically we will use `epoch.curr` as the initialized epoch,
429    /// Since state table will begin as empty.
430    /// In some cases like replicated state table, state table may not be empty initially,
431    /// as such we need to wait for `epoch.prev` checkpoint to complete,
432    /// hence this interface is made async.
433    fn init(&mut self, opts: InitOptions) -> impl StorageFuture<'_, ()>;
434
435    /// Updates the monotonically increasing write epoch to `new_epoch`.
436    /// All writes after this function is called will be tagged with `new_epoch`. In other words,
437    /// the previous write epoch is sealed.
438    fn seal_current_epoch(&mut self, next_epoch: u64, opts: SealCurrentEpochOptions);
439}
440
441pub trait StateStoreWriteVector: StateStoreWriteEpochControl + StaticSendSync {
442    fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()>;
443}
444
445pub struct VectorNearestOptions {
446    pub top_n: usize,
447    pub measure: DistanceMeasurement,
448    pub hnsw_ef_search: usize,
449}
450
451pub trait OnNearestItemFn<'a, O> = OnNearestItem<O> + Send + Sync + 'a;
452
453pub trait StateStoreReadVector: StaticSendSync {
454    fn nearest<'a, O: Send + 'a>(
455        &'a self,
456        vec: VectorRef<'a>,
457        options: VectorNearestOptions,
458        on_nearest_item_fn: impl OnNearestItemFn<'a, O>,
459    ) -> impl StorageFuture<'a, Vec<O>>;
460}
461
462/// Controls block-stream prefetch for range scans. Prefetching may increase the memory footprint
463/// of the CN process because the prefetched blocks cannot be evicted.
464#[derive(Default, Clone, Copy)]
465pub struct PrefetchOptions {
466    /// Enables prefetch for range-scan iterators. Point gets do not inherit this setting.
467    pub prefetch: bool,
468    /// Caller and tracing classification for large queries. Leaf SST iterators do not consume it.
469    pub for_large_query: bool,
470}
471
472impl PrefetchOptions {
473    pub fn prefetch_for_large_range_scan() -> Self {
474        Self {
475            prefetch: true,
476            for_large_query: true,
477        }
478    }
479
480    pub fn prefetch_for_small_range_scan() -> Self {
481        Self {
482            prefetch: true,
483            for_large_query: false,
484        }
485    }
486
487    pub fn new(prefetch: bool, for_large_query: bool) -> Self {
488        Self {
489            prefetch,
490            for_large_query,
491        }
492    }
493}
494
495impl From<TracedPrefetchOptions> for PrefetchOptions {
496    fn from(value: TracedPrefetchOptions) -> Self {
497        Self {
498            prefetch: value.prefetch,
499            for_large_query: value.for_large_query,
500        }
501    }
502}
503
504impl From<PrefetchOptions> for TracedPrefetchOptions {
505    fn from(value: PrefetchOptions) -> Self {
506        Self {
507            prefetch: value.prefetch,
508            for_large_query: value.for_large_query,
509        }
510    }
511}
512
513#[derive(Default, Clone)]
514pub struct ReadOptions {
515    /// A hint for prefix key to check the SST filter.
516    /// If the `prefix_hint` is not None, it should be included in
517    /// `key` or `key_range` in the read API.
518    pub prefix_hint: Option<Bytes>,
519    pub prefetch_options: PrefetchOptions,
520    pub cache_policy: CachePolicy,
521}
522
523impl From<TracedReadOptions> for ReadOptions {
524    fn from(value: TracedReadOptions) -> Self {
525        Self {
526            prefix_hint: value.prefix_hint.map(|b| b.into()),
527            prefetch_options: value.prefetch_options.into(),
528            cache_policy: value.cache_policy.into(),
529        }
530    }
531}
532
533impl ReadOptions {
534    pub fn into_traced_read_options(
535        self,
536        table_id: TableId,
537        epoch: Option<HummockReadEpoch>,
538        table_option: TableOption,
539    ) -> TracedReadOptions {
540        let value = self;
541        let (read_version_from_backup, read_committed) = match epoch {
542            None | Some(HummockReadEpoch::NoWait(_)) => (false, false),
543            Some(HummockReadEpoch::Backup(_)) => (true, false),
544            Some(HummockReadEpoch::Committed(_))
545            | Some(HummockReadEpoch::BatchQueryCommitted(_, _))
546            | Some(HummockReadEpoch::TimeTravel(_))
547            | Some(HummockReadEpoch::LatestCommitted(_)) => (false, true),
548        };
549        TracedReadOptions {
550            prefix_hint: value.prefix_hint.map(|b| b.into()),
551            prefetch_options: value.prefetch_options.into(),
552            cache_policy: value.cache_policy.into(),
553            retention_seconds: table_option.retention_seconds,
554            table_id: table_id.into(),
555            read_version_from_backup,
556            read_committed,
557        }
558    }
559}
560
561pub fn gen_min_epoch(base_epoch: u64, retention_seconds: Option<u32>) -> u64 {
562    match retention_seconds {
563        Some(retention_seconds_u32) => {
564            if base_epoch == MAX_EPOCH {
565                panic!("generate min epoch for MAX_EPOCH");
566            }
567            Epoch(base_epoch)
568                .subtract_ms(retention_seconds_u32 as u64 * 1000)
569                .0
570        }
571        None => 0,
572    }
573}
574
575pub trait CheckOldValueEquality = Fn(&Bytes, &Bytes) -> bool + Send + Sync;
576
577pub static CHECK_BYTES_EQUAL: LazyLock<Arc<dyn CheckOldValueEquality>> =
578    LazyLock::new(|| Arc::new(|first: &Bytes, second: &Bytes| first == second));
579
580#[derive(Default, Clone)]
581pub enum OpConsistencyLevel {
582    #[default]
583    Inconsistent,
584    ConsistentOldValue {
585        check_old_value: Arc<dyn CheckOldValueEquality>,
586        /// whether should store the old value
587        is_log_store: bool,
588    },
589}
590
591impl Debug for OpConsistencyLevel {
592    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
593        match self {
594            OpConsistencyLevel::Inconsistent => f.write_str("OpConsistencyLevel::Inconsistent"),
595            OpConsistencyLevel::ConsistentOldValue { is_log_store, .. } => f
596                .debug_struct("OpConsistencyLevel::ConsistentOldValue")
597                .field("is_log_store", is_log_store)
598                .finish(),
599        }
600    }
601}
602
603impl PartialEq<Self> for OpConsistencyLevel {
604    fn eq(&self, other: &Self) -> bool {
605        matches!(
606            (self, other),
607            (
608                OpConsistencyLevel::Inconsistent,
609                OpConsistencyLevel::Inconsistent
610            ) | (
611                OpConsistencyLevel::ConsistentOldValue {
612                    is_log_store: true,
613                    ..
614                },
615                OpConsistencyLevel::ConsistentOldValue {
616                    is_log_store: true,
617                    ..
618                },
619            ) | (
620                OpConsistencyLevel::ConsistentOldValue {
621                    is_log_store: false,
622                    ..
623                },
624                OpConsistencyLevel::ConsistentOldValue {
625                    is_log_store: false,
626                    ..
627                },
628            )
629        )
630    }
631}
632
633impl Eq for OpConsistencyLevel {}
634
635impl OpConsistencyLevel {
636    pub fn update(&mut self, new_level: &OpConsistencyLevel) {
637        assert_ne!(self, new_level);
638        *self = new_level.clone()
639    }
640}
641
642#[derive(Clone)]
643pub struct NewLocalOptions {
644    pub table_id: TableId,
645    pub fragment_id: FragmentId,
646    /// Whether the operation is consistent. The term `consistent` requires the following:
647    ///
648    /// 1. A key cannot be inserted or deleted for more than once, i.e. inserting to an existing
649    ///    key or deleting an non-existing key is not allowed.
650    ///
651    /// 2. The old value passed from
652    ///    `update` and `delete` should match the original stored value.
653    pub op_consistency_level: OpConsistencyLevel,
654    pub table_option: TableOption,
655
656    /// Indicate if this is replicated. If it is, we should not
657    /// upload its `ReadVersions`.
658    pub is_replicated: bool,
659
660    /// The vnode bitmap for the local state store instance
661    pub vnodes: Arc<Bitmap>,
662
663    pub upload_on_flush: bool,
664}
665
666impl From<TracedNewLocalOptions> for NewLocalOptions {
667    fn from(value: TracedNewLocalOptions) -> Self {
668        Self {
669            table_id: value.table_id.into(),
670            fragment_id: value.fragment_id.into(),
671            op_consistency_level: match value.op_consistency_level {
672                TracedOpConsistencyLevel::Inconsistent => OpConsistencyLevel::Inconsistent,
673                TracedOpConsistencyLevel::ConsistentOldValue => {
674                    OpConsistencyLevel::ConsistentOldValue {
675                        check_old_value: CHECK_BYTES_EQUAL.clone(),
676                        // TODO: for simplicity, set it to false
677                        is_log_store: false,
678                    }
679                }
680            },
681            table_option: value.table_option.into(),
682            is_replicated: value.is_replicated,
683            vnodes: Arc::new(value.vnodes.into()),
684            upload_on_flush: value.upload_on_flush,
685        }
686    }
687}
688
689impl From<NewLocalOptions> for TracedNewLocalOptions {
690    fn from(value: NewLocalOptions) -> Self {
691        Self {
692            table_id: value.table_id.into(),
693            fragment_id: value.fragment_id.into(),
694            op_consistency_level: match value.op_consistency_level {
695                OpConsistencyLevel::Inconsistent => TracedOpConsistencyLevel::Inconsistent,
696                OpConsistencyLevel::ConsistentOldValue { .. } => {
697                    TracedOpConsistencyLevel::ConsistentOldValue
698                }
699            },
700            table_option: value.table_option.into(),
701            is_replicated: value.is_replicated,
702            vnodes: value.vnodes.as_ref().clone().into(),
703            upload_on_flush: value.upload_on_flush,
704        }
705    }
706}
707
708impl NewLocalOptions {
709    pub fn new(
710        table_id: TableId,
711        fragment_id: FragmentId,
712        op_consistency_level: OpConsistencyLevel,
713        table_option: TableOption,
714        vnodes: Arc<Bitmap>,
715        upload_on_flush: bool,
716    ) -> Self {
717        NewLocalOptions {
718            table_id,
719            fragment_id,
720            op_consistency_level,
721            table_option,
722            is_replicated: false,
723            vnodes,
724            upload_on_flush,
725        }
726    }
727
728    pub fn new_replicated(
729        table_id: TableId,
730        fragment_id: FragmentId,
731        op_consistency_level: OpConsistencyLevel,
732        table_option: TableOption,
733        vnodes: Arc<Bitmap>,
734    ) -> Self {
735        NewLocalOptions {
736            table_id,
737            fragment_id,
738            op_consistency_level,
739            table_option,
740            is_replicated: true,
741            vnodes,
742            upload_on_flush: false,
743        }
744    }
745
746    pub fn for_test(table_id: TableId) -> Self {
747        Self {
748            table_id,
749            fragment_id: FragmentId::default(),
750            op_consistency_level: OpConsistencyLevel::Inconsistent,
751            table_option: TableOption {
752                retention_seconds: None,
753            },
754            is_replicated: false,
755            vnodes: Arc::new(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
756            upload_on_flush: true,
757        }
758    }
759}
760
761#[derive(Clone)]
762pub struct InitOptions {
763    pub epoch: EpochPair,
764}
765
766impl InitOptions {
767    pub fn new(epoch: EpochPair) -> Self {
768        Self { epoch }
769    }
770}
771
772impl From<InitOptions> for TracedInitOptions {
773    fn from(value: InitOptions) -> Self {
774        TracedInitOptions {
775            epoch: value.epoch.into(),
776        }
777    }
778}
779
780impl From<TracedInitOptions> for InitOptions {
781    fn from(value: TracedInitOptions) -> Self {
782        InitOptions {
783            epoch: value.epoch.into(),
784        }
785    }
786}
787
788#[derive(Clone, Debug)]
789pub struct SealCurrentEpochOptions {
790    pub table_watermarks: Option<(WatermarkDirection, Vec<VnodeWatermark>, WatermarkSerdeType)>,
791    pub switch_op_consistency_level: Option<OpConsistencyLevel>,
792}
793
794impl From<SealCurrentEpochOptions> for TracedSealCurrentEpochOptions {
795    fn from(value: SealCurrentEpochOptions) -> Self {
796        TracedSealCurrentEpochOptions {
797            table_watermarks: value.table_watermarks.map(
798                |(direction, watermarks, watermark_type)| {
799                    (
800                        direction == WatermarkDirection::Ascending,
801                        watermarks
802                            .into_iter()
803                            .map(|watermark| {
804                                let pb_watermark = PbVnodeWatermark::from(watermark);
805                                Message::encode_to_vec(&pb_watermark)
806                            })
807                            .collect(),
808                        PbWatermarkSerdeType::from(watermark_type) as i32,
809                    )
810                },
811            ),
812            switch_op_consistency_level: value
813                .switch_op_consistency_level
814                .map(|level| matches!(level, OpConsistencyLevel::ConsistentOldValue { .. })),
815        }
816    }
817}
818
819impl From<TracedSealCurrentEpochOptions> for SealCurrentEpochOptions {
820    fn from(value: TracedSealCurrentEpochOptions) -> SealCurrentEpochOptions {
821        SealCurrentEpochOptions {
822            table_watermarks: value.table_watermarks.map(
823                |(is_ascending, watermarks, watermark_serde_type)| {
824                    (
825                        if is_ascending {
826                            WatermarkDirection::Ascending
827                        } else {
828                            WatermarkDirection::Descending
829                        },
830                        watermarks
831                            .into_iter()
832                            .map(|serialized_watermark| {
833                                Message::decode(serialized_watermark.as_slice())
834                                    .map(|pb: PbVnodeWatermark| VnodeWatermark::from(pb))
835                                    .expect("should not failed")
836                            })
837                            .collect(),
838                        match PbWatermarkSerdeType::try_from(watermark_serde_type).unwrap() {
839                            PbWatermarkSerdeType::TypeUnspecified => unreachable!(),
840                            PbWatermarkSerdeType::PkPrefix => WatermarkSerdeType::PkPrefix,
841                            PbWatermarkSerdeType::NonPkPrefix => WatermarkSerdeType::NonPkPrefix,
842                            PbWatermarkSerdeType::Value => WatermarkSerdeType::Value,
843                        },
844                    )
845                },
846            ),
847            switch_op_consistency_level: value.switch_op_consistency_level.map(|enable| {
848                if enable {
849                    OpConsistencyLevel::ConsistentOldValue {
850                        check_old_value: CHECK_BYTES_EQUAL.clone(),
851                        is_log_store: false,
852                    }
853                } else {
854                    OpConsistencyLevel::Inconsistent
855                }
856            }),
857        }
858    }
859}
860
861impl SealCurrentEpochOptions {
862    pub fn for_test() -> Self {
863        Self {
864            table_watermarks: None,
865            switch_op_consistency_level: None,
866        }
867    }
868}
risingwave_storage/store.rs

risingwave_storage/
store.rs
