risingwave_stream/executor/
chain.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::executor::prelude::*;
16use crate::task::CreateMviewProgressReporter;
17
18/// [`ChainExecutor`] is an executor that enables synchronization between the existing stream and
19/// newly appended executors. Currently, [`ChainExecutor`] is mainly used to implement MV on MV
20/// feature. It pipes new data of existing MVs to newly created MV only all of the old data in the
21/// existing MVs are dispatched.
22pub struct ChainExecutor {
23    upstream: Executor,
24
25    progress: CreateMviewProgressReporter,
26}
27
28impl ChainExecutor {
29    pub fn new(upstream: Executor, progress: CreateMviewProgressReporter) -> Self {
30        Self { upstream, progress }
31    }
32
33    #[try_stream(ok = Message, error = StreamExecutorError)]
34    async fn execute_inner(mut self) {
35        let mut upstream = self.upstream.execute();
36
37        // 1. Poll the upstream to get the first barrier.
38        let barrier = expect_first_barrier(&mut upstream).await?;
39
40        // The first barrier message should be propagated.
41        yield Message::Barrier(barrier);
42
43        // 2. Continuously consume the upstream. Report completion on the first barrier after the
44        // initial barrier, before propagating it, so the progress is collected with that barrier.
45        let mut progress_finished = false;
46        #[for_await]
47        for msg in upstream {
48            let msg = msg?;
49            if !progress_finished && let Message::Barrier(barrier) = &msg {
50                self.progress.finish(barrier.epoch, 0);
51                progress_finished = true;
52            }
53            yield msg;
54        }
55    }
56}
57
58impl Execute for ChainExecutor {
59    fn execute(self: Box<Self>) -> super::BoxedMessageStream {
60        self.execute_inner().boxed()
61    }
62}
63
64#[cfg(test)]
65mod test {
66
67    use futures::StreamExt;
68    use risingwave_common::array::StreamChunk;
69    use risingwave_common::array::stream_chunk::StreamChunkTestExt;
70    use risingwave_common::catalog::{Field, Schema};
71    use risingwave_common::types::DataType;
72    use risingwave_common::util::epoch::test_epoch;
73    use risingwave_pb::stream_plan::Dispatcher;
74
75    use super::ChainExecutor;
76    use crate::executor::test_utils::MockSource;
77    use crate::executor::{AddMutation, Barrier, Execute, Message, Mutation, StreamKey};
78    use crate::task::CreateMviewProgressReporter;
79    use crate::task::barrier_test_utils::LocalBarrierTestEnv;
80
81    #[tokio::test]
82    async fn test_basic() {
83        let test_env = LocalBarrierTestEnv::for_test().await;
84        let barrier_manager = test_env.local_barrier_manager.clone();
85        let progress = CreateMviewProgressReporter::for_test(barrier_manager);
86        let actor_id = progress.actor_id();
87
88        let schema = Schema::new(vec![Field::unnamed(DataType::Int64)]);
89        let upstream = MockSource::with_messages(vec![
90            Message::Barrier(Barrier::new_test_barrier(test_epoch(1)).with_mutation(
91                Mutation::Add(AddMutation {
92                    adds: maplit::hashmap! {
93                        0.into() => vec![Dispatcher {
94                            downstream_actor_id: vec![actor_id],
95                            ..Default::default()
96                        }],
97                    },
98                    added_actors: maplit::hashset! { actor_id },
99                    ..Default::default()
100                }),
101            )),
102            Message::Chunk(StreamChunk::from_pretty("I\n + 3")),
103            Message::Chunk(StreamChunk::from_pretty("I\n + 4")),
104            Message::Barrier(Barrier::new_test_barrier(test_epoch(2))),
105        ])
106        .into_executor(schema.clone(), StreamKey::new());
107
108        let chain = ChainExecutor::new(upstream, progress);
109
110        let mut chain = chain.boxed().execute();
111        chain.next().await;
112
113        assert_eq!(
114            chain.next().await.transpose().unwrap(),
115            Some(Message::Chunk(StreamChunk::from_pretty("I\n + 3")))
116        );
117        assert_eq!(
118            chain.next().await.transpose().unwrap(),
119            Some(Message::Chunk(StreamChunk::from_pretty("I\n + 4")))
120        );
121        assert!(matches!(
122            chain.next().await.transpose().unwrap(),
123            Some(Message::Barrier(_))
124        ));
125    }
126}
risingwave_stream/executor/chain.rs

risingwave_stream/executor/
chain.rs
