feat(state_table): state table try flush for non-checkpoint barrier

src/compute/tests/integration_tests.rs

@@ -472,7 +472,7 @@ async fn test_row_seq_scan() -> Result<()> {
     ]));
 
     epoch.inc();
-    state.commit(epoch).await.unwrap();
+    state.commit(epoch, true).await.unwrap();
 
     let executor = Box::new(RowSeqScanExecutor::new(
         table,

src/stream/benches/bench_state_table.rs

@@ -119,7 +119,7 @@ async fn run_bench_state_table_inserts<const USE_WATERMARK_CACHE: bool>(
         state_table.insert(row);
     }
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 }
 
 fn bench_state_table_inserts(c: &mut Criterion) {
@@ -179,7 +179,7 @@ async fn run_bench_state_table_chunks<const USE_WATERMARK_CACHE: bool>(
         state_table.write_chunk(chunk);
     }
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 }
 
 fn bench_state_table_write_chunk(c: &mut Criterion) {

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

@@ -898,71 +898,87 @@ where
         }
     }
 
-    pub async fn commit(&mut self, new_epoch: EpochPair) -> StreamExecutorResult<()> {
+    pub async fn commit(
+        &mut self,
+        new_epoch: EpochPair,
+        is_checkpoint: bool,
+    ) -> StreamExecutorResult<()> {
         assert_eq!(self.epoch(), new_epoch.prev);
-        trace!(
-            table_id = %self.table_id,
-            epoch = ?self.epoch(),
-            "commit state table"
-        );
         // Tick the watermark buffer here because state table is expected to be committed once
         // per epoch.
-        self.watermark_buffer_strategy.tick();
-        if !self.is_dirty() {
-            // If the state table is not modified, go fast path.
-            self.local_store
-                .seal_current_epoch(new_epoch.curr, SealCurrentEpochOptions::new());
-            return Ok(());
-        } else {
-            self.seal_current_epoch(new_epoch.curr)
-                .instrument(tracing::info_span!("state_table_commit"))
-                .await?;
-        }
+        match is_checkpoint {
+            false => {
+                self.local_store.try_flush().await?;
+                if !self.is_dirty() {
+                    // If the state table is not modified, go fast path.
+                    self.local_store
+                        .seal_current_epoch(new_epoch.curr, SealCurrentEpochOptions::new());
+                    return Ok(());
+                } else {
+                    self.seal_current_epoch(new_epoch.curr)
+                        .instrument(tracing::info_span!("state_table_commit"))
+                        .await?;
+                }
+            }
+            true => {
+                self.watermark_buffer_strategy.tick();
+                if !self.is_dirty() {
+                    // If the state table is not modified, go fast path.
+                    self.local_store
+                        .seal_current_epoch(new_epoch.curr, SealCurrentEpochOptions::new());
+                    return Ok(());
+                } else {
+                    self.seal_current_epoch(new_epoch.curr)
+                        .instrument(tracing::info_span!("state_table_commit"))
+                        .await?;
+                }
 
-        // Refresh watermark cache if it is out of sync.
-        if USE_WATERMARK_CACHE && !self.watermark_cache.is_synced() {
-            if let Some(ref watermark) = self.prev_cleaned_watermark {
-                let range: (Bound<Once<Datum>>, Bound<Once<Datum>>) =
-                    (Included(once(Some(watermark.clone()))), Unbounded);
-                // NOTE(kwannoel): We buffer `pks` before inserting into watermark cache
-                // because we can't hold an immutable ref (via `iter_key_and_val_with_pk_range`)
-                // and a mutable ref (via `self.watermark_cache.insert`) at the same time.
-                // TODO(kwannoel): We can optimize it with:
-                // 1. Either use `RefCell`.
-                // 2. Or pass in a direct reference to LocalStateStore,
-                //    instead of referencing it indirectly from `self`.
-                //    Similar to how we do for pk_indices.
-                let mut pks = Vec::with_capacity(self.watermark_cache.capacity());
-                {
-                    let mut streams = vec![];
-                    for vnode in self.vnodes().iter_vnodes() {
-                        let stream = self
-                            .iter_with_vnode(vnode, &range, PrefetchOptions::default())
-                            .await?;
-                        streams.push(Box::pin(stream));
-                    }
-                    let merged_stream = merge_sort(streams);
-                    pin_mut!(merged_stream);
-
-                    #[for_await]
-                    for entry in merged_stream.take(self.watermark_cache.capacity()) {
-                        let keyed_row = entry?;
-                        let pk = self.pk_serde.deserialize(keyed_row.key())?;
-                        if !pk.is_null_at(0) {
-                            pks.push(pk);
+                // Refresh watermark cache if it is out of sync.
+                if USE_WATERMARK_CACHE && !self.watermark_cache.is_synced() {
+                    if let Some(ref watermark) = self.prev_cleaned_watermark {
+                        let range: (Bound<Once<Datum>>, Bound<Once<Datum>>) =
+                            (Included(once(Some(watermark.clone()))), Unbounded);
+                        // NOTE(kwannoel): We buffer `pks` before inserting into watermark cache
+                        // because we can't hold an immutable ref (via `iter_key_and_val_with_pk_range`)
+                        // and a mutable ref (via `self.watermark_cache.insert`) at the same time.
+                        // TODO(kwannoel): We can optimize it with:
+                        // 1. Either use `RefCell`.
+                        // 2. Or pass in a direct reference to LocalStateStore,
+                        //    instead of referencing it indirectly from `self`.
+                        //    Similar to how we do for pk_indices.
+                        let mut pks = Vec::with_capacity(self.watermark_cache.capacity());
+                        {
+                            let mut streams = vec![];
+                            for vnode in self.vnodes().iter_vnodes() {
+                                let stream = self
+                                    .iter_with_vnode(vnode, &range, PrefetchOptions::default())
+                                    .await?;
+                                streams.push(Box::pin(stream));
+                            }
+                            let merged_stream = merge_sort(streams);
+                            pin_mut!(merged_stream);
+
+                            #[for_await]
+                            for entry in merged_stream.take(self.watermark_cache.capacity()) {
+                                let keyed_row = entry?;
+                                let pk = self.pk_serde.deserialize(keyed_row.key())?;
+                                if !pk.is_null_at(0) {
+                                    pks.push(pk);
+                                }
+                            }
                         }
-                    }
-                }
 
-                let mut filler = self.watermark_cache.begin_syncing();
-                for pk in pks {
-                    filler.insert_unchecked(DefaultOrdered(pk), ());
-                }
-                filler.finish();
+                        let mut filler = self.watermark_cache.begin_syncing();
+                        for pk in pks {
+                            filler.insert_unchecked(DefaultOrdered(pk), ());
+                        }
+                        filler.finish();
 
-                let n_cache_entries = self.watermark_cache.len();
-                if n_cache_entries < self.watermark_cache.capacity() {
-                    self.watermark_cache.set_table_row_count(n_cache_entries);
+                        let n_cache_entries = self.watermark_cache.len();
+                        if n_cache_entries < self.watermark_cache.capacity() {
+                            self.watermark_cache.set_table_row_count(n_cache_entries);
+                        }
+                    }
                 }
             }
         }
@@ -975,7 +991,6 @@ where
     /// in the epoch will be visible
     pub fn commit_no_data_expected(&mut self, new_epoch: EpochPair) {
         assert_eq!(self.epoch(), new_epoch.prev);
-        assert!(!self.is_dirty());
         // Tick the watermark buffer here because state table is expected to be committed once
         // per epoch.
         self.watermark_buffer_strategy.tick();

src/stream/src/common/table/test_state_table.rs

@@ -76,7 +76,7 @@ async fn test_state_table_update_insert() {
     ]));
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     state_table.delete(OwnedRow::new(vec![
         Some(6_i32.into()),
@@ -132,7 +132,7 @@ async fn test_state_table_update_insert() {
     );
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     let row6_commit = state_table
         .get_row(&OwnedRow::new(vec![Some(6_i32.into())]))
@@ -169,7 +169,7 @@ async fn test_state_table_update_insert() {
     ]));
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     // one epoch: delete (1, 2, 3, 4), insert (5, 6, 7, None), delete(5, 6, 7, None)
     state_table.delete(OwnedRow::new(vec![
@@ -198,7 +198,7 @@ async fn test_state_table_update_insert() {
     assert_eq!(row1, None);
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     let row1_commit = state_table
         .get_row(&OwnedRow::new(vec![Some(1_i32.into())]))
@@ -263,7 +263,7 @@ async fn test_state_table_iter_with_prefix() {
     ]));
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     state_table.insert(OwnedRow::new(vec![
         Some(1_i32.into()),
@@ -392,7 +392,7 @@ async fn test_state_table_iter_with_pk_range() {
     ]));
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     state_table.insert(OwnedRow::new(vec![
         Some(1_i32.into()),
@@ -597,7 +597,7 @@ async fn test_state_table_iter_with_value_indices() {
     }
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     // write [3, 33, 333], [4, 44, 444], [5, 55, 555], [7, 77, 777], [8, 88, 888]into mem_table,
     // [3, 33, 3333], [6, 66, 666], [9, 99, 999] exists in
@@ -785,7 +785,7 @@ async fn test_state_table_iter_with_shuffle_value_indices() {
     }
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     // write [3, 33, 333], [4, 44, 444], [5, 55, 555], [7, 77, 777], [8, 88, 888]into mem_table,
     // [3, 33, 3333], [6, 66, 666], [9, 99, 999] exists in
@@ -1341,7 +1341,7 @@ async fn test_state_table_may_exist() {
     check_may_exist(&state_table, vec![1, 4], vec![2, 3, 6, 12]).await;
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
     let e1 = epoch.prev;
 
     // test may_exist with data only in immutable memtable (e1)
@@ -1382,7 +1382,7 @@ async fn test_state_table_may_exist() {
     check_may_exist(&state_table, vec![1, 4, 6], vec![2, 3, 12]).await;
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
     let e2 = epoch.prev;
 
     // test may_exist with data in immutable memtable (e2), committed ssts (e1)
@@ -1408,7 +1408,7 @@ async fn test_state_table_may_exist() {
     check_may_exist(&state_table, vec![1, 3, 4, 6], vec![2, 12]).await;
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
     let e3 = epoch.prev;
 
     // test may_exist with data in immutable memtable (e3), uncommitted ssts (e2), committed
@@ -1438,7 +1438,7 @@ async fn test_state_table_may_exist() {
     test_env.storage.try_wait_epoch_for_test(e2).await;
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
     let e4 = epoch.prev;
 
     let e3_res = test_env.storage.seal_and_sync_epoch(e3).await.unwrap();
@@ -1545,7 +1545,7 @@ async fn test_state_table_watermark_cache_ignore_null() {
     state_table.update_watermark(watermark, true);
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     let cache = state_table.get_watermark_cache();
     assert_eq!(cache.len(), 1);
@@ -1626,7 +1626,7 @@ async fn test_state_table_watermark_cache_write_chunk() {
     state_table.update_watermark(watermark, true);
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     let inserts_1 = vec![
         (
@@ -1735,7 +1735,7 @@ async fn test_state_table_watermark_cache_write_chunk() {
     state_table.update_watermark(watermark, true);
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     // After sync, we should scan all rows into watermark cache.
     let cache = state_table.get_watermark_cache();
@@ -1832,7 +1832,7 @@ async fn test_state_table_watermark_cache_refill() {
     state_table.update_watermark(watermark, true);
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     // After the first barrier, watermark cache won't be filled.
     let cache = state_table.get_watermark_cache();
@@ -1898,7 +1898,7 @@ async fn test_state_table_iter_prefix_and_sub_range() {
     ]));
 
     epoch.inc();
-    state_table.commit(epoch).await.unwrap();
+    state_table.commit(epoch, true).await.unwrap();
 
     let pk_prefix = OwnedRow::new(vec![Some(1_i32.into())]);
 

src/stream/src/common/table/test_storage_table.rs

@@ -109,7 +109,7 @@ async fn test_storage_table_value_indices() {
     ]));
 
     epoch.inc();
-    state.commit(epoch).await.unwrap();
+    state.commit(epoch, true).await.unwrap();
     test_env.commit_epoch(epoch.prev).await;
 
     let get_row1_res = table
@@ -222,7 +222,7 @@ async fn test_shuffled_column_id_for_storage_table_get_row() {
     ]));
 
     epoch.inc();
-    state.commit(epoch).await.unwrap();
+    state.commit(epoch, true).await.unwrap();
     test_env.commit_epoch(epoch.prev).await;
 
     let get_row1_res = table
@@ -329,7 +329,7 @@ async fn test_row_based_storage_table_point_get_in_batch_mode() {
         Some(222_i32.into()),
     ]));
     epoch.inc();
-    state.commit(epoch).await.unwrap();
+    state.commit(epoch, true).await.unwrap();
     test_env.commit_epoch(epoch.prev).await;
 
     let get_row1_res = table
@@ -444,7 +444,7 @@ async fn test_batch_scan_with_value_indices() {
     ]));
 
     epoch.inc();
-    state.commit(epoch).await.unwrap();
+    state.commit(epoch, true).await.unwrap();
     test_env.commit_epoch(epoch.prev).await;
 
     let iter = table