use barrier receiver

Author: Liang Zhao

src/stream/src/executor/now.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use async_stack_trace::StackTrace;
 use chrono::NaiveDateTime;
 use futures::{pin_mut, StreamExt};
 use futures_async_stream::try_stream;
@@ -23,14 +24,17 @@ use risingwave_common::util::chunk_coalesce::DataChunkBuilder;
 use risingwave_common::util::epoch::Epoch;
 use risingwave_storage::table::streaming_table::state_table::StateTable;
 use risingwave_storage::StateStore;
+use tokio::sync::mpsc::UnboundedReceiver;
 
 use super::{
-    expect_first_barrier, BoxedExecutor, BoxedMessageStream, Executor, Message, PkIndices,
-    PkIndicesRef, StreamExecutorError, Watermark,
+    Barrier, BoxedMessageStream, Executor, Message, PkIndices, PkIndicesRef, StreamExecutorError,
+    Watermark,
 };
 
 pub struct NowExecutor<S: StateStore> {
-    input: Option<BoxedExecutor>,
+    /// Receiver of barrier channel.
+    barrier_receiver: Option<UnboundedReceiver<Barrier>>,
+
     pk_indices: PkIndices,
     identity: String,
     schema: Schema,
@@ -39,15 +43,19 @@ pub struct NowExecutor<S: StateStore> {
 
 impl<S: StateStore> NowExecutor<S> {
     #[allow(dead_code)]
-    pub fn new(input: BoxedExecutor, executor_id: u64, state_table: StateTable<S>) -> Self {
+    pub fn new(
+        barrier_receiver: UnboundedReceiver<Barrier>,
+        executor_id: u64,
+        state_table: StateTable<S>,
+    ) -> Self {
         let schema = Schema::new(vec![Field {
             data_type: DataType::Timestamp,
             name: String::from("now"),
             sub_fields: vec![],
             type_name: String::default(),
         }]);
         Self {
-            input: Some(input),
+            barrier_receiver: Some(barrier_receiver),
             pk_indices: vec![0],
             identity: format!("NowExecutor {:X}", executor_id),
             schema,
@@ -57,10 +65,14 @@ impl<S: StateStore> NowExecutor<S> {
 
     #[try_stream(ok = Message, error = StreamExecutorError)]
     async fn into_stream(mut self) {
-        let mut input = self.input.take().unwrap().execute();
+        let mut barrier_receiver = self.barrier_receiver.take().unwrap();
 
         // Consume the first barrier message and initialize state table.
-        let barrier = expect_first_barrier(&mut input).await?;
+        let barrier = barrier_receiver
+            .recv()
+            .stack_trace("now_executor_recv_first_barrier")
+            .await
+            .unwrap();
         self.state_table.init_epoch(barrier.epoch);
 
         // The first barrier message should be propagated.
@@ -78,57 +90,51 @@ impl<S: StateStore> NowExecutor<S> {
 
         let mut last_timestamp = state_row.and_then(|row| row[0].clone());
 
-        #[for_await]
-        for msg in input {
-            if let Message::Barrier(barrier) = msg? {
-                if !barrier.is_update() {
-                    let time_millis = Epoch::from(barrier.epoch.curr).as_unix_millis();
-                    let timestamp = Some(ScalarImpl::NaiveDateTime(NaiveDateTimeWrapper::new(
-                        NaiveDateTime::from_timestamp(
-                            (time_millis / 1000) as i64,
-                            (time_millis % 1000 * 1_000_000) as u32,
-                        ),
-                    )));
-
-                    let mut data_chunk_builder = DataChunkBuilder::new(
-                        self.schema().data_types(),
-                        if last_timestamp.is_some() { 2 } else { 1 },
-                    );
-                    if last_timestamp.is_some() {
-                        let chunk_popped = data_chunk_builder
-                            .append_one_row_from_datums([&last_timestamp].into_iter());
-                        debug_assert!(chunk_popped.is_none());
-                    }
-                    let data_chunk = data_chunk_builder
-                        .append_one_row_from_datums([&timestamp].into_iter())
-                        .unwrap();
-                    let mut ops = if last_timestamp.is_some() {
-                        vec![Op::Delete]
-                    } else {
-                        vec![]
-                    };
-                    ops.push(Op::Insert);
-                    let stream_chunk = StreamChunk::from_parts(ops, data_chunk);
-                    yield Message::Chunk(stream_chunk);
-
-                    yield Message::Watermark(Watermark::new(
-                        0,
-                        timestamp.as_ref().unwrap().clone(),
-                    ));
-
-                    if last_timestamp.is_some() {
-                        self.state_table.delete(Row::new(vec![last_timestamp]));
-                    }
-                    self.state_table.insert(Row::new(vec![timestamp.clone()]));
-                    last_timestamp = timestamp;
-
-                    self.state_table.commit(barrier.epoch).await?;
+        while let Some(barrier) = barrier_receiver.recv().await {
+            if !barrier.is_update() {
+                let time_millis = Epoch::from(barrier.epoch.curr).as_unix_millis();
+                let timestamp = Some(ScalarImpl::NaiveDateTime(NaiveDateTimeWrapper::new(
+                    NaiveDateTime::from_timestamp(
+                        (time_millis / 1000) as i64,
+                        (time_millis % 1000 * 1_000_000) as u32,
+                    ),
+                )));
+
+                let mut data_chunk_builder = DataChunkBuilder::new(
+                    self.schema().data_types(),
+                    if last_timestamp.is_some() { 2 } else { 1 },
+                );
+                if last_timestamp.is_some() {
+                    let chunk_popped = data_chunk_builder
+                        .append_one_row_from_datums([&last_timestamp].into_iter());
+                    debug_assert!(chunk_popped.is_none());
+                }
+                let data_chunk = data_chunk_builder
+                    .append_one_row_from_datums([&timestamp].into_iter())
+                    .unwrap();
+                let mut ops = if last_timestamp.is_some() {
+                    vec![Op::Delete]
                 } else {
-                    self.state_table.commit_no_data_expected(barrier.epoch);
+                    vec![]
+                };
+                ops.push(Op::Insert);
+                let stream_chunk = StreamChunk::from_parts(ops, data_chunk);
+                yield Message::Chunk(stream_chunk);
+
+                yield Message::Watermark(Watermark::new(0, timestamp.as_ref().unwrap().clone()));
+
+                if last_timestamp.is_some() {
+                    self.state_table.delete(Row::new(vec![last_timestamp]));
                 }
+                self.state_table.insert(Row::new(vec![timestamp.clone()]));
+                last_timestamp = timestamp;
 
-                yield Message::Barrier(barrier);
+                self.state_table.commit(barrier.epoch).await?;
+            } else {
+                self.state_table.commit_no_data_expected(barrier.epoch);
             }
+
+            yield Message::Barrier(barrier);
         }
     }
 }