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

risingwave_stream/executor/
merge.rs
