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

src/storage/hummock_trace/src/collector.rs

@@ -316,6 +316,10 @@ impl TraceSpan {
         Self::new_global_op(Operation::Flush(delete_range), storage_type)
     }
 
+    pub fn new_try_flush_span(storage_type: StorageType) -> MayTraceSpan {
+        Self::new_global_op(Operation::Flush(vec![]), storage_type)
+    }
+
     pub fn new_meta_message_span(resp: SubscribeResponse) -> MayTraceSpan {
         Self::new_global_op(
             Operation::MetaMessage(Box::new(TracedSubResp::from(resp))),

src/storage/hummock_trace/src/record.rs

@@ -287,6 +287,7 @@ pub enum OperationResult {
     Get(TraceResult<Option<TracedBytes>>),
     Insert(TraceResult<()>),
     Delete(TraceResult<()>),
+    TryFlush(TraceResult<()>),
     Flush(TraceResult<usize>),
     Iter(TraceResult<()>),
     IterNext(TraceResult<Option<(TracedBytes, TracedBytes)>>),

src/storage/src/hummock/store/local_hummock_storage.rs

@@ -48,6 +48,7 @@ use crate::storage_value::StorageValue;
 use crate::store::*;
 use crate::StateStoreIter;
 
+const MEM_TABLE_SPILL_THRESHOLD: usize = 64 * 1024 * 1024;
 /// `LocalHummockStorage` is a handle for a state table shard to access data from and write data to
 /// the hummock state backend. It is created via `HummockStorage::new_local`.
 pub struct LocalHummockStorage {
@@ -321,6 +322,19 @@ impl LocalStateStore for LocalHummockStorage {
         .await
     }
 
+    async fn try_flush(&mut self) -> StorageResult<()> {
+        if self.mem_table.kv_size.size() > MEM_TABLE_SPILL_THRESHOLD {
+            tracing::info!(
+                "The size of mem table exceeds 64 Mb and spill occurs. table_id {}",
+                self.table_id.table_id()
+            );
+
+            self.flush(vec![]).await?;
+        }
+
+        Ok(())
+    }
+
     fn epoch(&self) -> u64 {
         self.epoch.expect("should have set the epoch")
     }
@@ -343,7 +357,6 @@ impl LocalStateStore for LocalHummockStorage {
     }
 
     fn seal_current_epoch(&mut self, next_epoch: u64) {
-        assert!(!self.is_dirty());
         let prev_epoch = self
             .epoch
             .replace(next_epoch)

src/storage/src/mem_table.rs

@@ -575,6 +575,19 @@ impl<S: StateStoreWrite + StateStoreRead> LocalStateStore for MemtableLocalState
             prev_epoch
         );
     }
+
+    async fn try_flush(&mut self) -> StorageResult<()> {
+        tracing::info!("In memory state store");
+        if self.mem_table.kv_size.size() > 64 * 1024 * 1024 {
+            tracing::info!(
+                "The size of mem table exceeds 64 Mb and spill occurs. table_id {}",
+                self.table_id.table_id()
+            );
+            self.flush(vec![]).await?;
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]

src/storage/src/monitor/monitored_store.rs

@@ -263,7 +263,6 @@ impl<S: LocalStateStore> LocalStateStore for MonitoredStateStore<S> {
         &mut self,
         delete_ranges: Vec<(Bound<Bytes>, Bound<Bytes>)>,
     ) -> impl Future<Output = StorageResult<usize>> + Send + '_ {
-        // TODO: collect metrics
         self.inner
             .flush(delete_ranges)
             .verbose_instrument_await("store_flush")
@@ -285,6 +284,12 @@ impl<S: LocalStateStore> LocalStateStore for MonitoredStateStore<S> {
         // TODO: may collect metrics
         self.inner.seal_current_epoch(next_epoch)
     }
+
+    fn try_flush(&mut self) -> impl Future<Output = StorageResult<()>> + Send + '_ {
+        self.inner
+            .try_flush()
+            .verbose_instrument_await("store_try_flush")
+    }
 }
 
 impl<S: StateStore> StateStore for MonitoredStateStore<S> {

src/storage/src/monitor/traced_store.rs

@@ -212,6 +212,13 @@ impl<S: LocalStateStore> LocalStateStore for TracedStateStore<S> {
         let _span = TraceSpan::new_seal_current_epoch_span(next_epoch, self.storage_type);
         self.inner.seal_current_epoch(next_epoch)
     }
+
+    async fn try_flush(&mut self) -> StorageResult<()> {
+        let span = TraceSpan::new_try_flush_span(self.storage_type);
+        let res = self.inner.try_flush().await;
+        span.may_send_result(OperationResult::TryFlush(res.as_ref().map(|o| *o).into()));
+        res
+    }
 }
 
 impl<S: StateStore> StateStore for TracedStateStore<S> {

src/storage/src/panic_store.rs

@@ -133,6 +133,11 @@ impl LocalStateStore for PanicStateStore {
     fn seal_current_epoch(&mut self, _next_epoch: u64) {
         panic!("should not operate on the panic state store!")
     }
+
+    #[allow(clippy::unused_async)]
+    async fn try_flush(&mut self) -> StorageResult<()> {
+        panic!("should not operate on the panic state store!");
+    }
 }
 
 impl StateStore for PanicStateStore {

src/storage/src/store.rs

@@ -234,6 +234,7 @@ pub trait LocalStateStore: StaticSendSync {
         delete_ranges: Vec<(Bound<Bytes>, Bound<Bytes>)>,
     ) -> impl Future<Output = StorageResult<usize>> + Send + '_;
 
+    fn try_flush(&mut self) -> impl Future<Output = StorageResult<()>> + Send + '_;
     fn epoch(&self) -> u64;
 
     fn is_dirty(&self) -> bool;

src/storage/src/store_impl.rs

@@ -452,6 +452,13 @@ pub mod verify {
             self.actual.flush(delete_ranges).await
         }
 
+        async fn try_flush(&mut self) -> StorageResult<()> {
+            if let Some(expected) = &mut self.expected {
+                expected.try_flush().await?;
+            }
+            self.actual.try_flush().await
+        }
+
         async fn init(&mut self, options: InitOptions) -> StorageResult<()> {
             self.actual.init(options.clone()).await?;
             if let Some(expected) = &mut self.expected {
@@ -828,6 +835,8 @@ pub mod boxed_state_store {
             delete_ranges: Vec<(Bound<Bytes>, Bound<Bytes>)>,
         ) -> StorageResult<usize>;
 
+        async fn try_flush(&mut self) -> StorageResult<()>;
+
         fn epoch(&self) -> u64;
 
         fn is_dirty(&self) -> bool;
@@ -883,6 +892,10 @@ pub mod boxed_state_store {
             self.flush(delete_ranges).await
         }
 
+        async fn try_flush(&mut self) -> StorageResult<()> {
+            self.try_flush().await
+        }
+
         fn epoch(&self) -> u64 {
             self.epoch()
         }
@@ -949,6 +962,10 @@ pub mod boxed_state_store {
             self.deref_mut().flush(delete_ranges)
         }
 
+        fn try_flush(&mut self) -> impl Future<Output = StorageResult<()>> + Send + '_ {
+            self.deref_mut().try_flush()
+        }
+
         fn epoch(&self) -> u64 {
             self.deref().epoch()
         }

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

@@ -903,7 +903,10 @@ where
     }
 
     pub async fn commit(&mut self, new_epoch: EpochPair) -> StreamExecutorResult<()> {
-        assert_eq!(self.epoch(), new_epoch.prev);
+        assert!(self.epoch() >= new_epoch.prev);
+        if self.epoch() >= new_epoch.curr {
+            panic!("Fail to commit, the epoch gap runs out");
+        }
         trace!(
             table_id = %self.table_id,
             epoch = ?self.epoch(),
@@ -1078,6 +1081,12 @@ where
         self.local_store.seal_current_epoch(next_epoch);
         Ok(())
     }
+
+    pub async fn try_flush(&mut self, next_epoch: u64) -> StreamExecutorResult<()> {
+        self.local_store.try_flush().await?;
+        self.local_store.seal_current_epoch(next_epoch);
+        Ok(())
+    }
 }
 
 // Iterator functions

src/stream/src/executor/hash_join.rs

@@ -30,6 +30,7 @@ use risingwave_common::util::epoch::EpochPair;
 use risingwave_common::util::iter_util::ZipEqDebug;
 use risingwave_expr::expr::BoxedExpression;
 use risingwave_expr::ExprError;
+use risingwave_pb::stream_plan::barrier::BarrierKind;
 use risingwave_storage::StateStore;
 use tokio::time::Instant;
 
@@ -795,7 +796,11 @@ impl<K: HashKey, S: StateStore, const T: JoinTypePrimitive> HashJoinExecutor<K,
                 }
                 AlignedMessage::Barrier(barrier) => {
                     let barrier_start_time = Instant::now();
-                    self.flush_data(barrier.epoch).await?;
+                    if barrier.kind == BarrierKind::Barrier {
+                        self.try_flush_data(barrier.epoch).await?;
+                    } else {
+                        self.flush_data(barrier.epoch).await?;
+                    }
 
                     // Update the vnode bitmap for state tables of both sides if asked.
                     if let Some(vnode_bitmap) = barrier.as_update_vnode_bitmap(self.ctx.id) {
@@ -847,6 +852,14 @@ impl<K: HashKey, S: StateStore, const T: JoinTypePrimitive> HashJoinExecutor<K,
         Ok(())
     }
 
+    async fn try_flush_data(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
+        // All changes to the state has been buffered in the mem-table of the state table. Just
+        // `commit` them here.
+        self.side_l.ht.try_flush(epoch).await?;
+        self.side_r.ht.try_flush(epoch).await?;
+        Ok(())
+    }
+
     // We need to manually evict the cache.
     fn evict_cache(
         side_update: &mut JoinSide<K, S>,

src/stream/src/executor/managed_state/join/mod.rs

@@ -464,6 +464,12 @@ impl<K: HashKey, S: StateStore> JoinHashMap<K, S> {
         Ok(())
     }
 
+    pub async fn try_flush(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
+        self.state.table.try_flush(epoch.curr).await?;
+        self.degree_state.table.try_flush(epoch.curr).await?;
+        Ok(())
+    }
+
     /// Insert a join row
     #[allow(clippy::unused_async)]
     pub async fn insert(&mut self, key: &K, value: JoinRow<impl Row>) -> StreamExecutorResult<()> {