Skip to main content

risingwave_stream/task/barrier_manager/
mod.rs

1// Copyright 2025 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
15pub mod cdc_progress;
16pub mod progress;
17
18use std::sync::Arc;
19
20pub use progress::CreateMviewProgressReporter;
21use risingwave_common::id::{SourceId, TableId};
22use risingwave_common::util::epoch::EpochPair;
23use risingwave_pb::connector_service::SinkMetadata;
24use risingwave_pb::id::{FragmentId, PartialGraphId, SinkId};
25use risingwave_pb::stream_service::PbIcebergPkIndexSinkRole;
26use tokio::sync::mpsc;
27use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
28
29use crate::error::{IntoUnexpectedExit, StreamError};
30use crate::executor::exchange::permit;
31use crate::executor::monitor::StreamingMetrics;
32use crate::executor::{Barrier, BarrierInner};
33use crate::task::barrier_manager::progress::BackfillState;
34use crate::task::cdc_progress::CdcTableBackfillState;
35use crate::task::{ActorId, StreamEnvironment};
36
37/// Events sent from actors via [`LocalBarrierManager`] to [`super::barrier_worker::managed_state::PartialGraphState`].
38///
39/// See [`crate::task`] for architecture overview.
40pub(super) enum LocalBarrierEvent {
41    ReportActorCollected {
42        actor_id: ActorId,
43        epoch: EpochPair,
44    },
45    ReportCreateProgress {
46        epoch: EpochPair,
47        fragment_id: FragmentId,
48        actor: ActorId,
49        state: BackfillState,
50    },
51    ReportSourceListFinished {
52        epoch: EpochPair,
53        actor_id: ActorId,
54        table_id: TableId,
55        associated_source_id: SourceId,
56    },
57    ReportSourceLoadFinished {
58        epoch: EpochPair,
59        actor_id: ActorId,
60        table_id: TableId,
61        associated_source_id: SourceId,
62    },
63    RefreshFinished {
64        epoch: EpochPair,
65        actor_id: ActorId,
66        table_id: TableId,
67    },
68    RegisterBarrierSender {
69        actor_id: ActorId,
70        barrier_sender: mpsc::UnboundedSender<Barrier>,
71    },
72    RegisterLocalUpstreamOutput {
73        actor_id: ActorId,
74        upstream_actor_id: ActorId,
75        upstream_partial_graph_id: PartialGraphId,
76        term_id: String,
77        tx: permit::Sender,
78    },
79    ReportCdcTableBackfillProgress {
80        actor_id: ActorId,
81        epoch: EpochPair,
82        state: CdcTableBackfillState,
83    },
84    ReportCdcSourceOffsetUpdated {
85        epoch: EpochPair,
86        actor_id: ActorId,
87        source_id: SourceId,
88    },
89    ReportIcebergPkIndexSinkMetadata {
90        epoch: EpochPair,
91        sink_id: SinkId,
92        actor_id: ActorId,
93        role: PbIcebergPkIndexSinkRole,
94        metadata: Option<SinkMetadata>,
95    },
96}
97
98/// Can send [`LocalBarrierEvent`] to [`super::barrier_worker::managed_state::PartialGraphState::poll_next_event`]
99///
100/// See [`crate::task`] for architecture overview.
101#[derive(Clone)]
102pub struct LocalBarrierManager {
103    barrier_event_sender: UnboundedSender<LocalBarrierEvent>,
104    actor_failure_sender: UnboundedSender<(ActorId, StreamError)>,
105    pub(crate) term_id: String,
106    pub(crate) env: StreamEnvironment,
107}
108
109impl LocalBarrierManager {
110    pub(super) fn new(
111        term_id: String,
112        env: StreamEnvironment,
113    ) -> (
114        Self,
115        UnboundedReceiver<LocalBarrierEvent>,
116        UnboundedReceiver<(ActorId, StreamError)>,
117    ) {
118        let (event_tx, event_rx) = unbounded_channel();
119        let (err_tx, err_rx) = unbounded_channel();
120        (
121            Self {
122                barrier_event_sender: event_tx,
123                actor_failure_sender: err_tx,
124                term_id,
125                env,
126            },
127            event_rx,
128            err_rx,
129        )
130    }
131
132    pub fn for_test() -> Self {
133        Self::new("114514".to_owned(), StreamEnvironment::for_test()).0
134    }
135
136    /// Event is handled by [`super::barrier_worker::managed_state::PartialGraphState::poll_next_event`]
137    fn send_event(&self, event: LocalBarrierEvent) {
138        // ignore error, because the current barrier manager maybe a stale one
139        let _ = self.barrier_event_sender.send(event);
140    }
141
142    /// When a [`crate::executor::StreamConsumer`] (typically [`crate::executor::DispatchExecutor`]) get a barrier, it should report
143    /// and collect this barrier with its own `actor_id` using this function.
144    pub fn collect<M>(&self, actor_id: ActorId, barrier: &BarrierInner<M>) {
145        self.send_event(LocalBarrierEvent::ReportActorCollected {
146            actor_id,
147            epoch: barrier.epoch,
148        })
149    }
150
151    /// When a actor exit unexpectedly, it should report this event using this function, so meta
152    /// will notice actor's exit while collecting.
153    pub fn notify_failure(&self, actor_id: ActorId, err: StreamError) {
154        let _ = self
155            .actor_failure_sender
156            .send((actor_id, err.into_unexpected_exit(actor_id)));
157    }
158
159    pub fn subscribe_barrier(&self, actor_id: ActorId) -> UnboundedReceiver<Barrier> {
160        let (tx, rx) = mpsc::unbounded_channel();
161        self.send_event(LocalBarrierEvent::RegisterBarrierSender {
162            actor_id,
163            barrier_sender: tx,
164        });
165        rx
166    }
167
168    pub fn register_local_upstream_output(
169        &self,
170        actor_id: ActorId,
171        upstream_actor_id: ActorId,
172        upstream_fragment_id: FragmentId,
173        upstream_partial_graph_id: PartialGraphId,
174        metrics: Arc<StreamingMetrics>,
175    ) -> permit::Receiver {
176        let upstream_fragment_id_str = upstream_fragment_id.to_string();
177        let fragment_channel_buffered_bytes = metrics
178            .fragment_channel_buffered_bytes
179            .with_guarded_label_values(&[&upstream_fragment_id_str]);
180        let (tx, rx) = permit::channel_from_config_with_metrics(
181            self.env.global_config(),
182            permit::ChannelMetrics {
183                sender_actor_channel_buffered_bytes: fragment_channel_buffered_bytes.clone(),
184                receiver_actor_channel_buffered_bytes: fragment_channel_buffered_bytes,
185            },
186        );
187        self.send_event(LocalBarrierEvent::RegisterLocalUpstreamOutput {
188            actor_id,
189            upstream_actor_id,
190            upstream_partial_graph_id,
191            term_id: self.term_id.clone(),
192            tx,
193        });
194        rx
195    }
196
197    pub fn report_source_list_finished(
198        &self,
199        epoch: EpochPair,
200        actor_id: ActorId,
201        table_id: TableId,
202        associated_source_id: SourceId,
203    ) {
204        self.send_event(LocalBarrierEvent::ReportSourceListFinished {
205            epoch,
206            actor_id,
207            table_id,
208            associated_source_id,
209        });
210    }
211
212    pub fn report_source_load_finished(
213        &self,
214        epoch: EpochPair,
215        actor_id: ActorId,
216        table_id: TableId,
217        associated_source_id: SourceId,
218    ) {
219        self.send_event(LocalBarrierEvent::ReportSourceLoadFinished {
220            epoch,
221            actor_id,
222            table_id,
223            associated_source_id,
224        });
225    }
226
227    pub fn report_refresh_finished(&self, epoch: EpochPair, actor_id: ActorId, table_id: TableId) {
228        self.send_event(LocalBarrierEvent::RefreshFinished {
229            epoch,
230            actor_id,
231            table_id,
232        });
233    }
234
235    pub fn report_cdc_source_offset_updated(
236        &self,
237        epoch: EpochPair,
238        actor_id: ActorId,
239        source_id: SourceId,
240    ) {
241        self.send_event(LocalBarrierEvent::ReportCdcSourceOffsetUpdated {
242            epoch,
243            actor_id,
244            source_id,
245        });
246    }
247
248    pub fn report_iceberg_pk_index_sink_metadata(
249        &self,
250        epoch: EpochPair,
251        sink_id: risingwave_common::id::SinkId,
252        actor_id: ActorId,
253        role: PbIcebergPkIndexSinkRole,
254        metadata: Option<SinkMetadata>,
255    ) {
256        self.send_event(LocalBarrierEvent::ReportIcebergPkIndexSinkMetadata {
257            epoch,
258            sink_id,
259            actor_id,
260            role,
261            metadata,
262        });
263    }
264}
265
266#[cfg(test)]
267impl LocalBarrierManager {
268    pub(super) fn spawn_for_test()
269    -> crate::task::barrier_worker::EventSender<crate::task::barrier_worker::LocalActorOperation>
270    {
271        use std::sync::Arc;
272        use std::sync::atomic::AtomicU64;
273
274        use crate::executor::monitor::StreamingMetrics;
275        use crate::task::barrier_worker::{EventSender, LocalBarrierWorker};
276
277        let (tx, rx) = unbounded_channel();
278        let _join_handle = LocalBarrierWorker::spawn(
279            StreamEnvironment::for_test(),
280            Arc::new(StreamingMetrics::unused()),
281            None,
282            Arc::new(AtomicU64::new(0)),
283            rx,
284        );
285        EventSender(tx)
286    }
287}