feat(source): throttling source based on storage stats

src/compute/tests/integration_tests.rs

@@ -172,7 +172,7 @@ async fn test_table_materialize() -> StreamResult<()> {
         None, // There is no external stream source.
         Arc::new(StreamingMetrics::unused()),
         barrier_rx,
-        u64::MAX,
+        vec![],
         1,
     );
 

src/storage/src/hummock/local_version/pinned_version.rs

@@ -155,6 +155,10 @@ impl PinnedVersion {
     pub fn version(&self) -> HummockVersion {
         self.version.deref().clone()
     }
+
+    pub fn version_ref(&self) -> &HummockVersion {
+        self.version.deref()
+    }
 }
 
 pub(crate) async fn start_pinned_version_worker(

src/storage/src/hummock/mod.rs

@@ -271,6 +271,17 @@ impl HummockStorage {
     pub fn get_pinned_version(&self) -> PinnedVersion {
         self.pinned_version.load().deref().deref().clone()
     }
+
+    pub fn need_write_throttling(&self) -> bool {
+        // TODO #7997 make it configurable
+        const MAX_L0_SUB_LEVEL_NUMBER: usize = 1000;
+        self.pinned_version
+            .load()
+            .version_ref()
+            .levels
+            .values()
+            .any(|levels| levels.l0.as_ref().unwrap().sub_levels.len() > MAX_L0_SUB_LEVEL_NUMBER)
+    }
 }
 
 #[cfg(any(test, feature = "test"))]

src/stream/src/executor/source/fs_source_executor.rs

@@ -27,14 +27,15 @@ use risingwave_connector::source::{
 use risingwave_source::source_desc::{FsSourceDesc, SourceDescBuilder};
 use risingwave_storage::StateStore;
 use tokio::sync::mpsc::UnboundedReceiver;
-use tokio::time::Instant;
 
 use super::executor_core::StreamSourceCore;
 use crate::error::StreamResult;
 use crate::executor::error::StreamExecutorError;
 use crate::executor::monitor::StreamingMetrics;
 use crate::executor::source::reader::SourceReaderStream;
+use crate::executor::throttler::SourceThrottlerImpl;
 use crate::executor::*;
+
 /// [`FsSourceExecutor`] is a streaming source, fir external file systems
 /// such as s3.
 pub struct FsSourceExecutor<S: StateStore> {
@@ -55,8 +56,7 @@ pub struct FsSourceExecutor<S: StateStore> {
     /// Receiver of barrier channel.
     barrier_receiver: Option<UnboundedReceiver<Barrier>>,
 
-    /// Expected barrier latency
-    expected_barrier_latency_ms: u64,
+    throttlers: Vec<SourceThrottlerImpl>,
 }
 
 impl<S: StateStore> FsSourceExecutor<S> {
@@ -68,7 +68,7 @@ impl<S: StateStore> FsSourceExecutor<S> {
         stream_source_core: StreamSourceCore<S>,
         metrics: Arc<StreamingMetrics>,
         barrier_receiver: UnboundedReceiver<Barrier>,
-        expected_barrier_latency_ms: u64,
+        throttlers: Vec<SourceThrottlerImpl>,
         executor_id: u64,
     ) -> StreamResult<Self> {
         Ok(Self {
@@ -79,7 +79,7 @@ impl<S: StateStore> FsSourceExecutor<S> {
             stream_source_core,
             metrics,
             barrier_receiver: Some(barrier_receiver),
-            expected_barrier_latency_ms,
+            throttlers,
         })
     }
 
@@ -336,18 +336,16 @@ impl<S: StateStore> FsSourceExecutor<S> {
 
         yield Message::Barrier(barrier);
 
-        // We allow data to flow for 5 * `expected_barrier_latency_ms` milliseconds, considering
-        // some other latencies like network and cost in Meta.
-        let max_wait_barrier_time_ms = self.expected_barrier_latency_ms as u128 * 5;
-        let mut last_barrier_time = Instant::now();
         let mut self_paused = false;
         let mut metric_row_per_barrier: u64 = 0;
         while let Some(msg) = stream.next().await {
             match msg? {
                 // This branch will be preferred.
                 Either::Left(barrier) => {
-                    last_barrier_time = Instant::now();
-                    if self_paused {
+                    for throttler in &mut self.throttlers {
+                        throttler.on_barrier();
+                    }
+                    if self_paused && self.throttlers.iter().all(|t| !t.should_pause()) {
                         stream.resume_source();
                         self_paused = false;
                     }
@@ -386,10 +384,7 @@ impl<S: StateStore> FsSourceExecutor<S> {
                     chunk,
                     split_offset_mapping,
                 }) => {
-                    if last_barrier_time.elapsed().as_millis() > max_wait_barrier_time_ms {
-                        // Exceeds the max wait barrier time, the source will be paused. Currently
-                        // we can guarantee the source is not paused since it received stream
-                        // chunks.
+                    if !self_paused && self.throttlers.iter().any(|t| t.should_pause()) {
                         self_paused = true;
                         stream.pause_source();
                     }

src/stream/src/executor/source/mod.rs

@@ -21,5 +21,6 @@ pub mod source_executor;
 
 mod reader;
 pub mod state_table_handler;
+pub mod throttler;
 
 pub use state_table_handler::*;

src/stream/src/executor/source/source_executor.rs

@@ -24,17 +24,13 @@ use risingwave_connector::source::{
 use risingwave_source::source_desc::{SourceDesc, SourceDescBuilder};
 use risingwave_storage::StateStore;
 use tokio::sync::mpsc::UnboundedReceiver;
-use tokio::time::Instant;
 
 use super::executor_core::StreamSourceCore;
 use crate::executor::monitor::StreamingMetrics;
 use crate::executor::source::reader::SourceReaderStream;
+use crate::executor::throttler::SourceThrottlerImpl;
 use crate::executor::*;
 
-/// A constant to multiply when calculating the maximum time to wait for a barrier. This is due to
-/// some latencies in network and cost in meta.
-const WAIT_BARRIER_MULTIPLE_TIMES: u128 = 5;
-
 pub struct SourceExecutor<S: StateStore> {
     ctx: ActorContextRef,
 
@@ -53,8 +49,7 @@ pub struct SourceExecutor<S: StateStore> {
     /// Receiver of barrier channel.
     barrier_receiver: Option<UnboundedReceiver<Barrier>>,
 
-    /// Expected barrier latency.
-    expected_barrier_latency_ms: u64,
+    throttlers: Vec<SourceThrottlerImpl>,
 }
 
 impl<S: StateStore> SourceExecutor<S> {
@@ -66,7 +61,7 @@ impl<S: StateStore> SourceExecutor<S> {
         stream_source_core: Option<StreamSourceCore<S>>,
         metrics: Arc<StreamingMetrics>,
         barrier_receiver: UnboundedReceiver<Barrier>,
-        expected_barrier_latency_ms: u64,
+        throttlers: Vec<SourceThrottlerImpl>,
         executor_id: u64,
     ) -> Self {
         Self {
@@ -77,7 +72,7 @@ impl<S: StateStore> SourceExecutor<S> {
             stream_source_core,
             metrics,
             barrier_receiver: Some(barrier_receiver),
-            expected_barrier_latency_ms,
+            throttlers,
         }
     }
 
@@ -300,19 +295,17 @@ impl<S: StateStore> SourceExecutor<S> {
 
         yield Message::Barrier(barrier);
 
-        // We allow data to flow for `WAIT_BARRIER_MULTIPLE_TIMES` * `expected_barrier_latency_ms`
-        // milliseconds, considering some other latencies like network and cost in Meta.
-        let max_wait_barrier_time_ms =
-            self.expected_barrier_latency_ms as u128 * WAIT_BARRIER_MULTIPLE_TIMES;
-        let mut last_barrier_time = Instant::now();
+        // Whether the source is paused.
         let mut self_paused = false;
         let mut metric_row_per_barrier: u64 = 0;
         while let Some(msg) = stream.next().await {
             match msg? {
                 // This branch will be preferred.
                 Either::Left(barrier) => {
-                    last_barrier_time = Instant::now();
-                    if self_paused {
+                    for throttler in &mut self.throttlers {
+                        throttler.on_barrier();
+                    }
+                    if self_paused && self.throttlers.iter().all(|t| !t.should_pause()) {
                         stream.resume_source();
                         self_paused = false;
                     }
@@ -356,10 +349,7 @@ impl<S: StateStore> SourceExecutor<S> {
                     chunk,
                     split_offset_mapping,
                 }) => {
-                    if last_barrier_time.elapsed().as_millis() > max_wait_barrier_time_ms {
-                        // Exceeds the max wait barrier time, the source will be paused. Currently
-                        // we can guarantee the source is not paused since it received stream
-                        // chunks.
+                    if !self_paused && self.throttlers.iter().any(|t| t.should_pause()) {
                         self_paused = true;
                         stream.pause_source();
                     }
@@ -551,7 +541,7 @@ mod tests {
             Some(core),
             Arc::new(StreamingMetrics::unused()),
             barrier_rx,
-            u64::MAX,
+            vec![],
             1,
         );
         let mut executor = Box::new(executor).execute();
@@ -643,7 +633,7 @@ mod tests {
             Some(core),
             Arc::new(StreamingMetrics::unused()),
             barrier_rx,
-            u64::MAX,
+            vec![],
             1,
         );
 

src/stream/src/executor/source/throttler.rs

@@ -0,0 +1,70 @@
+// 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.
+
+use risingwave_storage::StateStoreImpl;
+use tokio::time::Instant;
+
+pub enum SourceThrottlerImpl {
+    MaxWaitBarrier(MaxWaitBarrierThrottler),
+    StateStore(StateStoreImpl),
+}
+
+/// A constant to multiply when calculating the maximum time to wait for a barrier. This is due to
+/// some latencies in network and cost in meta.
+const WAIT_BARRIER_MULTIPLE_TIMES: u128 = 5;
+
+pub struct MaxWaitBarrierThrottler {
+    max_wait_barrier_time_ms: u128,
+    last_barrier_time: Instant,
+}
+
+impl MaxWaitBarrierThrottler {
+    pub fn new(barrier_interval_ms: u128) -> Self {
+        Self {
+            max_wait_barrier_time_ms: barrier_interval_ms * WAIT_BARRIER_MULTIPLE_TIMES,
+            last_barrier_time: Instant::now(),
+        }
+    }
+
+    fn should_pause(&self) -> bool {
+        // We allow data to flow for `WAIT_BARRIER_MULTIPLE_TIMES` *
+        // `expected_barrier_latency_ms` milliseconds, considering some
+        // other latencies like network and cost in Meta.
+        self.last_barrier_time.elapsed().as_millis() > self.max_wait_barrier_time_ms
+    }
+}
+
+impl SourceThrottlerImpl {
+    pub fn should_pause(&self) -> bool {
+        match self {
+            SourceThrottlerImpl::MaxWaitBarrier(inner) => inner.should_pause(),
+            SourceThrottlerImpl::StateStore(inner) => {
+                if let Some(hummock) = inner.as_hummock() {
+                    return hummock.need_write_throttling();
+                }
+                false
+            }
+        }
+    }
+
+    pub fn on_barrier(&mut self) {
+        #[allow(clippy::single_match)]
+        match self {
+            SourceThrottlerImpl::MaxWaitBarrier(inner) => {
+                inner.last_barrier_time = Instant::now();
+            }
+            _ => {}
+        }
+    }
+}

src/stream/src/from_proto/source.rs

@@ -23,6 +23,7 @@ use super::*;
 use crate::executor::source::StreamSourceCore;
 use crate::executor::source_executor::SourceExecutor;
 use crate::executor::state_table_handler::SourceStateTableHandler;
+use crate::executor::throttler::{MaxWaitBarrierThrottler, SourceThrottlerImpl};
 use crate::executor::FsSourceExecutor;
 
 const FS_CONNECTORS: &[&str] = &["s3"];
@@ -44,6 +45,11 @@ impl ExecutorBuilder for SourceExecutorBuilder {
             .lock_barrier_manager()
             .register_sender(params.actor_context.id, sender);
 
+        let mut throttlers = vec![];
+        throttlers.push(SourceThrottlerImpl::MaxWaitBarrier(
+            MaxWaitBarrierThrottler::new(stream.config.barrier_interval_ms as u128),
+        ));
+        throttlers.push(SourceThrottlerImpl::StateStore(stream.state_store.clone()));
         if let Some(source) = &node.source_inner {
             let source_id = TableId::new(source.source_id);
             let source_name = source.source_name.clone();
@@ -108,7 +114,7 @@ impl ExecutorBuilder for SourceExecutorBuilder {
                     stream_source_core,
                     params.executor_stats,
                     barrier_receiver,
-                    stream.config.barrier_interval_ms as u64,
+                    throttlers,
                     params.executor_id,
                 )?))
             } else {
@@ -119,7 +125,7 @@ impl ExecutorBuilder for SourceExecutorBuilder {
                     Some(stream_source_core),
                     params.executor_stats,
                     barrier_receiver,
-                    stream.config.barrier_interval_ms as u64,
+                    throttlers,
                     params.executor_id,
                 )))
             }
@@ -133,7 +139,7 @@ impl ExecutorBuilder for SourceExecutorBuilder {
                 None,
                 params.executor_stats,
                 barrier_receiver,
-                stream.config.barrier_interval_ms as u64,
+                throttlers,
                 params.executor_id,
             )))
         }

src/stream/src/task/stream_manager.rs

@@ -74,7 +74,7 @@ pub struct LocalStreamManagerCore {
     actor_monitor_tasks: HashMap<ActorId, ActorHandle>,
 
     /// The state store implement
-    state_store: StateStoreImpl,
+    pub(crate) state_store: StateStoreImpl,
 
     /// Metrics of the stream manager
     pub(crate) streaming_metrics: Arc<StreamingMetrics>,