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