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

risingwave_stream/executor/
merge.rs
