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        staging_table_id: TableId,
68    },
69    RegisterBarrierSender {
70        actor_id: ActorId,
71        barrier_sender: mpsc::UnboundedSender<Barrier>,
72    },
73    RegisterLocalUpstreamOutput {
74        actor_id: ActorId,
75        upstream_actor_id: ActorId,
76        upstream_partial_graph_id: PartialGraphId,
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            tx,
192        });
193        rx
194    }
195
196    pub fn report_source_list_finished(
197        &self,
198        epoch: EpochPair,
199        actor_id: ActorId,
200        table_id: TableId,
201        associated_source_id: SourceId,
202    ) {
203        self.send_event(LocalBarrierEvent::ReportSourceListFinished {
204            epoch,
205            actor_id,
206            table_id,
207            associated_source_id,
208        });
209    }
210
211    pub fn report_source_load_finished(
212        &self,
213        epoch: EpochPair,
214        actor_id: ActorId,
215        table_id: TableId,
216        associated_source_id: SourceId,
217    ) {
218        self.send_event(LocalBarrierEvent::ReportSourceLoadFinished {
219            epoch,
220            actor_id,
221            table_id,
222            associated_source_id,
223        });
224    }
225
226    pub fn report_refresh_finished(
227        &self,
228        epoch: EpochPair,
229        actor_id: ActorId,
230        table_id: TableId,
231        staging_table_id: TableId,
232    ) {
233        self.send_event(LocalBarrierEvent::RefreshFinished {
234            epoch,
235            actor_id,
236            table_id,
237            staging_table_id,
238        });
239    }
240
241    pub fn report_cdc_source_offset_updated(
242        &self,
243        epoch: EpochPair,
244        actor_id: ActorId,
245        source_id: SourceId,
246    ) {
247        self.send_event(LocalBarrierEvent::ReportCdcSourceOffsetUpdated {
248            epoch,
249            actor_id,
250            source_id,
251        });
252    }
253
254    pub fn report_iceberg_pk_index_sink_metadata(
255        &self,
256        epoch: EpochPair,
257        sink_id: risingwave_common::id::SinkId,
258        actor_id: ActorId,
259        role: PbIcebergPkIndexSinkRole,
260        metadata: Option<SinkMetadata>,
261    ) {
262        self.send_event(LocalBarrierEvent::ReportIcebergPkIndexSinkMetadata {
263            epoch,
264            sink_id,
265            actor_id,
266            role,
267            metadata,
268        });
269    }
270}
271
272#[cfg(test)]
273impl LocalBarrierManager {
274    pub(super) fn spawn_for_test()
275    -> crate::task::barrier_worker::EventSender<crate::task::barrier_worker::LocalActorOperation>
276    {
277        use std::sync::Arc;
278        use std::sync::atomic::AtomicU64;
279
280        use crate::executor::monitor::StreamingMetrics;
281        use crate::task::barrier_worker::{EventSender, LocalBarrierWorker};
282
283        let (tx, rx) = unbounded_channel();
284        let _join_handle = LocalBarrierWorker::spawn(
285            StreamEnvironment::for_test(),
286            Arc::new(StreamingMetrics::unused()),
287            None,
288            Arc::new(AtomicU64::new(0)),
289            rx,
290        );
291        EventSender(tx)
292    }
293}