refine update watermark

Author: Liang Zhao

src/stream/src/executor/hash_agg.rs

@@ -396,13 +396,17 @@ impl<K: HashKey, S: StateStore> HashAggExecutor<K, S> {
             ref total_lookup_count,
             ref metrics,
             ref chunk_size,
-            buffered_watermarks: ref _buffered_watermarks,
+            ref buffered_watermarks,
             ..
         }: &'a mut HashAggExecutorExtra<K, S>,
         agg_groups: &'a mut AggGroupMap<K, S>,
         epoch: EpochPair,
     ) {
-        // TODO("https://github.com/risingwavelabs/risingwave/issues/6112"): use buffered_watermarks[0] do some state cleaning
+        let state_clean_watermark = buffered_watermarks
+            .first()
+            .and_then(|opt_watermark| opt_watermark.as_ref())
+            .map(|watermark| watermark.val.clone());
+
         let actor_id_str = ctx.id.to_string();
         metrics
             .agg_lookup_miss_count
@@ -483,10 +487,16 @@ impl<K: HashKey, S: StateStore> HashAggExecutor<K, S> {
             }
 
             // Commit all state tables.
-            futures::future::try_join_all(
-                iter_table_storage(storages).map(|state_table| state_table.commit(epoch)),
-            )
+            futures::future::try_join_all(iter_table_storage(storages).map(|state_table| async {
+                if let Some(watermark) = state_clean_watermark.as_ref() {
+                    state_table.update_watermark(watermark.clone())
+                };
+                state_table.commit(epoch).await
+            }))
             .await?;
+            if let Some(watermark) = state_clean_watermark.as_ref() {
+                result_table.update_watermark(watermark.clone());
+            };
             result_table.commit(epoch).await?;
 
             // Evict cache to target capacity.
@@ -495,8 +505,14 @@ impl<K: HashKey, S: StateStore> HashAggExecutor<K, S> {
             // Nothing to flush.
             // Call commit on state table to increment the epoch.
             iter_table_storage(storages).for_each(|state_table| {
+                if let Some(watermark) = state_clean_watermark.as_ref() {
+                    state_table.update_watermark(watermark.clone())
+                };
                 state_table.commit_no_data_expected(epoch);
             });
+            if let Some(watermark) = state_clean_watermark.as_ref() {
+                result_table.update_watermark(watermark.clone());
+            };
             result_table.commit_no_data_expected(epoch);
             return Ok(());
         }