refactor(state-table): introduce watermark buffer strategy to state table (#7869)

Add an option to control the buffering strategy of the watermarks in state table.

Will be used in #7752

Approved-By: st1page
Approved-By: soundOfDestiny

Author: TennyZhuang

src/stream/src/common/table/mod.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 pub mod state_table;
+mod watermark;
 
 #[cfg(test)]
 pub mod test_state_table;

src/stream/src/common/table/state_table.rs

@@ -51,6 +51,7 @@ use risingwave_storage::table::{compute_chunk_vnode, compute_vnode, Distribution
 use risingwave_storage::StateStore;
 use tracing::trace;
 
+use super::watermark::{WatermarkBufferByEpoch, WatermarkBufferStrategy};
 use crate::executor::{StreamExecutorError, StreamExecutorResult};
 
 /// This num is arbitrary and we may want to improve this choice in the future.
@@ -59,7 +60,10 @@ const STATE_CLEANING_PERIOD_EPOCH: usize = 5;
 /// `StateTable` is the interface accessing relational data in KV(`StateStore`) with
 /// row-based encoding.
 #[derive(Clone)]
-pub struct StateTable<S: StateStore> {
+pub struct StateTable<
+    S: StateStore,
+    W: WatermarkBufferStrategy = WatermarkBufferByEpoch<STATE_CLEANING_PERIOD_EPOCH>,
+> {
     /// Id for this table.
     table_id: TableId,
 
@@ -119,12 +123,11 @@ pub struct StateTable<S: StateStore> {
     /// latest watermark
     cur_watermark: Option<ScalarImpl>,
 
-    /// number of commits with watermark since the last time we did state cleaning by watermark.
-    num_wmked_commits_since_last_clean: usize,
+    watermark_buffer_strategy: W,
 }
 
 // initialize
-impl<S: StateStore> StateTable<S> {
+impl<S: StateStore, W: WatermarkBufferStrategy> StateTable<S, W> {
     /// Create state table from table catalog and store.
     pub async fn from_table_catalog(
         table_catalog: &Table,
@@ -243,7 +246,7 @@ impl<S: StateStore> StateTable<S> {
             epoch: None,
             last_watermark: None,
             cur_watermark: None,
-            num_wmked_commits_since_last_clean: 0,
+            watermark_buffer_strategy: W::default(),
         }
     }
 
@@ -424,7 +427,7 @@ impl<S: StateStore> StateTable<S> {
             epoch: None,
             last_watermark: None,
             cur_watermark: None,
-            num_wmked_commits_since_last_clean: 0,
+            watermark_buffer_strategy: W::default(),
         }
     }
 
@@ -595,6 +598,7 @@ impl<S: StateStore> StateTable<S> {
         std::mem::replace(&mut self.vnodes, new_vnodes)
     }
 }
+
 // write
 impl<S: StateStore> StateTable<S> {
     #[expect(clippy::boxed_local)]
@@ -751,7 +755,7 @@ impl<S: StateStore> StateTable<S> {
         assert_eq!(self.epoch(), new_epoch.prev);
         assert!(!self.is_dirty());
         if self.cur_watermark.is_some() {
-            self.num_wmked_commits_since_last_clean += 1;
+            self.watermark_buffer_strategy.tick();
         }
         self.update_epoch(new_epoch);
     }
@@ -762,15 +766,19 @@ impl<S: StateStore> StateTable<S> {
         buffer: BTreeMap<Bytes, KeyOp>,
         epoch: u64,
     ) -> StreamExecutorResult<()> {
-        let watermark = self.cur_watermark.as_ref().and_then(|cur_watermark_ref| {
-            self.num_wmked_commits_since_last_clean += 1;
-
-            if self.num_wmked_commits_since_last_clean >= STATE_CLEANING_PERIOD_EPOCH {
-                Some(cur_watermark_ref)
+        let watermark = {
+            if let Some(watermark) = self.cur_watermark.take() {
+                self.watermark_buffer_strategy.tick();
+                if !self.watermark_buffer_strategy.apply() {
+                    self.cur_watermark = Some(watermark);
+                    None
+                } else {
+                    Some(watermark)
+                }
             } else {
                 None
             }
-        });
+        };
 
         let mut write_batch = self.local_store.start_write_batch(WriteOptions {
             epoch,
@@ -784,7 +792,7 @@ impl<S: StateStore> StateTable<S> {
         };
         let range_end_suffix = watermark.map(|watermark| {
             serialize_pk(
-                row::once(Some(watermark.clone())),
+                row::once(Some(watermark)),
                 prefix_serializer.as_ref().unwrap(),
             )
         });
@@ -835,10 +843,6 @@ impl<S: StateStore> StateTable<S> {
             }
         }
         write_batch.ingest().await?;
-        if watermark.is_some() {
-            self.last_watermark = self.cur_watermark.take();
-            self.num_wmked_commits_since_last_clean = 0;
-        }
         Ok(())
     }
 
@@ -950,7 +954,7 @@ fn get_second<T, U>(arg: StreamExecutorResult<(T, U)>) -> StreamExecutorResult<U
 }
 
 // Iterator functions
-impl<S: StateStore> StateTable<S> {
+impl<S: StateStore, W: WatermarkBufferStrategy> StateTable<S, W> {
     /// This function scans rows from the relational table.
     pub async fn iter(&self) -> StreamExecutorResult<RowStream<'_, S>> {
         self.iter_with_pk_prefix(row::empty()).await

src/stream/src/common/table/watermark.rs

@@ -0,0 +1,60 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// Strategy to decide how to buffer the watermarks, used for state cleaning.
+pub trait WatermarkBufferStrategy: Default {
+    /// Trigger when a epoch is committed.
+    fn tick(&mut self);
+
+    /// Whether to clear the buffer.
+    ///
+    /// Returns true to indicate that the buffer should be cleared and the strategy states reset.
+    fn apply(&mut self) -> bool;
+}
+
+/// No buffer, apply watermark to memory immediately.
+/// Use the strategy when you want to apply the watermark immediately.
+#[derive(Default, Debug)]
+pub struct WatermarkNoBuffer;
+
+impl WatermarkBufferStrategy for WatermarkNoBuffer {
+    fn tick(&mut self) {}
+
+    fn apply(&mut self) -> bool {
+        true
+    }
+}
+
+/// Buffer the watermark by a epoch period.
+/// The strategy reduced the delete-range calls to storage.
+#[derive(Default, Debug)]
+pub struct WatermarkBufferByEpoch<const PERIOD: usize> {
+    /// number of epochs since the last time we did state cleaning by watermark.
+    buffered_epochs_cnt: usize,
+}
+
+impl<const PERIOD: usize> WatermarkBufferStrategy for WatermarkBufferByEpoch<PERIOD> {
+    fn tick(&mut self) {
+        self.buffered_epochs_cnt += 1;
+    }
+
+    fn apply(&mut self) -> bool {
+        if self.buffered_epochs_cnt >= PERIOD {
+            self.buffered_epochs_cnt = 0;
+            true
+        } else {
+            false
+        }
+    }
+}