refactor

Author: Chiro11

src/stream/src/executor/exchange/input.rs

@@ -15,12 +15,10 @@
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
-use anyhow::anyhow;
 use either::Either;
 use local_input::LocalInputStreamInner;
 use pin_project::pin_project;
 use risingwave_common::util::addr::{HostAddr, is_local_address};
-use tokio::sync::mpsc;
 
 use super::permit::Receiver;
 use crate::executor::prelude::*;
@@ -89,25 +87,6 @@ pub(crate) fn apply_dispatcher_barrier(
         .extend(dispatcher_barrier.passed_actors);
 }
 
-pub(crate) async fn process_dispatcher_msg(
-    dispatcher_msg: DispatcherMessage,
-    barrier_rx: &mut mpsc::UnboundedReceiver<Barrier>,
-) -> StreamExecutorResult<Message> {
-    let msg = match dispatcher_msg {
-        DispatcherMessage::Chunk(chunk) => Message::Chunk(chunk),
-        DispatcherMessage::Barrier(barrier) => {
-            let mut recv_barrier = barrier_rx
-                .recv()
-                .await
-                .ok_or_else(|| anyhow!("end of barrier recv"))?;
-            apply_dispatcher_barrier(&mut recv_barrier, barrier);
-            Message::Barrier(recv_barrier)
-        }
-        DispatcherMessage::Watermark(watermark) => Message::Watermark(watermark),
-    };
-    Ok(msg)
-}
-
 impl LocalInput {
     pub fn new(channel: Receiver, upstream_actor_id: ActorId) -> Self {
         Self {

src/stream/src/executor/merge.rs

@@ -16,18 +16,14 @@ use std::collections::VecDeque;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
-use anyhow::Context as _;
-use futures::future::try_join_all;
 use risingwave_common::array::StreamChunkBuilder;
 use risingwave_common::config::MetricLevel;
 use tokio::sync::mpsc;
 use tokio::time::Instant;
 
 use super::exchange::input::BoxedActorInput;
 use super::*;
-use crate::executor::exchange::input::{
-    assert_equal_dispatcher_barrier, new_input, process_dispatcher_msg,
-};
+use crate::executor::exchange::input::{apply_dispatcher_barrier, assert_equal_dispatcher_barrier};
 use crate::executor::prelude::*;
 use crate::task::LocalBarrierManager;
 
@@ -242,31 +238,43 @@ impl MergeExecutor {
         // Channels that're blocked by the barrier to align.
         let mut start_time = Instant::now();
         pin_mut!(select_all);
-        while let Some(msg) = select_all.next().await {
+
+        let mut barrier_buffer = DispatchBarrierBuffer {
+            buffer: VecDeque::new(),
+            curr_fragment_id: self.upstream_fragment_id,
+            barrier_rx: self.barrier_rx,
+            actor_id,
+            local_barrier_manager: self.local_barrier_manager,
+            metrics: self.metrics.clone(),
+            fragment_id: self.fragment_id,
+        };
+
+        loop {
+            let msg = barrier_buffer.select_once(&mut select_all).await?;
             metrics
                 .actor_input_buffer_blocking_duration_ns
                 .inc_by(start_time.elapsed().as_nanos() as u64);
-            let msg: DispatcherMessage = msg?;
-            let mut msg: Message = process_dispatcher_msg(msg, &mut self.barrier_rx).await?;
 
-            match &mut msg {
-                Message::Watermark(_) => {
-                    // Do nothing.
-                }
-                Message::Chunk(chunk) => {
+            let msg = match msg {
+                DispatcherMessage::Watermark(watermark) => Message::Watermark(watermark),
+                DispatcherMessage::Chunk(chunk) => {
                     metrics.actor_in_record_cnt.inc_by(chunk.cardinality() as _);
+                    Message::Chunk(chunk)
                 }
-                Message::Barrier(barrier) => {
+                DispatcherMessage::Barrier(barrier) => {
                     tracing::debug!(
                         target: "events::stream::barrier::path",
                         actor_id = actor_id,
                         "receiver receives barrier from path: {:?}",
                         barrier.passed_actors
                     );
-                    barrier.passed_actors.push(actor_id);
+                    let (mut recv_barrier, new_inputs) =
+                        barrier_buffer.pop_barrier_with_inputs().await?;
+                    apply_dispatcher_barrier(&mut recv_barrier, barrier);
+                    recv_barrier.passed_actors.push(actor_id);
 
                     if let Some(Mutation::Update(UpdateMutation { dispatchers, .. })) =
-                        barrier.mutation.as_deref()
+                        recv_barrier.mutation.as_deref()
                         && select_all
                             .upstream_input_ids()
                             .any(|actor_id| dispatchers.contains_key(&actor_id))
@@ -275,8 +283,8 @@ impl MergeExecutor {
                         select_all.flush_buffered_watermarks();
                     }
 
-                    if let Some(update) =
-                        barrier.as_update_merge(self.actor_context.id, self.upstream_fragment_id)
+                    if let Some(update) = recv_barrier
+                        .as_update_merge(self.actor_context.id, self.upstream_fragment_id)
                     {
                         let new_upstream_fragment_id = update
                             .new_upstream_fragment_id
@@ -291,28 +299,12 @@ impl MergeExecutor {
                         // `Watermark` of upstream may become stale after upstream scaling.
                         select_all.flush_buffered_watermarks();
 
-                        if !update.added_upstream_actors.is_empty() {
-                            // Create new upstreams receivers.
-                            let mut new_upstreams: Vec<_> = try_join_all(
-                                update.added_upstream_actors.iter().map(|upstream_actor| {
-                                    new_input(
-                                        &self.local_barrier_manager,
-                                        self.metrics.clone(),
-                                        self.actor_context.id,
-                                        self.fragment_id,
-                                        upstream_actor,
-                                        new_upstream_fragment_id,
-                                    )
-                                }),
-                            )
-                            .await
-                            .context("failed to create upstream receivers")?;
-
+                        if let Some(mut new_upstreams) = new_inputs {
                             // Poll the first barrier from the new upstreams. It must be the same as
                             // the one we polled from original upstreams.
                             for upstream in &mut new_upstreams {
                                 let new_barrier = expect_first_barrier(upstream).await?;
-                                assert_equal_dispatcher_barrier(barrier, &new_barrier);
+                                assert_equal_dispatcher_barrier(&recv_barrier, &new_barrier);
                             }
 
                             // Add the new upstreams to select.
@@ -332,12 +324,16 @@ impl MergeExecutor {
                         );
                     }
 
-                    if barrier.is_stop(actor_id) {
+                    let is_stop = recv_barrier.is_stop(actor_id);
+                    let msg = Message::Barrier(recv_barrier);
+                    if is_stop {
                         yield msg;
                         break;
                     }
+
+                    msg
                 }
-            }
+            };
 
             yield msg;
             start_time = Instant::now();

src/stream/src/executor/mod.rs

@@ -14,8 +14,9 @@
 
 mod prelude;
 
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use std::fmt::Debug;
+use std::future::pending;
 use std::hash::Hash;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -24,6 +25,7 @@ use std::vec;
 
 use await_tree::InstrumentAwait;
 use enum_as_inner::EnumAsInner;
+use futures::future::try_join_all;
 use futures::stream::{BoxStream, FusedStream, FuturesUnordered, StreamFuture};
 use futures::{Stream, StreamExt};
 use itertools::Itertools;
@@ -56,12 +58,14 @@ use risingwave_pb::stream_plan::{
     SubscriptionUpstreamInfo, ThrottleMutation,
 };
 use smallvec::SmallVec;
+use tokio::sync::mpsc;
 use tokio::time::Instant;
 
 use crate::error::StreamResult;
-use crate::executor::exchange::input::BoxedInput;
+use crate::executor::exchange::input::{BoxedActorInput, BoxedInput, new_input};
+use crate::executor::prelude::StreamingMetrics;
 use crate::executor::watermark::BufferedWatermarks;
-use crate::task::{ActorId, FragmentId};
+use crate::task::{ActorId, FragmentId, LocalBarrierManager};
 
 mod actor;
 mod barrier_align;
@@ -1713,3 +1717,87 @@ impl<InputId: Clone + Ord + Hash + std::fmt::Debug, M> DynamicReceivers<InputId,
         self.blocked.is_empty() && self.active.is_empty()
     }
 }
+
+pub(crate) struct DispatchBarrierBuffer {
+    pub buffer: VecDeque<(Barrier, Option<Vec<BoxedActorInput>>)>,
+    pub curr_fragment_id: FragmentId,
+    pub barrier_rx: mpsc::UnboundedReceiver<Barrier>,
+    pub actor_id: ActorId,
+    // cloned/moved from MergeExecutor, just used to create new inputs
+    pub local_barrier_manager: LocalBarrierManager,
+    pub metrics: Arc<StreamingMetrics>,
+    pub fragment_id: FragmentId,
+}
+
+impl DispatchBarrierBuffer {
+    pub async fn select_once(
+        &mut self,
+        stream: &mut (impl Stream<Item = StreamExecutorResult<DispatcherMessage>> + Unpin),
+    ) -> StreamExecutorResult<DispatcherMessage> {
+        tokio::select! {
+            biased;
+
+            msg = stream.next() => {
+                msg.ok_or(StreamExecutorError::channel_closed(
+                    "upstream executor closed unexpectedly",
+                ))?
+            }
+
+            Err(e) = self.concurrently_fetch_barrier_rx() => {
+                Err(e)
+            }
+        }
+    }
+
+    pub async fn pop_barrier_with_inputs(
+        &mut self,
+    ) -> StreamExecutorResult<(Barrier, Option<Vec<BoxedActorInput>>)> {
+        if !self.buffer.is_empty() {
+            let barrier = self.buffer.pop_front().unwrap();
+            Ok(barrier)
+        } else {
+            let barrier = self
+                .barrier_rx
+                .recv()
+                .await
+                .context("barrier channel closed")?;
+            self.pre_apply_barrier(barrier).await
+        }
+    }
+
+    async fn concurrently_fetch_barrier_rx(&mut self) -> StreamExecutorResult<!> {
+        while let Some(barrier) = self.barrier_rx.recv().await {
+            let barrier_with_inputs = self.pre_apply_barrier(barrier).await?;
+            self.buffer.push_back(barrier_with_inputs);
+        }
+        pending::<!>().await
+    }
+
+    async fn pre_apply_barrier(
+        &mut self,
+        barrier: Barrier,
+    ) -> StreamExecutorResult<(Barrier, Option<Vec<BoxedActorInput>>)> {
+        let new_actor_inputs = if let Some(update) =
+            barrier.as_update_merge(self.actor_id, self.curr_fragment_id)
+            && !update.added_upstream_actors.is_empty()
+        {
+            let new_inputs =
+                try_join_all(update.added_upstream_actors.iter().map(|upstream_actor| {
+                    new_input(
+                        &self.local_barrier_manager,
+                        self.metrics.clone(),
+                        self.actor_id,
+                        self.fragment_id,
+                        upstream_actor,
+                        self.curr_fragment_id,
+                    )
+                }))
+                .await
+                .context("failed to create upstream receivers")?;
+            Some(new_inputs)
+        } else {
+            None
+        };
+        Ok((barrier, new_actor_inputs))
+    }
+}