risingwave_stream/task/barrier_worker/
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
15use std::collections::hash_map::Entry;
16use std::collections::{HashMap, HashSet};
17use std::fmt::Display;
18use std::future::{pending, poll_fn};
19use std::sync::Arc;
20use std::task::Poll;
21
22use anyhow::anyhow;
23use await_tree::{InstrumentAwait, SpanExt};
24use futures::future::{BoxFuture, join, join_all};
25use futures::stream::{BoxStream, FuturesOrdered};
26use futures::{FutureExt, StreamExt, TryFutureExt};
27use itertools::Itertools;
28use risingwave_pb::stream_plan::barrier::BarrierKind;
29use risingwave_pb::stream_service::barrier_complete_response::{
30    PbCdcSourceOffsetUpdated, PbCdcTableBackfillProgress, PbCreateMviewProgress,
31    PbIcebergPkIndexSinkMetadata, PbListFinishedSource, PbLoadFinishedSource, PbLocalSstableInfo,
32    PbRefreshFinishedActor,
33};
34use risingwave_rpc_client::error::{ToTonicStatus, TonicStatusWrapper};
35use risingwave_storage::store_impl::AsHummock;
36use thiserror_ext::AsReport;
37use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
38use tokio::sync::oneshot;
39use tokio::task::JoinHandle;
40use tokio::{select, spawn};
41use tonic::{Code, Status};
42use tracing::warn;
43
44use self::managed_state::ManagedBarrierState;
45use crate::error::{ScoredStreamError, StreamError, StreamResult};
46#[cfg(test)]
47use crate::task::LocalBarrierManager;
48use crate::task::managed_state::{BarrierToComplete, ResetPartialGraphOutput};
49use crate::task::{
50    ActorId, AtomicU64Ref, CONFIG_OVERRIDE_CACHE_DEFAULT_CAPACITY, ConfigOverrideCache, FragmentId,
51    PartialGraphId, StreamActorManager, StreamEnvironment, UpDownActorIds,
52};
53pub mod managed_state;
54#[cfg(test)]
55mod tests;
56
57use risingwave_common::metrics::LabelGuardedHistogram;
58use risingwave_hummock_sdk::table_stats::to_prost_table_stats_map;
59use risingwave_hummock_sdk::{LocalSstableInfo, SyncResult};
60use risingwave_pb::stream_service::streaming_control_stream_request::{
61    InitRequest, Request, ResetPartialGraphsRequest,
62};
63use risingwave_pb::stream_service::streaming_control_stream_response::{
64    InitResponse, ReportPartialGraphFailureResponse, ResetPartialGraphResponse, Response,
65    ShutdownResponse,
66};
67use risingwave_pb::stream_service::{
68    BarrierCompleteResponse, InjectBarrierRequest, PbScoredError, StreamingControlStreamRequest,
69    StreamingControlStreamResponse, streaming_control_stream_response,
70};
71
72use crate::executor::Barrier;
73use crate::executor::exchange::permit::Receiver;
74use crate::executor::monitor::StreamingMetrics;
75use crate::task::barrier_worker::managed_state::{
76    ManagedBarrierStateDebugInfo, ManagedBarrierStateEvent, PartialGraphState, PartialGraphStatus,
77};
78
79/// If enabled, all actors will be grouped in the same tracing span within one epoch.
80/// Note that this option will significantly increase the overhead of tracing.
81pub const ENABLE_BARRIER_AGGREGATION: bool = false;
82
83/// Collect result of some barrier on current compute node. Will be reported to the meta service in [`LocalBarrierWorker::on_epoch_completed`].
84#[derive(Debug)]
85pub struct BarrierCompleteResult {
86    /// The result returned from `sync` of `StateStore`.
87    pub sync_result: Option<SyncResult>,
88
89    /// The updated creation progress of materialized view after this barrier.
90    pub create_mview_progress: Vec<PbCreateMviewProgress>,
91
92    /// The source IDs that have finished listing data for refreshable batch sources.
93    pub list_finished_source_ids: Vec<PbListFinishedSource>,
94
95    /// The source IDs that have finished loading data for refreshable batch sources.
96    pub load_finished_source_ids: Vec<PbLoadFinishedSource>,
97
98    pub cdc_table_backfill_progress: Vec<PbCdcTableBackfillProgress>,
99
100    /// CDC sources that have updated their offset at least once.
101    pub cdc_source_offset_updated: Vec<PbCdcSourceOffsetUpdated>,
102
103    /// Iceberg pk-index sink metadata reports collected during this barrier.
104    pub iceberg_pk_index_sink_metadata: Vec<PbIcebergPkIndexSinkMetadata>,
105
106    /// The materialize actors that have finished their part of a table refresh.
107    pub refresh_finished_actors: Vec<PbRefreshFinishedActor>,
108}
109
110/// Lives in [`crate::task::barrier_worker::LocalBarrierWorker`],
111/// Communicates with `ControlStreamManager` in meta.
112/// Handles [`risingwave_pb::stream_service::streaming_control_stream_request::Request`].
113pub(super) struct ControlStreamHandle {
114    #[expect(clippy::type_complexity)]
115    pair: Option<(
116        UnboundedSender<Result<StreamingControlStreamResponse, Status>>,
117        BoxStream<'static, Result<StreamingControlStreamRequest, Status>>,
118    )>,
119}
120
121impl ControlStreamHandle {
122    fn empty() -> Self {
123        Self { pair: None }
124    }
125
126    pub(super) fn new(
127        sender: UnboundedSender<Result<StreamingControlStreamResponse, Status>>,
128        request_stream: BoxStream<'static, Result<StreamingControlStreamRequest, Status>>,
129    ) -> Self {
130        Self {
131            pair: Some((sender, request_stream)),
132        }
133    }
134
135    pub(super) fn connected(&self) -> bool {
136        self.pair.is_some()
137    }
138
139    fn reset_stream_with_err(&mut self, err: Status) {
140        if let Some((sender, _)) = self.pair.take() {
141            // Note: `TonicStatusWrapper` provides a better error report.
142            let err = TonicStatusWrapper::new(err);
143            warn!(error = %err.as_report(), "control stream reset with error");
144
145            let err = err.into_inner();
146            if sender.send(Err(err)).is_err() {
147                warn!("failed to notify reset of control stream");
148            }
149        }
150    }
151
152    /// Send `Shutdown` message to the control stream and wait for the stream to be closed
153    /// by the meta service.
154    async fn shutdown_stream(&mut self) {
155        if let Some((sender, _)) = self.pair.take() {
156            if sender
157                .send(Ok(StreamingControlStreamResponse {
158                    response: Some(streaming_control_stream_response::Response::Shutdown(
159                        ShutdownResponse::default(),
160                    )),
161                }))
162                .is_err()
163            {
164                warn!("failed to notify shutdown of control stream");
165            } else {
166                tracing::info!("waiting for meta service to close control stream...");
167
168                // Wait for the stream to be closed, to ensure that the `Shutdown` message has
169                // been acknowledged by the meta service for more precise error report.
170                //
171                // This is because the meta service will reset the control stream manager and
172                // drop the connection to us upon recovery. As a result, the receiver part of
173                // this sender will also be dropped, causing the stream to close.
174                sender.closed().await;
175            }
176        } else {
177            debug!("control stream has been reset, ignore shutdown");
178        }
179    }
180
181    pub(super) fn ack_reset_partial_graph(
182        &mut self,
183        partial_graph_id: PartialGraphId,
184        root_err: Option<ScoredStreamError>,
185    ) {
186        self.send_response(Response::ResetPartialGraph(ResetPartialGraphResponse {
187            partial_graph_id,
188            root_err: root_err.map(|err| PbScoredError {
189                err_msg: err.error.to_report_string(),
190                score: err.score.0,
191            }),
192        }));
193    }
194
195    fn send_response(&mut self, response: streaming_control_stream_response::Response) {
196        if let Some((sender, _)) = self.pair.as_ref() {
197            if sender
198                .send(Ok(StreamingControlStreamResponse {
199                    response: Some(response),
200                }))
201                .is_err()
202            {
203                self.pair = None;
204                warn!("failed to send the response; the control stream was reset");
205            }
206        } else {
207            debug!(?response, "control stream has been reset. ignore response");
208        }
209    }
210
211    async fn next_request(&mut self) -> StreamingControlStreamRequest {
212        if let Some((_, stream)) = &mut self.pair {
213            match stream.next().await {
214                Some(Ok(request)) => {
215                    return request;
216                }
217                Some(Err(e)) => self.reset_stream_with_err(
218                    anyhow!(TonicStatusWrapper::new(e)) // wrap the status to provide better error report
219                        .context("failed to get request")
220                        .to_status_unnamed(Code::Internal),
221                ),
222                None => self.reset_stream_with_err(Status::internal("end of stream")),
223            }
224        }
225        pending().await
226    }
227}
228
229pub(super) enum TakeReceiverRequest {
230    Remote {
231        result_sender: oneshot::Sender<StreamResult<Receiver>>,
232        upstream_fragment_id: FragmentId,
233    },
234    Local(permit::Sender),
235}
236
237/// Sent from [`crate::task::stream_manager::LocalStreamManager`] to [`crate::task::barrier_worker::LocalBarrierWorker::run`].
238///
239/// See [`crate::task`] for architecture overview.
240#[derive(strum_macros::Display)]
241pub(super) enum LocalActorOperation {
242    NewControlStream {
243        handle: ControlStreamHandle,
244        init_request: InitRequest,
245    },
246    TakeReceiver {
247        partial_graph_id: PartialGraphId,
248        term_id: String,
249        ids: UpDownActorIds,
250        request: TakeReceiverRequest,
251    },
252    #[cfg(test)]
253    GetCurrentLocalBarrierManager(oneshot::Sender<LocalBarrierManager>),
254    #[cfg(test)]
255    TakePendingNewOutputRequest(ActorId, oneshot::Sender<Vec<(ActorId, NewOutputRequest)>>),
256    #[cfg(test)]
257    Flush(oneshot::Sender<()>),
258    InspectState {
259        result_sender: oneshot::Sender<String>,
260    },
261    Shutdown {
262        result_sender: oneshot::Sender<()>,
263    },
264}
265
266pub(super) struct LocalBarrierWorkerDebugInfo<'a> {
267    managed_barrier_state:
268        HashMap<PartialGraphId, (String, Option<ManagedBarrierStateDebugInfo<'a>>)>,
269    has_control_stream_connected: bool,
270}
271
272impl Display for LocalBarrierWorkerDebugInfo<'_> {
273    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
274        writeln!(
275            f,
276            "\nhas_control_stream_connected: {}",
277            self.has_control_stream_connected
278        )?;
279
280        for (partial_graph_id, (status, managed_barrier_state)) in &self.managed_barrier_state {
281            writeln!(
282                f,
283                "partial graph {} status: {} managed_barrier_state:\n{}",
284                partial_graph_id,
285                status,
286                managed_barrier_state
287                    .as_ref()
288                    .map(ToString::to_string)
289                    .unwrap_or_default()
290            )?;
291        }
292        Ok(())
293    }
294}
295
296/// [`LocalBarrierWorker`] manages barrier control flow.
297/// Specifically, [`LocalBarrierWorker`] serves barrier injection from meta server, sends the
298/// barriers to and collects them from all actors, and finally reports the progress.
299///
300/// Runs event loop in [`Self::run`]. Handles events sent by [`crate::task::LocalStreamManager`].
301///
302/// See [`crate::task`] for architecture overview.
303pub(super) struct LocalBarrierWorker {
304    /// Current barrier collection state.
305    pub(super) state: ManagedBarrierState,
306
307    /// Futures will be finished in the order of epoch in ascending order.
308    await_epoch_completed_futures:
309        HashMap<PartialGraphId, FuturesOrdered<AwaitEpochCompletedFuture>>,
310
311    control_stream_handle: ControlStreamHandle,
312
313    pub(super) actor_manager: Arc<StreamActorManager>,
314}
315
316impl LocalBarrierWorker {
317    pub(super) fn new(actor_manager: Arc<StreamActorManager>) -> Self {
318        Self {
319            state: Default::default(),
320            await_epoch_completed_futures: Default::default(),
321            control_stream_handle: ControlStreamHandle::empty(),
322            actor_manager,
323        }
324    }
325
326    fn to_debug_info(&self) -> LocalBarrierWorkerDebugInfo<'_> {
327        LocalBarrierWorkerDebugInfo {
328            managed_barrier_state: self
329                .state
330                .partial_graphs
331                .iter()
332                .map(|(partial_graph_id, status)| {
333                    (*partial_graph_id, {
334                        match status {
335                            PartialGraphStatus::ReceivedExchangeRequest(_) => {
336                                ("ReceivedExchangeRequest".to_owned(), None)
337                            }
338                            PartialGraphStatus::Running(state) => {
339                                ("running".to_owned(), Some(state.to_debug_info()))
340                            }
341                            PartialGraphStatus::Suspended(state) => {
342                                (format!("suspended: {:?}", state.suspend_time), None)
343                            }
344                            PartialGraphStatus::Resetting => ("resetting".to_owned(), None),
345                            PartialGraphStatus::Unspecified => {
346                                unreachable!()
347                            }
348                        }
349                    })
350                })
351                .collect(),
352            has_control_stream_connected: self.control_stream_handle.connected(),
353        }
354    }
355
356    async fn next_completed_epoch(
357        futures: &mut HashMap<PartialGraphId, FuturesOrdered<AwaitEpochCompletedFuture>>,
358    ) -> (PartialGraphId, Barrier, StreamResult<BarrierCompleteResult>) {
359        poll_fn(|cx| {
360            for (partial_graph_id, futures) in &mut *futures {
361                if let Poll::Ready(Some((barrier, result))) = futures.poll_next_unpin(cx) {
362                    return Poll::Ready((*partial_graph_id, barrier, result));
363                }
364            }
365            Poll::Pending
366        })
367        .await
368    }
369
370    async fn run(mut self, mut actor_op_rx: UnboundedReceiver<LocalActorOperation>) {
371        loop {
372            select! {
373                biased;
374                event = self.state.next_event() => {
375                    match event {
376                        ManagedBarrierStateEvent::BarrierCollected{
377                            partial_graph_id,
378                            barrier,
379                        } => {
380                            // update await_epoch_completed_futures
381                            // handled below in next_completed_epoch
382                            self.complete_barrier(partial_graph_id, barrier.epoch.prev);
383                        }
384                        ManagedBarrierStateEvent::ActorError{
385                            partial_graph_id,
386                            actor_id,
387                            err,
388                        } => {
389                            self.on_partial_graph_failure(partial_graph_id, Some(actor_id), err, "recv actor failure");
390                        }
391                        ManagedBarrierStateEvent::PartialGraphsReset(output) => {
392                            for (partial_graph_id, output) in output {
393                                self.ack_partial_graph_reset(partial_graph_id, Some(output));
394                            }
395                        }
396                        ManagedBarrierStateEvent::RegisterLocalUpstreamOutput{
397                            actor_id,
398                            upstream_actor_id,
399                            upstream_partial_graph_id,
400                            term_id,
401                            tx
402                        } => {
403                            self.handle_actor_op(LocalActorOperation::TakeReceiver {
404                                partial_graph_id: upstream_partial_graph_id,
405                                term_id,
406                                ids: (upstream_actor_id, actor_id),
407                                request: TakeReceiverRequest::Local(tx),
408                            });
409                        }
410                    }
411                }
412                (partial_graph_id, barrier, result) = Self::next_completed_epoch(&mut self.await_epoch_completed_futures) => {
413                    match result {
414                        Ok(result) => {
415                            self.on_epoch_completed(partial_graph_id, barrier.epoch.prev, result);
416                        }
417                        Err(err) => {
418                            // TODO: may only report as partial graph failure instead of reset the stream
419                            // when the HummockUploader support partial recovery. Currently the HummockUploader
420                            // enter `Err` state and stop working until a global recovery to clear the uploader.
421                            self.control_stream_handle.reset_stream_with_err(Status::internal(format!("failed to complete epoch: {} {:?} {:?}", partial_graph_id, barrier.epoch, err.as_report())));
422                        }
423                    }
424                },
425                actor_op = actor_op_rx.recv() => {
426                    if let Some(actor_op) = actor_op {
427                        match actor_op {
428                            LocalActorOperation::NewControlStream { handle, init_request  } => {
429                                self.control_stream_handle.reset_stream_with_err(Status::internal("control stream has been reset to a new one"));
430                                self.reset(init_request).await;
431                                self.control_stream_handle = handle;
432                                self.control_stream_handle.send_response(streaming_control_stream_response::Response::Init(InitResponse {}));
433                            }
434                            LocalActorOperation::Shutdown { result_sender } => {
435                                if self.state.partial_graphs.values().any(|graph| {
436                                    match graph {
437                                        PartialGraphStatus::Running(graph) => {
438                                            !graph.actor_states.is_empty()
439                                        }
440                                        PartialGraphStatus::Suspended(_) | PartialGraphStatus::Resetting |
441                                            PartialGraphStatus::ReceivedExchangeRequest(_) => {
442                                            false
443                                        }
444                                        PartialGraphStatus::Unspecified => {
445                                            unreachable!()
446                                        }
447                                    }
448                                }) {
449                                    tracing::warn!(
450                                        "shutdown with running actors, scaling or migration will be triggered"
451                                    );
452                                }
453                                self.control_stream_handle.shutdown_stream().await;
454                                let _ = result_sender.send(());
455                            }
456                            actor_op => {
457                                self.handle_actor_op(actor_op);
458                            }
459                        }
460                    }
461                    else {
462                        break;
463                    }
464                },
465                request = self.control_stream_handle.next_request() => {
466                    let result = self.handle_streaming_control_request(request.request.expect("non empty"));
467                    if let Err((partial_graph_id, err)) = result {
468                        self.on_partial_graph_failure(partial_graph_id, None, err, "failed to inject barrier");
469                    }
470                },
471            }
472        }
473    }
474
475    fn handle_streaming_control_request(
476        &mut self,
477        request: Request,
478    ) -> Result<(), (PartialGraphId, StreamError)> {
479        match request {
480            Request::InjectBarrier(req) => {
481                let partial_graph_id = req.partial_graph_id;
482                let result: StreamResult<()> = try {
483                    let barrier = Barrier::from_protobuf(req.get_barrier().unwrap())
484                        .map_err(StreamError::from)?;
485                    self.send_barrier(&barrier, req)?;
486                };
487                result.map_err(|e| (partial_graph_id, e))?;
488                Ok(())
489            }
490            Request::RemovePartialGraph(req) => {
491                self.remove_partial_graphs(req.partial_graph_ids);
492                Ok(())
493            }
494            Request::CreatePartialGraph(req) => {
495                self.add_partial_graph(req.partial_graph_id, req.term_id);
496                Ok(())
497            }
498            Request::ResetPartialGraphs(req) => {
499                self.reset_partial_graphs(req);
500                Ok(())
501            }
502            Request::Init(_) => {
503                unreachable!()
504            }
505        }
506    }
507
508    fn handle_actor_op(&mut self, actor_op: LocalActorOperation) {
509        match actor_op {
510            LocalActorOperation::NewControlStream { .. } | LocalActorOperation::Shutdown { .. } => {
511                unreachable!("event {actor_op} should be handled separately in async context")
512            }
513            LocalActorOperation::TakeReceiver {
514                partial_graph_id,
515                term_id,
516                ids,
517                request,
518            } => {
519                let err = match self.state.partial_graphs.entry(partial_graph_id) {
520                    Entry::Occupied(mut entry) => match entry.get_mut() {
521                        PartialGraphStatus::ReceivedExchangeRequest(pending_requests) => {
522                            pending_requests.push((term_id, ids, request));
523                            return;
524                        }
525                        PartialGraphStatus::Running(graph) => {
526                            if graph.local_barrier_manager.term_id != term_id {
527                                warn!(
528                                    %partial_graph_id,
529                                    ?ids,
530                                    request_term_id = term_id,
531                                    current_term_id = graph.local_barrier_manager.term_id,
532                                    "take receiver on unmatched partial graph term"
533                                );
534                                anyhow!(
535                                    "take receiver {:?} on unmatched partial graph term {} to current term {}",
536                                    ids,
537                                    term_id,
538                                    graph.local_barrier_manager.term_id
539                                )
540                            } else {
541                                let (upstream_actor_id, actor_id) = ids;
542                                graph.new_actor_output_request(
543                                    actor_id,
544                                    upstream_actor_id,
545                                    request,
546                                );
547                                return;
548                            }
549                        }
550                        PartialGraphStatus::Suspended(_) => anyhow!("partial graph suspended"),
551                        PartialGraphStatus::Resetting => anyhow!("partial graph resetting"),
552                        PartialGraphStatus::Unspecified => unreachable!(),
553                    },
554                    Entry::Vacant(entry) => {
555                        entry.insert(PartialGraphStatus::ReceivedExchangeRequest(vec![(
556                            term_id, ids, request,
557                        )]));
558                        return;
559                    }
560                };
561                if let TakeReceiverRequest::Remote { result_sender, .. } = request {
562                    let _ = result_sender.send(Err(err.into()));
563                }
564            }
565            #[cfg(test)]
566            LocalActorOperation::GetCurrentLocalBarrierManager(sender) => {
567                let partial_graph_status = self
568                    .state
569                    .partial_graphs
570                    .get(&crate::task::TEST_PARTIAL_GRAPH_ID)
571                    .unwrap();
572                let partial_graph_state = risingwave_common::must_match!(partial_graph_status, PartialGraphStatus::Running(database_state) => database_state);
573                let _ = sender.send(partial_graph_state.local_barrier_manager.clone());
574            }
575            #[cfg(test)]
576            LocalActorOperation::TakePendingNewOutputRequest(actor_id, sender) => {
577                let partial_graph_status = self
578                    .state
579                    .partial_graphs
580                    .get_mut(&crate::task::TEST_PARTIAL_GRAPH_ID)
581                    .unwrap();
582
583                let partial_graph_state = risingwave_common::must_match!(partial_graph_status, PartialGraphStatus::Running(database_state) => database_state);
584                assert!(!partial_graph_state.actor_states.contains_key(&actor_id));
585                let requests = partial_graph_state
586                    .actor_pending_new_output_requests
587                    .remove(&actor_id)
588                    .unwrap();
589                let _ = sender.send(requests);
590            }
591            #[cfg(test)]
592            LocalActorOperation::Flush(sender) => {
593                use futures::FutureExt;
594                while let Some(request) = self.control_stream_handle.next_request().now_or_never() {
595                    self.handle_streaming_control_request(
596                        request.request.expect("should not be empty"),
597                    )
598                    .unwrap();
599                }
600                while let Some(event) = self.state.next_event().now_or_never() {
601                    match event {
602                        ManagedBarrierStateEvent::BarrierCollected {
603                            barrier,
604                            partial_graph_id,
605                        } => {
606                            self.complete_barrier(partial_graph_id, barrier.epoch.prev);
607                        }
608                        ManagedBarrierStateEvent::ActorError { .. }
609                        | ManagedBarrierStateEvent::PartialGraphsReset { .. }
610                        | ManagedBarrierStateEvent::RegisterLocalUpstreamOutput { .. } => {
611                            unreachable!()
612                        }
613                    }
614                }
615                sender.send(()).unwrap()
616            }
617            LocalActorOperation::InspectState { result_sender } => {
618                let debug_info = self.to_debug_info();
619                let _ = result_sender.send(debug_info.to_string());
620            }
621        }
622    }
623}
624
625mod await_epoch_completed_future {
626    use std::future::Future;
627
628    use futures::FutureExt;
629    use futures::future::BoxFuture;
630    use risingwave_hummock_sdk::SyncResult;
631    use risingwave_pb::stream_service::barrier_complete_response::{
632        PbCdcSourceOffsetUpdated, PbCdcTableBackfillProgress, PbCreateMviewProgress,
633        PbIcebergPkIndexSinkMetadata, PbListFinishedSource, PbLoadFinishedSource,
634        PbRefreshFinishedActor,
635    };
636
637    use crate::error::StreamResult;
638    use crate::executor::Barrier;
639    use crate::task::{BarrierCompleteResult, await_tree_key};
640
641    pub(super) type AwaitEpochCompletedFuture =
642        impl Future<Output = (Barrier, StreamResult<BarrierCompleteResult>)> + 'static;
643
644    #[define_opaque(AwaitEpochCompletedFuture)]
645    #[expect(clippy::too_many_arguments)]
646    pub(super) fn instrument_complete_barrier_future(
647        complete_barrier_future: Option<BoxFuture<'static, StreamResult<SyncResult>>>,
648        barrier: Barrier,
649        barrier_await_tree_reg: Option<&await_tree::Registry>,
650        create_mview_progress: Vec<PbCreateMviewProgress>,
651        list_finished_source_ids: Vec<PbListFinishedSource>,
652        load_finished_source_ids: Vec<PbLoadFinishedSource>,
653        cdc_table_backfill_progress: Vec<PbCdcTableBackfillProgress>,
654        cdc_source_offset_updated: Vec<PbCdcSourceOffsetUpdated>,
655        iceberg_pk_index_sink_metadata: Vec<PbIcebergPkIndexSinkMetadata>,
656        refresh_finished_actors: Vec<PbRefreshFinishedActor>,
657    ) -> AwaitEpochCompletedFuture {
658        let prev_epoch = barrier.epoch.prev;
659        let future = async move {
660            if let Some(future) = complete_barrier_future {
661                let result = future.await;
662                result.map(Some)
663            } else {
664                Ok(None)
665            }
666        }
667        .map(move |result| {
668            (
669                barrier,
670                result.map(|sync_result| BarrierCompleteResult {
671                    sync_result,
672                    create_mview_progress,
673                    list_finished_source_ids,
674                    load_finished_source_ids,
675                    cdc_table_backfill_progress,
676                    cdc_source_offset_updated,
677                    iceberg_pk_index_sink_metadata,
678                    refresh_finished_actors,
679                }),
680            )
681        });
682        if let Some(reg) = barrier_await_tree_reg {
683            reg.register(
684                await_tree_key::BarrierAwait { prev_epoch },
685                format!("SyncEpoch({})", prev_epoch),
686            )
687            .instrument(future)
688            .left_future()
689        } else {
690            future.right_future()
691        }
692    }
693}
694
695use await_epoch_completed_future::*;
696use risingwave_common::catalog::TableId;
697use risingwave_pb::hummock::vector_index_delta::PbVectorIndexAdds;
698use risingwave_storage::{StateStoreImpl, dispatch_state_store};
699
700use crate::executor::exchange::permit;
701
702fn sync_epoch(
703    state_store: &StateStoreImpl,
704    barrier_sync_latency: LabelGuardedHistogram,
705    prev_epoch: u64,
706    table_ids: HashSet<TableId>,
707) -> BoxFuture<'static, StreamResult<SyncResult>> {
708    let timer = barrier_sync_latency.start_timer();
709
710    let state_store = state_store.clone();
711    let future = async move {
712        dispatch_state_store!(state_store, hummock, {
713            hummock.sync(vec![(prev_epoch, table_ids)]).await
714        })
715    };
716
717    future
718        .instrument_await(await_tree::span!("sync_epoch (epoch {})", prev_epoch))
719        .inspect_ok(move |_| {
720            let _guard = &barrier_sync_latency;
721            timer.observe_duration();
722        })
723        .map_err(move |e| {
724            tracing::error!(
725                prev_epoch,
726                error = %e.as_report(),
727                "Failed to sync state store",
728            );
729            e.into()
730        })
731        .boxed()
732}
733
734impl LocalBarrierWorker {
735    fn complete_barrier(&mut self, partial_graph_id: PartialGraphId, prev_epoch: u64) {
736        {
737            let Some(graph_state) = self
738                .state
739                .partial_graphs
740                .get_mut(&partial_graph_id)
741                .expect("should exist")
742                .state_for_request()
743            else {
744                return;
745            };
746            let BarrierToComplete {
747                barrier,
748                table_ids,
749                create_mview_progress,
750                list_finished_source_ids,
751                load_finished_source_ids,
752                cdc_table_backfill_progress,
753                cdc_source_offset_updated,
754                iceberg_pk_index_sink_metadata,
755                refresh_finished_actors,
756            } = graph_state.pop_barrier_to_complete(prev_epoch);
757
758            let complete_barrier_future = match &barrier.kind {
759                BarrierKind::Unspecified => unreachable!(),
760                BarrierKind::Initial => {
761                    tracing::info!(
762                        epoch = prev_epoch,
763                        "ignore sealing data for the first barrier"
764                    );
765                    tracing::info!(?prev_epoch, "ignored syncing data for the first barrier");
766                    None
767                }
768                BarrierKind::Barrier => None,
769                BarrierKind::Checkpoint => Some(sync_epoch(
770                    &self.actor_manager.env.state_store(),
771                    graph_state.graph_state.barrier_sync_latency(),
772                    prev_epoch,
773                    table_ids.expect("should be Some on BarrierKind::Checkpoint"),
774                )),
775            };
776
777            self.await_epoch_completed_futures
778                .entry(partial_graph_id)
779                .or_default()
780                .push_back({
781                    instrument_complete_barrier_future(
782                        complete_barrier_future,
783                        barrier,
784                        self.actor_manager.await_tree_reg.as_ref(),
785                        create_mview_progress,
786                        list_finished_source_ids,
787                        load_finished_source_ids,
788                        cdc_table_backfill_progress,
789                        cdc_source_offset_updated,
790                        iceberg_pk_index_sink_metadata,
791                        refresh_finished_actors,
792                    )
793                });
794        }
795    }
796
797    fn on_epoch_completed(
798        &mut self,
799        partial_graph_id: PartialGraphId,
800        epoch: u64,
801        result: BarrierCompleteResult,
802    ) {
803        let BarrierCompleteResult {
804            create_mview_progress,
805            sync_result,
806            list_finished_source_ids,
807            load_finished_source_ids,
808            cdc_table_backfill_progress,
809            cdc_source_offset_updated,
810            iceberg_pk_index_sink_metadata,
811            refresh_finished_actors,
812        } = result;
813
814        let (synced_sstables, table_watermarks, old_value_ssts, vector_index_adds) = sync_result
815            .map(|sync_result| {
816                (
817                    sync_result.uncommitted_ssts,
818                    sync_result.table_watermarks,
819                    sync_result.old_value_ssts,
820                    sync_result.vector_index_adds,
821                )
822            })
823            .unwrap_or_default();
824
825        let result = {
826            {
827                streaming_control_stream_response::Response::CompleteBarrier(
828                    BarrierCompleteResponse {
829                        request_id: "todo".to_owned(),
830                        partial_graph_id,
831                        epoch,
832                        status: None,
833                        create_mview_progress,
834                        synced_sstables: synced_sstables
835                            .into_iter()
836                            .map(
837                                |LocalSstableInfo {
838                                     sst_info,
839                                     table_stats,
840                                     created_at,
841                                 }| PbLocalSstableInfo {
842                                    sst: Some(sst_info.into()),
843                                    table_stats_map: to_prost_table_stats_map(table_stats),
844                                    created_at,
845                                },
846                            )
847                            .collect_vec(),
848                        worker_id: self.actor_manager.env.worker_id(),
849                        table_watermarks: table_watermarks
850                            .into_iter()
851                            .map(|(key, value)| (key, value.into()))
852                            .collect(),
853                        old_value_sstables: old_value_ssts
854                            .into_iter()
855                            .map(|sst| sst.sst_info.into())
856                            .collect(),
857                        list_finished_sources: list_finished_source_ids,
858                        load_finished_sources: load_finished_source_ids,
859                        cdc_source_offset_updated,
860                        vector_index_adds: vector_index_adds
861                            .into_iter()
862                            .map(|(table_id, adds)| {
863                                (
864                                    table_id,
865                                    PbVectorIndexAdds {
866                                        adds: adds.into_iter().map(|add| add.into()).collect(),
867                                    },
868                                )
869                            })
870                            .collect(),
871                        cdc_table_backfill_progress,
872                        refresh_finished_actors,
873                        iceberg_pk_index_sink_metadata,
874                    },
875                )
876            }
877        };
878
879        self.control_stream_handle.send_response(result);
880    }
881
882    /// Broadcast a barrier to all senders. Save a receiver which will get notified when this
883    /// barrier is finished, in managed mode.
884    ///
885    /// Note that the error returned here is typically a [`StreamError::barrier_send`], which is not
886    /// the root cause of the failure. The caller should then call `try_find_root_failure`
887    /// to find the root cause.
888    fn send_barrier(
889        &mut self,
890        barrier: &Barrier,
891        request: InjectBarrierRequest,
892    ) -> StreamResult<()> {
893        debug!(
894            target: "events::stream::barrier::manager::send",
895            "send barrier {:?}, actor_ids_to_collect = {:?}",
896            barrier,
897            request.actor_ids_to_collect
898        );
899
900        let status = self
901            .state
902            .partial_graphs
903            .get_mut(&request.partial_graph_id)
904            .expect("should exist");
905        if let Some(state) = status.state_for_request() {
906            state.transform_to_issued(barrier, request)?;
907        }
908        Ok(())
909    }
910
911    fn remove_partial_graphs(
912        &mut self,
913        partial_graph_ids: impl IntoIterator<Item = PartialGraphId>,
914    ) {
915        for partial_graph_id in partial_graph_ids {
916            if let Some(mut graph) = self.state.partial_graphs.remove(&partial_graph_id) {
917                if let Some(graph) = graph.state_for_request() {
918                    assert!(
919                        graph.graph_state.is_empty(),
920                        "non empty graph to be removed: {}",
921                        &graph.graph_state
922                    );
923                }
924            } else {
925                warn!(
926                    partial_graph_id = %partial_graph_id,
927                    "no partial graph to remove"
928                );
929            }
930        }
931    }
932
933    fn add_partial_graph(&mut self, partial_graph_id: PartialGraphId, term_id: String) {
934        match self.state.partial_graphs.entry(partial_graph_id) {
935            Entry::Occupied(entry) => {
936                let status = entry.into_mut();
937                if let PartialGraphStatus::ReceivedExchangeRequest(pending_requests) = status {
938                    let mut graph = PartialGraphState::new(
939                        partial_graph_id,
940                        term_id.clone(),
941                        self.actor_manager.clone(),
942                    );
943                    for (request_term_id, (upstream_actor_id, actor_id), request) in
944                        pending_requests.drain(..)
945                    {
946                        if request_term_id == term_id {
947                            graph.new_actor_output_request(actor_id, upstream_actor_id, request);
948                        } else {
949                            warn!(
950                                %partial_graph_id,
951                                %upstream_actor_id,
952                                %actor_id,
953                                request_term_id,
954                                current_term_id = term_id,
955                                "reject buffered exchange request with stale partial graph term"
956                            );
957                            if let TakeReceiverRequest::Remote { result_sender, .. } = request {
958                                let _ = result_sender.send(Err(anyhow!(
959                                    "take receiver ({upstream_actor_id}, {actor_id}) on stale partial graph term {request_term_id}; current term is {term_id}"
960                                )
961                                .into()));
962                            }
963                        }
964                    }
965                    *status = PartialGraphStatus::Running(graph);
966                } else {
967                    panic!("duplicated partial graph: {}", partial_graph_id);
968                }
969
970                status
971            }
972            Entry::Vacant(entry) => entry.insert(PartialGraphStatus::Running(
973                PartialGraphState::new(partial_graph_id, term_id, self.actor_manager.clone()),
974            )),
975        };
976    }
977
978    fn reset_partial_graphs(&mut self, req: ResetPartialGraphsRequest) {
979        let mut table_ids_to_clear = HashSet::new();
980        let mut reset_futures = HashMap::new();
981        for partial_graph_id in req.partial_graph_ids {
982            if let Some(status) = self.state.partial_graphs.get_mut(&partial_graph_id) {
983                let reset_future = status.start_reset(
984                    partial_graph_id,
985                    self.await_epoch_completed_futures.remove(&partial_graph_id),
986                    &mut table_ids_to_clear,
987                );
988                reset_futures.insert(partial_graph_id, reset_future);
989            } else {
990                self.ack_partial_graph_reset(partial_graph_id, None);
991            }
992        }
993        if reset_futures.is_empty() {
994            assert!(table_ids_to_clear.is_empty());
995            return;
996        }
997        let state_store = self.actor_manager.env.state_store();
998        self.state.resetting_graphs.push(spawn(async move {
999            let outputs =
1000                join_all(
1001                    reset_futures
1002                        .into_iter()
1003                        .map(|(partial_graph_id, future)| async move {
1004                            (partial_graph_id, future.await)
1005                        }),
1006                )
1007                .await;
1008            if !table_ids_to_clear.is_empty()
1009                && let Some(hummock) = state_store.as_hummock()
1010            {
1011                hummock.clear_tables(table_ids_to_clear).await;
1012            }
1013            outputs
1014        }));
1015    }
1016
1017    fn ack_partial_graph_reset(
1018        &mut self,
1019        partial_graph_id: PartialGraphId,
1020        reset_output: Option<ResetPartialGraphOutput>,
1021    ) {
1022        info!(
1023            %partial_graph_id,
1024            "partial graph reset successfully"
1025        );
1026        assert!(!self.state.partial_graphs.contains_key(&partial_graph_id));
1027        self.await_epoch_completed_futures.remove(&partial_graph_id);
1028        self.control_stream_handle.ack_reset_partial_graph(
1029            partial_graph_id,
1030            reset_output.and_then(|output| output.root_err),
1031        );
1032    }
1033
1034    /// When some other failure happens (like failed to send barrier), the error is reported using
1035    /// this function. The control stream will be responded with a message to notify about the error,
1036    /// and the global barrier worker will later reset and rerun the partial graph.
1037    fn on_partial_graph_failure(
1038        &mut self,
1039        partial_graph_id: PartialGraphId,
1040        failed_actor: Option<ActorId>,
1041        err: StreamError,
1042        message: impl Into<String>,
1043    ) {
1044        let message = message.into();
1045        error!(%partial_graph_id, ?failed_actor, message, err = ?err.as_report(), "suspend partial graph on error");
1046        let completing_futures = self.await_epoch_completed_futures.remove(&partial_graph_id);
1047        self.state
1048            .partial_graphs
1049            .get_mut(&partial_graph_id)
1050            .expect("should exist")
1051            .suspend(failed_actor, err, completing_futures);
1052        self.control_stream_handle
1053            .send_response(Response::ReportPartialGraphFailure(
1054                ReportPartialGraphFailureResponse { partial_graph_id },
1055            ));
1056    }
1057
1058    /// Force stop all actors on this worker, and then drop their resources.
1059    async fn reset(&mut self, _init_request: InitRequest) {
1060        join(
1061            join_all(
1062                self.state
1063                    .partial_graphs
1064                    .values_mut()
1065                    .map(|graph| graph.abort()),
1066            ),
1067            async {
1068                while let Some(join_result) = self.state.resetting_graphs.next().await {
1069                    join_result.expect("failed to join reset partial graphs handle");
1070                }
1071            },
1072        )
1073        .await;
1074        if let Some(m) = self.actor_manager.await_tree_reg.as_ref() {
1075            m.clear();
1076        }
1077
1078        if let Some(hummock) = self.actor_manager.env.state_store().as_hummock() {
1079            hummock
1080                .clear_shared_buffer()
1081                .instrument_await("store_clear_shared_buffer".verbose())
1082                .await
1083        }
1084        self.actor_manager.env.dml_manager_ref().clear();
1085        *self = Self::new(self.actor_manager.clone());
1086        self.actor_manager.env.client_pool().invalidate_all();
1087    }
1088
1089    /// Create a [`LocalBarrierWorker`] with managed mode.
1090    pub fn spawn(
1091        env: StreamEnvironment,
1092        streaming_metrics: Arc<StreamingMetrics>,
1093        await_tree_reg: Option<await_tree::Registry>,
1094        watermark_epoch: AtomicU64Ref,
1095        actor_op_rx: UnboundedReceiver<LocalActorOperation>,
1096    ) -> JoinHandle<()> {
1097        let runtime = {
1098            let mut builder = tokio::runtime::Builder::new_multi_thread();
1099            if let Some(worker_threads_num) = env.global_config().actor_runtime_worker_threads_num {
1100                builder.worker_threads(worker_threads_num);
1101            }
1102            builder
1103                .thread_name("rw-streaming")
1104                .enable_all()
1105                .build()
1106                .unwrap()
1107        };
1108
1109        let actor_manager = Arc::new(StreamActorManager {
1110            env,
1111            streaming_metrics,
1112            watermark_epoch,
1113            await_tree_reg,
1114            runtime: runtime.into(),
1115            config_override_cache: ConfigOverrideCache::new(CONFIG_OVERRIDE_CACHE_DEFAULT_CAPACITY),
1116        });
1117        let worker = LocalBarrierWorker::new(actor_manager);
1118        tokio::spawn(worker.run(actor_op_rx))
1119    }
1120}
1121
1122pub(super) struct EventSender<T>(pub(super) UnboundedSender<T>);
1123
1124impl<T> Clone for EventSender<T> {
1125    fn clone(&self) -> Self {
1126        Self(self.0.clone())
1127    }
1128}
1129
1130impl<T> EventSender<T> {
1131    pub(super) fn send_event(&self, event: T) {
1132        self.0.send(event).expect("should be able to send event")
1133    }
1134
1135    pub(super) async fn send_and_await<RSP>(
1136        &self,
1137        make_event: impl FnOnce(oneshot::Sender<RSP>) -> T,
1138    ) -> StreamResult<RSP> {
1139        let (tx, rx) = oneshot::channel();
1140        let event = make_event(tx);
1141        self.send_event(event);
1142        rx.await
1143            .map_err(|_| anyhow!("barrier manager maybe reset").into())
1144    }
1145}
1146
1147pub(crate) enum NewOutputRequest {
1148    Local(permit::Sender),
1149    Remote(permit::Sender),
1150}
1151
1152#[cfg(test)]
1153pub(crate) mod barrier_test_utils {
1154    use assert_matches::assert_matches;
1155    use futures::StreamExt;
1156    use risingwave_pb::stream_service::streaming_control_stream_request::{
1157        InitRequest, PbCreatePartialGraphRequest,
1158    };
1159    use risingwave_pb::stream_service::{
1160        InjectBarrierRequest, PbStreamingControlStreamRequest, StreamingControlStreamRequest,
1161        StreamingControlStreamResponse, streaming_control_stream_request,
1162        streaming_control_stream_response,
1163    };
1164    use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
1165    use tokio::sync::oneshot;
1166    use tokio_stream::wrappers::UnboundedReceiverStream;
1167    use tonic::Status;
1168
1169    use crate::executor::Barrier;
1170    use crate::task::barrier_worker::{ControlStreamHandle, EventSender, LocalActorOperation};
1171    use crate::task::{ActorId, LocalBarrierManager, NewOutputRequest, TEST_PARTIAL_GRAPH_ID};
1172
1173    pub(crate) struct LocalBarrierTestEnv {
1174        pub local_barrier_manager: LocalBarrierManager,
1175        pub(super) actor_op_tx: EventSender<LocalActorOperation>,
1176        pub request_tx: UnboundedSender<Result<StreamingControlStreamRequest, Status>>,
1177        pub response_rx: UnboundedReceiver<Result<StreamingControlStreamResponse, Status>>,
1178    }
1179
1180    impl LocalBarrierTestEnv {
1181        pub(crate) async fn for_test() -> Self {
1182            let actor_op_tx = LocalBarrierManager::spawn_for_test();
1183
1184            let (request_tx, request_rx) = unbounded_channel();
1185            let (response_tx, mut response_rx) = unbounded_channel();
1186
1187            request_tx
1188                .send(Ok(PbStreamingControlStreamRequest {
1189                    request: Some(
1190                        streaming_control_stream_request::Request::CreatePartialGraph(
1191                            PbCreatePartialGraphRequest {
1192                                partial_graph_id: TEST_PARTIAL_GRAPH_ID,
1193                                term_id: "for_test".into(),
1194                            },
1195                        ),
1196                    ),
1197                }))
1198                .unwrap();
1199
1200            actor_op_tx.send_event(LocalActorOperation::NewControlStream {
1201                handle: ControlStreamHandle::new(
1202                    response_tx,
1203                    UnboundedReceiverStream::new(request_rx).boxed(),
1204                ),
1205                init_request: InitRequest {},
1206            });
1207
1208            assert_matches!(
1209                response_rx.recv().await.unwrap().unwrap().response.unwrap(),
1210                streaming_control_stream_response::Response::Init(_)
1211            );
1212
1213            let local_barrier_manager = actor_op_tx
1214                .send_and_await(LocalActorOperation::GetCurrentLocalBarrierManager)
1215                .await
1216                .unwrap();
1217
1218            Self {
1219                local_barrier_manager,
1220                actor_op_tx,
1221                request_tx,
1222                response_rx,
1223            }
1224        }
1225
1226        pub(crate) fn inject_barrier(
1227            &self,
1228            barrier: &Barrier,
1229            actor_to_collect: impl IntoIterator<Item = ActorId>,
1230        ) {
1231            self.request_tx
1232                .send(Ok(StreamingControlStreamRequest {
1233                    request: Some(streaming_control_stream_request::Request::InjectBarrier(
1234                        InjectBarrierRequest {
1235                            request_id: "".to_owned(),
1236                            barrier: Some(barrier.to_protobuf()),
1237                            actor_ids_to_collect: actor_to_collect.into_iter().collect(),
1238                            table_ids_to_sync: vec![],
1239                            partial_graph_id: TEST_PARTIAL_GRAPH_ID,
1240                            actors_to_build: vec![],
1241                        },
1242                    )),
1243                }))
1244                .unwrap();
1245        }
1246
1247        pub(crate) async fn flush_all_events(&self) {
1248            Self::flush_all_events_impl(&self.actor_op_tx).await
1249        }
1250
1251        pub(super) async fn flush_all_events_impl(actor_op_tx: &EventSender<LocalActorOperation>) {
1252            let (tx, rx) = oneshot::channel();
1253            actor_op_tx.send_event(LocalActorOperation::Flush(tx));
1254            rx.await.unwrap()
1255        }
1256
1257        pub(crate) async fn take_pending_new_output_requests(
1258            &self,
1259            actor_id: ActorId,
1260        ) -> Vec<(ActorId, NewOutputRequest)> {
1261            self.actor_op_tx
1262                .send_and_await(|tx| LocalActorOperation::TakePendingNewOutputRequest(actor_id, tx))
1263                .await
1264                .unwrap()
1265        }
1266    }
1267}
risingwave_stream/task/barrier_worker/mod.rs

risingwave_stream/task/barrier_worker/
mod.rs
