risingwave_stream/executor/
merge.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 std::collections::VecDeque;
16use std::pin::Pin;
17use std::task::{Context, Poll};
18
19use risingwave_common::array::StreamChunkBuilder;
20use tokio::sync::mpsc;
21
22use super::exchange::input::BoxedActorInput;
23use super::*;
24use crate::executor::prelude::*;
25use crate::task::LocalBarrierManager;
26
27pub type SelectReceivers = DynamicReceivers<ActorId, ()>;
28
29pub type MergeUpstream = BufferChunks<SelectReceivers>;
30pub type SingletonUpstream = BoxedActorInput;
31
32pub(crate) enum MergeExecutorUpstream {
33    Singleton(SingletonUpstream),
34    Merge(MergeUpstream),
35}
36
37pub(crate) struct MergeExecutorInput {
38    upstream: MergeExecutorUpstream,
39    actor_context: ActorContextRef,
40    upstream_fragment_id: UpstreamFragmentId,
41    local_barrier_manager: LocalBarrierManager,
42    executor_stats: Arc<StreamingMetrics>,
43    pub(crate) info: ExecutorInfo,
44}
45
46impl MergeExecutorInput {
47    pub(crate) fn new(
48        upstream: MergeExecutorUpstream,
49        actor_context: ActorContextRef,
50        upstream_fragment_id: UpstreamFragmentId,
51        local_barrier_manager: LocalBarrierManager,
52        executor_stats: Arc<StreamingMetrics>,
53        info: ExecutorInfo,
54    ) -> Self {
55        Self {
56            upstream,
57            actor_context,
58            upstream_fragment_id,
59            local_barrier_manager,
60            executor_stats,
61            info,
62        }
63    }
64
65    pub(crate) fn into_executor(self, barrier_rx: mpsc::UnboundedReceiver<Barrier>) -> Executor {
66        let fragment_id = self.actor_context.fragment_id;
67        let executor = match self.upstream {
68            MergeExecutorUpstream::Singleton(input) => ReceiverExecutor::new(
69                self.actor_context,
70                fragment_id,
71                self.upstream_fragment_id,
72                input,
73                self.local_barrier_manager,
74                self.executor_stats,
75                barrier_rx,
76            )
77            .boxed(),
78            MergeExecutorUpstream::Merge(inputs) => MergeExecutor::new(
79                self.actor_context,
80                fragment_id,
81                self.upstream_fragment_id,
82                inputs,
83                self.local_barrier_manager,
84                self.executor_stats,
85                barrier_rx,
86            )
87            .boxed(),
88        };
89        (self.info, executor).into()
90    }
91}
92
93impl Stream for MergeExecutorInput {
94    type Item = DispatcherMessageStreamItem;
95
96    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
97        match &mut self.get_mut().upstream {
98            MergeExecutorUpstream::Singleton(input) => input.poll_next_unpin(cx),
99            MergeExecutorUpstream::Merge(inputs) => inputs.poll_next_unpin(cx),
100        }
101    }
102}
103
104mod upstream {
105    use super::*;
106
107    /// Trait unifying operations on [`MergeUpstream`] and [`SingletonUpstream`], so that we can
108    /// reuse code between [`MergeExecutor`] and [`ReceiverExecutor`].
109    pub trait Upstream:
110        Stream<Item = StreamExecutorResult<DispatcherMessage>> + Unpin + Send + 'static
111    {
112        fn upstream_input_ids(&self) -> impl Iterator<Item = ActorId> + '_;
113        fn flush_buffered_watermarks(&mut self);
114        fn update(&mut self, to_add: Vec<BoxedActorInput>, to_remove: &HashSet<ActorId>);
115    }
116
117    impl Upstream for MergeUpstream {
118        fn upstream_input_ids(&self) -> impl Iterator<Item = ActorId> + '_ {
119            self.inner.upstream_input_ids()
120        }
121
122        fn flush_buffered_watermarks(&mut self) {
123            self.inner.flush_buffered_watermarks();
124        }
125
126        fn update(&mut self, to_add: Vec<BoxedActorInput>, to_remove: &HashSet<ActorId>) {
127            if !to_add.is_empty() {
128                self.inner.add_upstreams_from(to_add);
129            }
130            if !to_remove.is_empty() {
131                self.inner.remove_upstreams(to_remove);
132            }
133        }
134    }
135
136    impl Upstream for SingletonUpstream {
137        fn upstream_input_ids(&self) -> impl Iterator<Item = ActorId> + '_ {
138            std::iter::once(self.id())
139        }
140
141        fn flush_buffered_watermarks(&mut self) {
142            // For single input, we won't buffer watermarks so there's nothing to do.
143        }
144
145        fn update(&mut self, to_add: Vec<BoxedActorInput>, to_remove: &HashSet<ActorId>) {
146            assert_eq!(
147                to_remove,
148                &HashSet::from_iter([self.id()]),
149                "the removed upstream actor should be the same as the current input"
150            );
151
152            // Replace the single input.
153            *self = Itertools::exactly_one(to_add.into_iter())
154                .expect("receiver should have exactly one new upstream");
155        }
156    }
157}
158use upstream::Upstream;
159
160/// The core of `MergeExecutor` and `ReceiverExecutor`.
161pub struct MergeExecutorInner<U> {
162    /// The context of the actor.
163    actor_context: ActorContextRef,
164
165    /// Upstream channels.
166    upstream: U,
167
168    /// Belonged fragment id.
169    fragment_id: FragmentId,
170
171    /// Upstream fragment id.
172    upstream_fragment_id: FragmentId,
173
174    local_barrier_manager: LocalBarrierManager,
175
176    /// Streaming metrics.
177    metrics: Arc<StreamingMetrics>,
178
179    barrier_rx: mpsc::UnboundedReceiver<Barrier>,
180}
181
182impl<U> MergeExecutorInner<U> {
183    pub fn new(
184        ctx: ActorContextRef,
185        fragment_id: FragmentId,
186        upstream_fragment_id: FragmentId,
187        upstream: U,
188        local_barrier_manager: LocalBarrierManager,
189        metrics: Arc<StreamingMetrics>,
190        barrier_rx: mpsc::UnboundedReceiver<Barrier>,
191    ) -> Self {
192        Self {
193            actor_context: ctx,
194            upstream,
195            fragment_id,
196            upstream_fragment_id,
197            local_barrier_manager,
198            metrics,
199            barrier_rx,
200        }
201    }
202}
203
204/// `MergeExecutor` merges data from multiple upstream actors and aligns them with barriers.
205pub type MergeExecutor = MergeExecutorInner<MergeUpstream>;
206
207impl MergeExecutor {
208    #[cfg(test)]
209    pub fn for_test(
210        actor_id: impl Into<ActorId>,
211        inputs: Vec<super::exchange::permit::Receiver>,
212        local_barrier_manager: crate::task::LocalBarrierManager,
213        schema: Schema,
214        chunk_size: usize,
215        barrier_rx: Option<mpsc::UnboundedReceiver<Barrier>>,
216    ) -> Self {
217        let actor_id = actor_id.into();
218        use super::exchange::input::LocalInput;
219        use crate::executor::exchange::input::ActorInput;
220
221        let barrier_rx =
222            barrier_rx.unwrap_or_else(|| local_barrier_manager.subscribe_barrier(actor_id));
223
224        let metrics = StreamingMetrics::unused();
225        let actor_ctx = ActorContext::for_test(actor_id);
226        let upstream = Self::new_merge_upstream(
227            inputs
228                .into_iter()
229                .enumerate()
230                .map(|(idx, input)| LocalInput::new(input, ActorId::new(idx as u32)).boxed_input())
231                .collect(),
232            &metrics,
233            &actor_ctx,
234            chunk_size,
235            schema,
236        );
237
238        Self::new(
239            actor_ctx,
240            514.into(),
241            1919.into(),
242            upstream,
243            local_barrier_manager,
244            metrics.into(),
245            barrier_rx,
246        )
247    }
248
249    pub(crate) fn new_merge_upstream(
250        upstreams: Vec<BoxedActorInput>,
251        metrics: &StreamingMetrics,
252        actor_context: &ActorContext,
253        chunk_size: usize,
254        schema: Schema,
255    ) -> MergeUpstream {
256        let merge_barrier_align_duration = Some(
257            metrics
258                .merge_barrier_align_duration
259                .with_guarded_label_values(&[
260                    &actor_context.id.to_string(),
261                    &actor_context.fragment_id.to_string(),
262                ]),
263        );
264
265        BufferChunks::new(
266            // Futures of all active upstreams.
267            SelectReceivers::new(upstreams, None, merge_barrier_align_duration),
268            chunk_size,
269            schema,
270        )
271    }
272}
273
274impl<U> MergeExecutorInner<U>
275where
276    U: Upstream,
277{
278    #[try_stream(ok = Message, error = StreamExecutorError)]
279    pub(crate) async fn execute_inner(mut self: Box<Self>) {
280        let mut upstream = self.upstream;
281        let actor_id = self.actor_context.id;
282
283        let mut metrics = self.metrics.new_actor_input_metrics(
284            actor_id,
285            self.fragment_id,
286            self.upstream_fragment_id,
287        );
288
289        let mut barrier_buffer = DispatchBarrierBuffer::new(
290            self.barrier_rx,
291            actor_id,
292            self.upstream_fragment_id,
293            self.local_barrier_manager,
294            self.metrics.clone(),
295            self.fragment_id,
296            self.actor_context.config.clone(),
297        );
298
299        loop {
300            let msg = barrier_buffer
301                .await_next_message(&mut upstream, &metrics)
302                .await?;
303            let msg = match msg {
304                DispatcherMessage::Watermark(watermark) => Message::Watermark(watermark),
305                DispatcherMessage::Chunk(chunk) => {
306                    metrics.actor_in_record_cnt.inc_by(chunk.cardinality() as _);
307                    Message::Chunk(chunk)
308                }
309                DispatcherMessage::Barrier(barrier) => {
310                    let (barrier, new_inputs) =
311                        barrier_buffer.pop_barrier_with_inputs(barrier).await?;
312
313                    if let Some(Mutation::Update(UpdateMutation { dispatchers, .. })) =
314                        barrier.mutation.as_deref()
315                        && upstream
316                            .upstream_input_ids()
317                            .any(|actor_id| dispatchers.contains_key(&actor_id))
318                    {
319                        // `Watermark` of upstream may become stale after downstream scaling.
320                        upstream.flush_buffered_watermarks();
321                    }
322
323                    if let Some(update) =
324                        barrier.as_update_merge(self.actor_context.id, self.upstream_fragment_id)
325                    {
326                        let new_upstream_fragment_id = update
327                            .new_upstream_fragment_id
328                            .unwrap_or(self.upstream_fragment_id);
329                        let removed_upstream_actor_id: HashSet<_> =
330                            if update.new_upstream_fragment_id.is_some() {
331                                upstream.upstream_input_ids().collect()
332                            } else {
333                                update.removed_upstream_actor_id.iter().copied().collect()
334                            };
335
336                        // `Watermark` of upstream may become stale after upstream scaling.
337                        upstream.flush_buffered_watermarks();
338
339                        // Add and remove upstreams.
340                        upstream.update(new_inputs.unwrap_or_default(), &removed_upstream_actor_id);
341
342                        self.upstream_fragment_id = new_upstream_fragment_id;
343                        metrics = self.metrics.new_actor_input_metrics(
344                            actor_id,
345                            self.fragment_id,
346                            self.upstream_fragment_id,
347                        );
348                    }
349
350                    let is_stop = barrier.is_stop(actor_id);
351                    let msg = Message::Barrier(barrier);
352                    if is_stop {
353                        yield msg;
354                        break;
355                    }
356
357                    msg
358                }
359            };
360
361            yield msg;
362        }
363    }
364}
365
366impl<U> Execute for MergeExecutorInner<U>
367where
368    U: Upstream,
369{
370    fn execute(self: Box<Self>) -> BoxedMessageStream {
371        self.execute_inner().boxed()
372    }
373}
374
375/// A wrapper that buffers the `StreamChunk`s from upstream until no more ready items are available.
376/// Besides, any message other than `StreamChunk` will trigger the buffered `StreamChunk`s
377/// to be emitted immediately along with the message itself.
378pub struct BufferChunks<S: Stream> {
379    inner: S,
380    chunk_builder: StreamChunkBuilder,
381
382    /// The items to be emitted. Whenever there's something here, we should return a `Poll::Ready` immediately.
383    pending_items: VecDeque<S::Item>,
384}
385
386impl<S: Stream> BufferChunks<S> {
387    pub(super) fn new(inner: S, chunk_size: usize, schema: Schema) -> Self {
388        assert!(chunk_size > 0);
389        let chunk_builder = StreamChunkBuilder::new(chunk_size, schema.data_types());
390        Self {
391            inner,
392            chunk_builder,
393            pending_items: VecDeque::new(),
394        }
395    }
396}
397
398impl<S: Stream> std::ops::Deref for BufferChunks<S> {
399    type Target = S;
400
401    fn deref(&self) -> &Self::Target {
402        &self.inner
403    }
404}
405
406impl<S: Stream> std::ops::DerefMut for BufferChunks<S> {
407    fn deref_mut(&mut self) -> &mut Self::Target {
408        &mut self.inner
409    }
410}
411
412impl<S: Stream> Stream for BufferChunks<S>
413where
414    S: Stream<Item = DispatcherMessageStreamItem> + Unpin,
415{
416    type Item = S::Item;
417
418    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
419        loop {
420            if let Some(item) = self.pending_items.pop_front() {
421                return Poll::Ready(Some(item));
422            }
423
424            match self.inner.poll_next_unpin(cx) {
425                Poll::Pending => {
426                    return if let Some(chunk_out) = self.chunk_builder.take() {
427                        Poll::Ready(Some(Ok(MessageInner::Chunk(chunk_out))))
428                    } else {
429                        Poll::Pending
430                    };
431                }
432
433                Poll::Ready(Some(result)) => {
434                    if let Ok(MessageInner::Chunk(chunk)) = result {
435                        for row in chunk.records() {
436                            if let Some(chunk_out) = self.chunk_builder.append_record(row) {
437                                self.pending_items
438                                    .push_back(Ok(MessageInner::Chunk(chunk_out)));
439                            }
440                        }
441                    } else {
442                        return if let Some(chunk_out) = self.chunk_builder.take() {
443                            self.pending_items.push_back(result);
444                            Poll::Ready(Some(Ok(MessageInner::Chunk(chunk_out))))
445                        } else {
446                            Poll::Ready(Some(result))
447                        };
448                    }
449                }
450
451                Poll::Ready(None) => {
452                    // See also the comments in `DynamicReceivers::poll_next`.
453                    unreachable!("Merge should always have upstream inputs");
454                }
455            }
456        }
457    }
458}
459
460#[cfg(test)]
461mod tests {
462    use std::sync::atomic::{AtomicBool, Ordering};
463    use std::time::Duration;
464
465    use assert_matches::assert_matches;
466    use futures::FutureExt;
467    use futures::future::try_join_all;
468    use risingwave_common::array::Op;
469    use risingwave_common::util::epoch::test_epoch;
470    use risingwave_pb::task_service::stream_exchange_service_server::{
471        StreamExchangeService, StreamExchangeServiceServer,
472    };
473    use risingwave_pb::task_service::{GetStreamRequest, GetStreamResponse, PbPermits};
474    use tokio::time::sleep;
475    use tokio_stream::wrappers::ReceiverStream;
476    use tonic::{Request, Response, Status, Streaming};
477
478    use super::*;
479    use crate::executor::exchange::input::{ActorInput, LocalInput, RemoteInput};
480    use crate::executor::exchange::permit::channel_for_test;
481    use crate::executor::{BarrierInner as Barrier, MessageInner as Message};
482    use crate::task::barrier_test_utils::LocalBarrierTestEnv;
483    use crate::task::test_utils::helper_make_local_actor;
484    use crate::task::{NewOutputRequest, TEST_PARTIAL_GRAPH_ID};
485
486    fn build_test_chunk(size: u64) -> StreamChunk {
487        let ops = vec![Op::Insert; size as usize];
488        StreamChunk::new(ops, vec![])
489    }
490
491    #[tokio::test]
492    async fn test_buffer_chunks() {
493        let test_env = LocalBarrierTestEnv::for_test().await;
494
495        let (tx, rx) = channel_for_test();
496        let input = LocalInput::new(rx, 1.into()).boxed_input();
497        let mut buffer = BufferChunks::new(input, 100, Schema::new(vec![]));
498
499        // Send a chunk
500        tx.send(Message::Chunk(build_test_chunk(10)).into())
501            .await
502            .unwrap();
503        assert_matches!(buffer.next().await.unwrap().unwrap(), Message::Chunk(chunk) => {
504            assert_eq!(chunk.ops().len() as u64, 10);
505        });
506
507        // Send 2 chunks and expect them to be merged.
508        tx.send(Message::Chunk(build_test_chunk(10)).into())
509            .await
510            .unwrap();
511        tx.send(Message::Chunk(build_test_chunk(10)).into())
512            .await
513            .unwrap();
514        assert_matches!(buffer.next().await.unwrap().unwrap(), Message::Chunk(chunk) => {
515            assert_eq!(chunk.ops().len() as u64, 20);
516        });
517
518        // Send a watermark.
519        tx.send(
520            Message::Watermark(Watermark {
521                col_idx: 0,
522                data_type: DataType::Int64,
523                val: ScalarImpl::Int64(233),
524            })
525            .into(),
526        )
527        .await
528        .unwrap();
529        assert_matches!(buffer.next().await.unwrap().unwrap(), Message::Watermark(watermark) => {
530            assert_eq!(watermark.val, ScalarImpl::Int64(233));
531        });
532
533        // Send 2 chunks before a watermark. Expect the 2 chunks to be merged and the watermark to be emitted.
534        tx.send(Message::Chunk(build_test_chunk(10)).into())
535            .await
536            .unwrap();
537        tx.send(Message::Chunk(build_test_chunk(10)).into())
538            .await
539            .unwrap();
540        tx.send(
541            Message::Watermark(Watermark {
542                col_idx: 0,
543                data_type: DataType::Int64,
544                val: ScalarImpl::Int64(233),
545            })
546            .into(),
547        )
548        .await
549        .unwrap();
550        assert_matches!(buffer.next().await.unwrap().unwrap(), Message::Chunk(chunk) => {
551            assert_eq!(chunk.ops().len() as u64, 20);
552        });
553        assert_matches!(buffer.next().await.unwrap().unwrap(), Message::Watermark(watermark) => {
554            assert_eq!(watermark.val, ScalarImpl::Int64(233));
555        });
556
557        // Send a barrier.
558        let barrier = Barrier::new_test_barrier(test_epoch(1));
559        test_env.inject_barrier(&barrier, [2.into()]);
560        tx.send(Message::Barrier(barrier.clone().into_dispatcher()).into())
561            .await
562            .unwrap();
563        assert_matches!(buffer.next().await.unwrap().unwrap(), Message::Barrier(Barrier { epoch: barrier_epoch, .. }) => {
564            assert_eq!(barrier_epoch.curr, test_epoch(1));
565        });
566
567        // Send 2 chunks before a barrier. Expect the 2 chunks to be merged and the barrier to be emitted.
568        tx.send(Message::Chunk(build_test_chunk(10)).into())
569            .await
570            .unwrap();
571        tx.send(Message::Chunk(build_test_chunk(10)).into())
572            .await
573            .unwrap();
574        let barrier = Barrier::new_test_barrier(test_epoch(2));
575        test_env.inject_barrier(&barrier, [2.into()]);
576        tx.send(Message::Barrier(barrier.clone().into_dispatcher()).into())
577            .await
578            .unwrap();
579        assert_matches!(buffer.next().await.unwrap().unwrap(), Message::Chunk(chunk) => {
580            assert_eq!(chunk.ops().len() as u64, 20);
581        });
582        assert_matches!(buffer.next().await.unwrap().unwrap(), Message::Barrier(Barrier { epoch: barrier_epoch, .. }) => {
583            assert_eq!(barrier_epoch.curr, test_epoch(2));
584        });
585    }
586
587    #[tokio::test]
588    async fn test_merger() {
589        const CHANNEL_NUMBER: usize = 10;
590        let mut txs = Vec::with_capacity(CHANNEL_NUMBER);
591        let mut rxs = Vec::with_capacity(CHANNEL_NUMBER);
592        for _i in 0..CHANNEL_NUMBER {
593            let (tx, rx) = channel_for_test();
594            txs.push(tx);
595            rxs.push(rx);
596        }
597        let barrier_test_env = LocalBarrierTestEnv::for_test().await;
598        let actor_id = 233.into();
599        let mut handles = Vec::with_capacity(CHANNEL_NUMBER);
600
601        let epochs = (10..1000u64)
602            .step_by(10)
603            .map(|idx| (idx, test_epoch(idx)))
604            .collect_vec();
605        let mut prev_epoch = 0;
606        let prev_epoch = &mut prev_epoch;
607        let barriers: HashMap<_, _> = epochs
608            .iter()
609            .map(|(_, epoch)| {
610                let barrier = Barrier::with_prev_epoch_for_test(*epoch, *prev_epoch);
611                *prev_epoch = *epoch;
612                barrier_test_env.inject_barrier(&barrier, [actor_id]);
613                (*epoch, barrier)
614            })
615            .collect();
616        let b2 = Barrier::with_prev_epoch_for_test(test_epoch(1000), *prev_epoch)
617            .with_mutation(Mutation::Stop(StopMutation::default()));
618        barrier_test_env.inject_barrier(&b2, [actor_id]);
619        barrier_test_env.flush_all_events().await;
620
621        for (tx_id, tx) in txs.into_iter().enumerate() {
622            let epochs = epochs.clone();
623            let barriers = barriers.clone();
624            let b2 = b2.clone();
625            let handle = tokio::spawn(async move {
626                for (idx, epoch) in epochs {
627                    if idx % 20 == 0 {
628                        tx.send(Message::Chunk(build_test_chunk(10)).into())
629                            .await
630                            .unwrap();
631                    } else {
632                        tx.send(
633                            Message::Watermark(Watermark {
634                                col_idx: (idx as usize / 20 + tx_id) % CHANNEL_NUMBER,
635                                data_type: DataType::Int64,
636                                val: ScalarImpl::Int64(idx as i64),
637                            })
638                            .into(),
639                        )
640                        .await
641                        .unwrap();
642                    }
643                    tx.send(Message::Barrier(barriers[&epoch].clone().into_dispatcher()).into())
644                        .await
645                        .unwrap();
646                    sleep(Duration::from_millis(1)).await;
647                }
648                tx.send(Message::Barrier(b2.clone().into_dispatcher()).into())
649                    .await
650                    .unwrap();
651            });
652            handles.push(handle);
653        }
654
655        let merger = MergeExecutor::for_test(
656            actor_id,
657            rxs,
658            barrier_test_env.local_barrier_manager.clone(),
659            Schema::new(vec![]),
660            100,
661            None,
662        );
663        let mut merger = merger.boxed().execute();
664        for (idx, epoch) in epochs {
665            if idx % 20 == 0 {
666                // expect 1 or more chunks with 100 rows in total
667                let mut count = 0usize;
668                while count < 100 {
669                    assert_matches!(merger.next().await.unwrap().unwrap(), Message::Chunk(chunk) => {
670                        count += chunk.ops().len();
671                    });
672                }
673                assert_eq!(count, 100);
674            } else if idx as usize / 20 >= CHANNEL_NUMBER - 1 {
675                // expect n watermarks
676                for _ in 0..CHANNEL_NUMBER {
677                    assert_matches!(merger.next().await.unwrap().unwrap(), Message::Watermark(watermark) => {
678                        assert_eq!(watermark.val, ScalarImpl::Int64((idx - 20 * (CHANNEL_NUMBER as u64 - 1)) as i64));
679                    });
680                }
681            }
682            // expect a barrier
683            assert_matches!(merger.next().await.unwrap().unwrap(), Message::Barrier(Barrier{epoch:barrier_epoch,..}) => {
684                assert_eq!(barrier_epoch.curr, epoch);
685            });
686        }
687        assert_matches!(
688            merger.next().await.unwrap().unwrap(),
689            Message::Barrier(Barrier {
690                mutation,
691                ..
692            }) if mutation.as_deref().unwrap().is_stop_mutation()
693        );
694
695        for handle in handles {
696            handle.await.unwrap();
697        }
698    }
699
700    #[tokio::test]
701    async fn test_configuration_change() {
702        let actor_id = 233.into();
703        let (untouched, old, new) = (234.into(), 235.into(), 238.into()); // upstream actors
704        let barrier_test_env = LocalBarrierTestEnv::for_test().await;
705        let metrics = Arc::new(StreamingMetrics::unused());
706
707        // untouched -> actor_id
708        // old -> actor_id
709        // new -> actor_id
710
711        let (upstream_fragment_id, fragment_id) = (10.into(), 18.into());
712
713        let actor_ctx = ActorContext::for_test(actor_id);
714
715        let inputs: Vec<_> =
716            try_join_all([untouched, old].into_iter().map(async |upstream_actor_id| {
717                new_input(
718                    &barrier_test_env.local_barrier_manager,
719                    metrics.clone(),
720                    actor_id,
721                    fragment_id,
722                    &helper_make_local_actor(upstream_actor_id),
723                    upstream_fragment_id,
724                    actor_ctx.config.clone(),
725                )
726                .await
727            }))
728            .await
729            .unwrap();
730
731        let merge_updates = maplit::hashmap! {
732            (actor_id, upstream_fragment_id) => MergeUpdate {
733                actor_id,
734                upstream_fragment_id,
735                new_upstream_fragment_id: None,
736                added_upstream_actors: vec![helper_make_local_actor(new)],
737                removed_upstream_actor_id: vec![old],
738            }
739        };
740
741        let b1 = Barrier::new_test_barrier(test_epoch(1)).with_mutation(Mutation::Update(
742            UpdateMutation {
743                merges: merge_updates,
744                ..Default::default()
745            },
746        ));
747        barrier_test_env.inject_barrier(&b1, [actor_id]);
748        barrier_test_env.flush_all_events().await;
749
750        let barrier_rx = barrier_test_env
751            .local_barrier_manager
752            .subscribe_barrier(actor_id);
753        let upstream = MergeExecutor::new_merge_upstream(
754            inputs,
755            &metrics,
756            &actor_ctx,
757            100,
758            Schema::empty().clone(),
759        );
760
761        let mut merge = MergeExecutor::new(
762            actor_ctx,
763            fragment_id,
764            upstream_fragment_id,
765            upstream,
766            barrier_test_env.local_barrier_manager.clone(),
767            metrics.clone(),
768            barrier_rx,
769        )
770        .boxed()
771        .execute();
772
773        let mut txs = HashMap::new();
774        macro_rules! send {
775            ($actors:expr, $msg:expr) => {
776                for actor in $actors {
777                    txs.get(&actor).unwrap().send($msg).await.unwrap();
778                }
779            };
780        }
781
782        macro_rules! assert_recv_pending {
783            () => {
784                assert!(
785                    merge
786                        .next()
787                        .now_or_never()
788                        .flatten()
789                        .transpose()
790                        .unwrap()
791                        .is_none()
792                );
793            };
794        }
795        macro_rules! recv {
796            () => {
797                merge.next().await.transpose().unwrap()
798            };
799        }
800
801        macro_rules! collect_upstream_tx {
802            ($actors:expr) => {
803                for upstream_id in $actors {
804                    let mut output_requests = barrier_test_env
805                        .take_pending_new_output_requests(upstream_id)
806                        .await;
807                    assert_eq!(output_requests.len(), 1);
808                    let (downstream_actor_id, request) = output_requests.pop().unwrap();
809                    assert_eq!(downstream_actor_id, actor_id);
810                    let NewOutputRequest::Local(tx) = request else {
811                        unreachable!()
812                    };
813                    txs.insert(upstream_id, tx);
814                }
815            };
816        }
817
818        assert_recv_pending!();
819        barrier_test_env.flush_all_events().await;
820
821        // 2. Take downstream receivers.
822        collect_upstream_tx!([untouched, old]);
823
824        // 3. Send a chunk.
825        send!([untouched, old], Message::Chunk(build_test_chunk(1)).into());
826        assert_eq!(2, recv!().unwrap().as_chunk().unwrap().cardinality()); // We should be able to receive the chunk twice.
827        assert_recv_pending!();
828
829        send!(
830            [untouched, old],
831            Message::Barrier(b1.clone().into_dispatcher()).into()
832        );
833        assert_recv_pending!(); // We should not receive the barrier, since merger is waiting for the new upstream new.
834
835        collect_upstream_tx!([new]);
836
837        send!([new], Message::Barrier(b1.clone().into_dispatcher()).into());
838        recv!().unwrap().as_barrier().unwrap(); // We should now receive the barrier.
839
840        // 5. Send a chunk.
841        send!([untouched, new], Message::Chunk(build_test_chunk(1)).into());
842        assert_eq!(2, recv!().unwrap().as_chunk().unwrap().cardinality()); // We should be able to receive the chunk twice.
843        assert_recv_pending!();
844    }
845
846    struct FakeExchangeService {
847        rpc_called: Arc<AtomicBool>,
848    }
849
850    fn exchange_client_test_barrier() -> crate::executor::Barrier {
851        Barrier::new_test_barrier(test_epoch(1))
852    }
853
854    #[async_trait::async_trait]
855    impl StreamExchangeService for FakeExchangeService {
856        type GetStreamStream = ReceiverStream<std::result::Result<GetStreamResponse, Status>>;
857
858        async fn get_stream(
859            &self,
860            _request: Request<Streaming<GetStreamRequest>>,
861        ) -> std::result::Result<Response<Self::GetStreamStream>, Status> {
862            let (tx, rx) = tokio::sync::mpsc::channel(10);
863            self.rpc_called.store(true, Ordering::SeqCst);
864            // send stream_chunk
865            let stream_chunk = StreamChunk::default().to_protobuf();
866            tx.send(Ok(GetStreamResponse {
867                message: Some(PbStreamMessageBatch {
868                    stream_message_batch: Some(
869                        risingwave_pb::stream_plan::stream_message_batch::StreamMessageBatch::StreamChunk(
870                            stream_chunk,
871                        ),
872                    ),
873                }),
874                permits: Some(PbPermits::default()),
875            }))
876            .await
877            .unwrap();
878            // send barrier
879            let barrier = exchange_client_test_barrier();
880            tx.send(Ok(GetStreamResponse {
881                message: Some(PbStreamMessageBatch {
882                    stream_message_batch: Some(
883                        risingwave_pb::stream_plan::stream_message_batch::StreamMessageBatch::BarrierBatch(
884                            BarrierBatch {
885                                barriers: vec![barrier.to_protobuf()],
886                            },
887                        ),
888                    ),
889                }),
890                permits: Some(PbPermits::default()),
891            }))
892            .await
893            .unwrap();
894            Ok(Response::new(ReceiverStream::new(rx)))
895        }
896    }
897
898    #[tokio::test]
899    async fn test_stream_exchange_client() {
900        let rpc_called = Arc::new(AtomicBool::new(false));
901        let server_run = Arc::new(AtomicBool::new(false));
902        let addr = "127.0.0.1:12348".parse().unwrap();
903
904        // Start a server.
905        let (shutdown_send, shutdown_recv) = tokio::sync::oneshot::channel();
906        let exchange_svc = StreamExchangeServiceServer::new(FakeExchangeService {
907            rpc_called: rpc_called.clone(),
908        });
909        let cp_server_run = server_run.clone();
910        let join_handle = tokio::spawn(async move {
911            cp_server_run.store(true, Ordering::SeqCst);
912            tonic::transport::Server::builder()
913                .add_service(exchange_svc)
914                .serve_with_shutdown(addr, async move {
915                    shutdown_recv.await.unwrap();
916                })
917                .await
918                .unwrap();
919        });
920
921        sleep(Duration::from_secs(1)).await;
922        assert!(server_run.load(Ordering::SeqCst));
923
924        let test_env = LocalBarrierTestEnv::for_test().await;
925
926        let remote_input = {
927            RemoteInput::new(
928                &test_env.local_barrier_manager,
929                addr.into(),
930                TEST_PARTIAL_GRAPH_ID,
931                (0.into(), 0.into()),
932                (0.into(), 0.into()),
933                Arc::new(StreamingMetrics::unused()),
934                Arc::new(StreamingConfig::default()),
935            )
936            .await
937            .unwrap()
938        };
939
940        test_env.inject_barrier(&exchange_client_test_barrier(), [remote_input.id()]);
941
942        pin_mut!(remote_input);
943
944        assert_matches!(remote_input.next().await.unwrap().unwrap(), Message::Chunk(chunk) => {
945            let (ops, columns, visibility) = chunk.into_inner();
946            assert!(ops.is_empty());
947            assert!(columns.is_empty());
948            assert!(visibility.is_empty());
949        });
950        assert_matches!(remote_input.next().await.unwrap().unwrap(), Message::Barrier(Barrier { epoch: barrier_epoch, .. }) => {
951            assert_eq!(barrier_epoch.curr, test_epoch(1));
952        });
953        assert!(rpc_called.load(Ordering::SeqCst));
954
955        shutdown_send.send(()).unwrap();
956        join_handle.await.unwrap();
957    }
958}
risingwave_stream/executor/merge.rs

risingwave_stream/executor/
merge.rs
