risingwave_meta/manager/sink_coordination/
coordinator_worker.rs

1// Copyright 2023 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::{BTreeMap, HashMap, HashSet, VecDeque};
16use std::fmt::Debug;
17use std::future::{Future, poll_fn};
18use std::pin::pin;
19use std::task::Poll;
20use std::time::{Duration, Instant};
21
22use anyhow::anyhow;
23use await_tree::InstrumentAwait;
24use futures::future::{Either, pending, select};
25use futures::pin_mut;
26use itertools::Itertools;
27use risingwave_common::bail;
28use risingwave_common::bitmap::Bitmap;
29use risingwave_connector::connector_common::IcebergSinkCompactionUpdate;
30use risingwave_connector::dispatch_sink;
31use risingwave_connector::sink::catalog::SinkId;
32use risingwave_connector::sink::{
33    Sink, SinkCommitCoordinator, SinkCommittedEpochSubscriber, SinkError, SinkParam, build_sink,
34};
35use risingwave_meta_model::pending_sink_state::SinkState;
36use risingwave_pb::connector_service::{SinkMetadata, coordinate_request};
37use risingwave_pb::stream_plan::PbSinkSchemaChange;
38use sea_orm::DatabaseConnection;
39use thiserror_ext::AsReport;
40use tokio::select;
41use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
42use tokio::time::sleep;
43use tokio_retry::strategy::{ExponentialBackoff, jitter};
44use tonic::Status;
45use tracing::{error, warn};
46
47use crate::manager::exactly_once_util::{
48    clean_aborted_records, commit_and_prune_epoch, list_sink_states_ordered_by_epoch,
49    persist_pre_commit_metadata,
50};
51use crate::manager::sink_coordination::handle::SinkWriterCoordinationHandle;
52
53async fn run_future_with_periodic_fn<F: Future>(
54    future: F,
55    interval: Duration,
56    mut f: impl FnMut(),
57) -> F::Output {
58    pin_mut!(future);
59    loop {
60        match select(&mut future, pin!(sleep(interval))).await {
61            Either::Left((output, _)) => {
62                break output;
63            }
64            Either::Right(_) => f(),
65        }
66    }
67}
68
69type HandleId = usize;
70
71#[derive(Default)]
72struct AligningRequests<R> {
73    requests: Vec<R>,
74    handle_ids: HashSet<HandleId>,
75    committed_bitmap: Option<Bitmap>, // lazy-initialized on first request
76}
77
78impl<R> AligningRequests<R> {
79    fn add_new_request(
80        &mut self,
81        handle_id: HandleId,
82        request: R,
83        vnode_bitmap: &Bitmap,
84    ) -> anyhow::Result<()>
85    where
86        R: Debug,
87    {
88        let committed_bitmap = self
89            .committed_bitmap
90            .get_or_insert_with(|| Bitmap::zeros(vnode_bitmap.len()));
91        assert_eq!(committed_bitmap.len(), vnode_bitmap.len());
92
93        let check_bitmap = (&*committed_bitmap) & vnode_bitmap;
94        if check_bitmap.count_ones() > 0 {
95            return Err(anyhow!(
96                "duplicate vnode {:?}. request vnode: {:?}, prev vnode: {:?}. pending request: {:?}, request: {:?}",
97                check_bitmap.iter_ones().collect_vec(),
98                vnode_bitmap,
99                committed_bitmap,
100                self.requests,
101                request
102            ));
103        }
104        *committed_bitmap |= vnode_bitmap;
105        self.requests.push(request);
106        assert!(self.handle_ids.insert(handle_id));
107        Ok(())
108    }
109
110    fn aligned(&self) -> bool {
111        self.committed_bitmap.as_ref().is_some_and(|b| b.all())
112    }
113}
114
115type RetryBackoffFuture = std::pin::Pin<Box<tokio::time::Sleep>>;
116type RetryBackoffStrategy = impl Iterator<Item = RetryBackoffFuture> + Send + 'static;
117
118struct TwoPhaseCommitHandler {
119    db: DatabaseConnection,
120    sink_id: SinkId,
121    curr_hummock_committed_epoch: u64,
122    job_committed_epoch_rx: UnboundedReceiver<u64>,
123    last_committed_epoch: Option<u64>,
124    pending_epochs: VecDeque<(u64, Option<Vec<u8>>, Option<PbSinkSchemaChange>)>,
125    prepared_epochs: VecDeque<(u64, Option<Vec<u8>>, Option<PbSinkSchemaChange>)>,
126    backoff_state: Option<(RetryBackoffFuture, RetryBackoffStrategy)>,
127}
128
129impl TwoPhaseCommitHandler {
130    fn new(
131        db: DatabaseConnection,
132        sink_id: SinkId,
133        initial_hummock_committed_epoch: u64,
134        job_committed_epoch_rx: UnboundedReceiver<u64>,
135        last_committed_epoch: Option<u64>,
136    ) -> Self {
137        Self {
138            db,
139            sink_id,
140            curr_hummock_committed_epoch: initial_hummock_committed_epoch,
141            job_committed_epoch_rx,
142            last_committed_epoch,
143            pending_epochs: VecDeque::new(),
144            prepared_epochs: VecDeque::new(),
145            backoff_state: None,
146        }
147    }
148
149    #[define_opaque(RetryBackoffStrategy)]
150    fn get_retry_backoff_strategy() -> RetryBackoffStrategy {
151        ExponentialBackoff::from_millis(10)
152            .max_delay(Duration::from_secs(60))
153            .map(jitter)
154            .map(|delay| Box::pin(tokio::time::sleep(delay)))
155    }
156
157    async fn next_to_commit(
158        &mut self,
159    ) -> anyhow::Result<(u64, Option<Vec<u8>>, Option<PbSinkSchemaChange>)> {
160        loop {
161            let wait_backoff = async {
162                if self.prepared_epochs.is_empty() {
163                    pending::<()>().await;
164                } else if let Some((backoff_fut, _)) = &mut self.backoff_state {
165                    backoff_fut.await;
166                }
167            };
168
169            select! {
170                _ = wait_backoff => {
171                    let item = self.prepared_epochs.front().cloned().expect("non-empty");
172                    return Ok(item);
173                }
174
175                recv_epoch = self.job_committed_epoch_rx.recv() => {
176                    let Some(recv_epoch) = recv_epoch else {
177                        return Err(anyhow!(
178                            "Hummock committed epoch sender closed unexpectedly"
179                        ));
180                    };
181                    self.curr_hummock_committed_epoch = recv_epoch;
182                    while let Some((epoch, metadata, schema_change)) = self.pending_epochs.pop_front_if(|(epoch, _, _)| *epoch <= recv_epoch) {
183                        if let Some((last_epoch, _, _)) = self.prepared_epochs.back() {
184                            assert!(epoch > *last_epoch, "prepared epochs must be in increasing order");
185                        }
186                        self.prepared_epochs.push_back((epoch, metadata, schema_change));
187                    }
188                }
189            }
190        }
191    }
192
193    fn push_new_item(
194        &mut self,
195        epoch: u64,
196        metadata: Option<Vec<u8>>,
197        schema_change: Option<PbSinkSchemaChange>,
198    ) {
199        if epoch > self.curr_hummock_committed_epoch {
200            if let Some((last_epoch, _, _)) = self.pending_epochs.back() {
201                assert!(
202                    epoch > *last_epoch,
203                    "pending epochs must be in increasing order"
204                );
205            }
206            self.pending_epochs
207                .push_back((epoch, metadata, schema_change));
208        } else {
209            assert!(self.pending_epochs.is_empty());
210            if let Some((last_epoch, _, _)) = self.prepared_epochs.back() {
211                assert!(
212                    epoch > *last_epoch,
213                    "prepared epochs must be in increasing order"
214                );
215            }
216            self.prepared_epochs
217                .push_back((epoch, metadata, schema_change));
218        }
219    }
220
221    async fn ack_committed(&mut self, epoch: u64) -> anyhow::Result<()> {
222        self.backoff_state = None;
223        let (last_epoch, _, _) = self.prepared_epochs.pop_front().expect("non-empty");
224        assert_eq!(last_epoch, epoch);
225
226        commit_and_prune_epoch(&self.db, self.sink_id, epoch, self.last_committed_epoch).await?;
227        self.last_committed_epoch = Some(epoch);
228        Ok(())
229    }
230
231    fn failed_committed(&mut self, epoch: u64, err: SinkError) {
232        assert_eq!(self.prepared_epochs.front().expect("non-empty").0, epoch,);
233        if let Some((prev_fut, strategy)) = &mut self.backoff_state {
234            let new_fut = strategy.next().expect("infinite");
235            *prev_fut = new_fut;
236        } else {
237            let mut strategy = Self::get_retry_backoff_strategy();
238            let backoff_fut = strategy.next().expect("infinite");
239            self.backoff_state = Some((backoff_fut, strategy));
240        }
241        tracing::error!(
242            error = %err.as_report(),
243            %self.sink_id,
244            "failed to commit epoch {}, Retrying after backoff",
245            epoch,
246        );
247    }
248
249    fn is_empty(&self) -> bool {
250        self.pending_epochs.is_empty() && self.prepared_epochs.is_empty()
251    }
252
253    /// Whether there exists an uncommitted schema change.
254    ///
255    /// Per current design, if a `schema_change` exists, it should be attached to the latest
256    /// uncommitted item across `pending_epochs` and `prepared_epochs`.
257    fn has_uncommitted_schema_change(&self) -> bool {
258        if let Some((_, _, schema_change)) = self.pending_epochs.back() {
259            schema_change.is_some()
260        } else if let Some((_, _, schema_change)) = self.prepared_epochs.back() {
261            schema_change.is_some()
262        } else {
263            false
264        }
265    }
266}
267
268struct CoordinationHandleManager {
269    param: SinkParam,
270    writer_handles: HashMap<HandleId, SinkWriterCoordinationHandle>,
271    next_handle_id: HandleId,
272    request_rx: UnboundedReceiver<SinkWriterCoordinationHandle>,
273}
274
275impl CoordinationHandleManager {
276    fn start(
277        &mut self,
278        log_store_rewind_start_epoch: Option<u64>,
279        handle_ids: impl IntoIterator<Item = HandleId>,
280    ) -> anyhow::Result<()> {
281        for handle_id in handle_ids {
282            let handle = self
283                .writer_handles
284                .get_mut(&handle_id)
285                .ok_or_else(|| anyhow!("fail to find handle for {} to start", handle_id,))?;
286            handle.start(log_store_rewind_start_epoch).map_err(|_| {
287                anyhow!(
288                    "fail to start {:?} for handle {}",
289                    log_store_rewind_start_epoch,
290                    handle_id
291                )
292            })?;
293        }
294        Ok(())
295    }
296
297    fn ack_aligned_initial_epoch(&mut self, aligned_initial_epoch: u64) -> anyhow::Result<()> {
298        for (handle_id, handle) in &mut self.writer_handles {
299            handle
300                .ack_aligned_initial_epoch(aligned_initial_epoch)
301                .map_err(|_| {
302                    anyhow!(
303                        "fail to ack_aligned_initial_epoch {:?} for handle {}",
304                        aligned_initial_epoch,
305                        handle_id
306                    )
307                })?;
308        }
309        Ok(())
310    }
311
312    fn ack_commit(
313        &mut self,
314        epoch: u64,
315        handle_ids: impl IntoIterator<Item = HandleId>,
316    ) -> anyhow::Result<()> {
317        for handle_id in handle_ids {
318            let handle = self.writer_handles.get_mut(&handle_id).ok_or_else(|| {
319                anyhow!(
320                    "fail to find handle for {} when ack commit on epoch {}",
321                    handle_id,
322                    epoch
323                )
324            })?;
325            handle.ack_commit(epoch).map_err(|_| {
326                anyhow!(
327                    "fail to ack commit on epoch {} for handle {}",
328                    epoch,
329                    handle_id
330                )
331            })?;
332        }
333        Ok(())
334    }
335
336    async fn next_request_inner(
337        writer_handles: &mut HashMap<HandleId, SinkWriterCoordinationHandle>,
338    ) -> anyhow::Result<(HandleId, coordinate_request::Msg)> {
339        poll_fn(|cx| {
340            for (handle_id, handle) in writer_handles.iter_mut() {
341                if let Poll::Ready(result) = handle.poll_next_request(cx) {
342                    return Poll::Ready(result.map(|request| (*handle_id, request)));
343                }
344            }
345            Poll::Pending
346        })
347        .await
348    }
349}
350
351enum CoordinationHandleManagerEvent {
352    NewHandle,
353    UpdateVnodeBitmap,
354    Stop,
355    CommitRequest {
356        epoch: u64,
357        metadata: SinkMetadata,
358        schema_change: Option<PbSinkSchemaChange>,
359    },
360    AlignInitialEpoch(u64),
361}
362
363impl CoordinationHandleManagerEvent {
364    fn name(&self) -> &'static str {
365        match self {
366            CoordinationHandleManagerEvent::NewHandle => "NewHandle",
367            CoordinationHandleManagerEvent::UpdateVnodeBitmap => "UpdateVnodeBitmap",
368            CoordinationHandleManagerEvent::Stop => "Stop",
369            CoordinationHandleManagerEvent::CommitRequest { .. } => "CommitRequest",
370            CoordinationHandleManagerEvent::AlignInitialEpoch(_) => "AlignInitialEpoch",
371        }
372    }
373}
374
375impl CoordinationHandleManager {
376    async fn next_event(&mut self) -> anyhow::Result<(HandleId, CoordinationHandleManagerEvent)> {
377        select! {
378            handle = self.request_rx.recv() => {
379                let handle = handle.ok_or_else(|| anyhow!("end of writer request stream"))?;
380                if handle.param() != &self.param {
381                    warn!(prev_param = ?self.param, new_param = ?handle.param(), "sink param mismatch");
382                }
383                let handle_id = self.next_handle_id;
384                self.next_handle_id += 1;
385                self.writer_handles.insert(handle_id, handle);
386                Ok((handle_id, CoordinationHandleManagerEvent::NewHandle))
387            }
388            result = Self::next_request_inner(&mut self.writer_handles) => {
389                let (handle_id, request) = result?;
390                let event = match request {
391                    coordinate_request::Msg::CommitRequest(request) => {
392                        CoordinationHandleManagerEvent::CommitRequest {
393                            epoch: request.epoch,
394                            metadata: request.metadata.ok_or_else(|| anyhow!("empty sink metadata"))?,
395                            schema_change: request.schema_change,
396                        }
397                    }
398                    coordinate_request::Msg::AlignInitialEpochRequest(epoch) => {
399                        CoordinationHandleManagerEvent::AlignInitialEpoch(epoch)
400                    }
401                    coordinate_request::Msg::UpdateVnodeRequest(_) => {
402                        CoordinationHandleManagerEvent::UpdateVnodeBitmap
403                    }
404                    coordinate_request::Msg::Stop(_) => {
405                        CoordinationHandleManagerEvent::Stop
406                    }
407                    coordinate_request::Msg::StartRequest(_) => {
408                        unreachable!("should have been handled");
409                    }
410                };
411                Ok((handle_id, event))
412            }
413        }
414    }
415
416    fn vnode_bitmap(&self, handle_id: HandleId) -> &Bitmap {
417        self.writer_handles[&handle_id].vnode_bitmap()
418    }
419
420    fn stop_handle(&mut self, handle_id: HandleId) -> anyhow::Result<()> {
421        self.writer_handles
422            .remove(&handle_id)
423            .expect("should exist")
424            .stop()
425    }
426
427    async fn wait_init_handles(&mut self) -> anyhow::Result<HashSet<HandleId>> {
428        assert!(self.writer_handles.is_empty());
429        let mut init_requests = AligningRequests::default();
430        while !init_requests.aligned() {
431            let (handle_id, event) = self.next_event().await?;
432            let unexpected_event = match event {
433                CoordinationHandleManagerEvent::NewHandle => {
434                    init_requests.add_new_request(handle_id, (), self.vnode_bitmap(handle_id))?;
435                    continue;
436                }
437                event => event.name(),
438            };
439            return Err(anyhow!(
440                "expect new handle during init, but got {}",
441                unexpected_event
442            ));
443        }
444        Ok(init_requests.handle_ids)
445    }
446
447    async fn alter_parallelisms(
448        &mut self,
449        altered_handles: impl Iterator<Item = HandleId>,
450    ) -> anyhow::Result<HashSet<HandleId>> {
451        let mut requests = AligningRequests::default();
452        for handle_id in altered_handles {
453            requests.add_new_request(handle_id, (), self.vnode_bitmap(handle_id))?;
454        }
455        let mut remaining_handles: HashSet<_> = self
456            .writer_handles
457            .keys()
458            .filter(|handle_id| !requests.handle_ids.contains(handle_id))
459            .cloned()
460            .collect();
461        while !remaining_handles.is_empty() || !requests.aligned() {
462            let (handle_id, event) = self.next_event().await?;
463            match event {
464                CoordinationHandleManagerEvent::NewHandle => {
465                    requests.add_new_request(handle_id, (), self.vnode_bitmap(handle_id))?;
466                }
467                CoordinationHandleManagerEvent::UpdateVnodeBitmap => {
468                    assert!(remaining_handles.remove(&handle_id));
469                    requests.add_new_request(handle_id, (), self.vnode_bitmap(handle_id))?;
470                }
471                CoordinationHandleManagerEvent::Stop => {
472                    assert!(remaining_handles.remove(&handle_id));
473                    self.stop_handle(handle_id)?;
474                }
475                CoordinationHandleManagerEvent::CommitRequest { epoch, .. } => {
476                    bail!(
477                        "receive commit request on epoch {} from handle {} during alter parallelism",
478                        epoch,
479                        handle_id
480                    );
481                }
482                CoordinationHandleManagerEvent::AlignInitialEpoch(epoch) => {
483                    bail!(
484                        "receive AlignInitialEpoch on epoch {} from handle {} during alter parallelism",
485                        epoch,
486                        handle_id
487                    );
488                }
489            }
490        }
491        Ok(requests.handle_ids)
492    }
493}
494
495/// Represents the coordinator worker's state machine for handling schema changes.
496///
497/// - `Running`: Normal operation, handles can be started immediately
498/// - `WaitingForFlushed`: Waiting for all pending two-phase commits to complete before starting new handles. This
499///   ensures new sink executors load the correct schema.
500enum CoordinatorWorkerState {
501    Running,
502    WaitingForFlushed(HashSet<HandleId>),
503}
504
505pub struct CoordinatorWorker {
506    handle_manager: CoordinationHandleManager,
507    /// Last epoch whose commit has been acknowledged to sink writers.
508    ///
509    /// On recovery, pending sink states are treated as already acknowledged to writers, so this is
510    /// initialized to the latest pending epoch if any. Otherwise, it starts from the latest
511    /// committed epoch persisted in `pending_sink_state`.
512    last_writer_acked_epoch: Option<u64>,
513    curr_state: CoordinatorWorkerState,
514}
515
516enum CoordinatorWorkerEvent {
517    HandleManagerEvent(HandleId, CoordinationHandleManagerEvent),
518    ReadyToCommit(u64, Option<Vec<u8>>, Option<PbSinkSchemaChange>),
519}
520
521impl CoordinatorWorker {
522    pub async fn run(
523        param: SinkParam,
524        request_rx: UnboundedReceiver<SinkWriterCoordinationHandle>,
525        db: DatabaseConnection,
526        subscriber: SinkCommittedEpochSubscriber,
527        iceberg_compact_stat_sender: UnboundedSender<IcebergSinkCompactionUpdate>,
528    ) {
529        let sink = match build_sink(param.clone()) {
530            Ok(sink) => sink,
531            Err(e) => {
532                error!(
533                    error = %e.as_report(),
534                    "unable to build sink with param {:?}",
535                    param
536                );
537                return;
538            }
539        };
540
541        dispatch_sink!(sink, sink, {
542            let coordinator =
543                match Box::pin(sink.new_coordinator(Some(iceberg_compact_stat_sender))).await {
544                    Ok(coordinator) => coordinator,
545                    Err(e) => {
546                        error!(
547                            error = %e.as_report(),
548                            "unable to build coordinator with param {:?}",
549                            param
550                        );
551                        return;
552                    }
553                };
554            Self::execute_coordinator(db, param, request_rx, coordinator, subscriber).await
555        });
556    }
557
558    pub async fn execute_coordinator(
559        db: DatabaseConnection,
560        param: SinkParam,
561        request_rx: UnboundedReceiver<SinkWriterCoordinationHandle>,
562        coordinator: SinkCommitCoordinator,
563        subscriber: SinkCommittedEpochSubscriber,
564    ) {
565        let mut worker = CoordinatorWorker {
566            handle_manager: CoordinationHandleManager {
567                param,
568                writer_handles: HashMap::new(),
569                next_handle_id: 0,
570                request_rx,
571            },
572            last_writer_acked_epoch: None,
573            curr_state: CoordinatorWorkerState::Running,
574        };
575
576        if let Err(e) = worker.run_coordination(db, coordinator, subscriber).await {
577            for handle in worker.handle_manager.writer_handles.into_values() {
578                handle.abort(Status::internal(format!(
579                    "failed to run coordination: {:?}",
580                    e.as_report()
581                )))
582            }
583        }
584    }
585
586    async fn try_handle_init_requests(
587        &mut self,
588        pending_handle_ids: &HashSet<HandleId>,
589        two_phase_handler: &mut TwoPhaseCommitHandler,
590    ) -> anyhow::Result<()> {
591        assert!(matches!(self.curr_state, CoordinatorWorkerState::Running));
592        if two_phase_handler.has_uncommitted_schema_change() {
593            // Delay handling init requests until all pending epochs are flushed.
594            self.curr_state = CoordinatorWorkerState::WaitingForFlushed(pending_handle_ids.clone());
595        } else {
596            self.handle_init_requests_impl(pending_handle_ids.clone())
597                .await?;
598        }
599        Ok(())
600    }
601
602    async fn handle_init_requests_impl(
603        &mut self,
604        pending_handle_ids: impl IntoIterator<Item = HandleId>,
605    ) -> anyhow::Result<()> {
606        let log_store_rewind_start_epoch = self.last_writer_acked_epoch;
607        self.handle_manager
608            .start(log_store_rewind_start_epoch, pending_handle_ids)?;
609        if log_store_rewind_start_epoch.is_none() {
610            let mut align_requests = AligningRequests::default();
611            while !align_requests.aligned() {
612                let (handle_id, event) = self.handle_manager.next_event().await?;
613                match event {
614                    CoordinationHandleManagerEvent::AlignInitialEpoch(initial_epoch) => {
615                        align_requests.add_new_request(
616                            handle_id,
617                            initial_epoch,
618                            self.handle_manager.vnode_bitmap(handle_id),
619                        )?;
620                    }
621                    other => {
622                        return Err(anyhow!("expect AlignInitialEpoch but got {}", other.name()));
623                    }
624                }
625            }
626            let aligned_initial_epoch = align_requests
627                .requests
628                .into_iter()
629                .max()
630                .expect("non-empty");
631            self.handle_manager
632                .ack_aligned_initial_epoch(aligned_initial_epoch)?;
633        }
634        Ok(())
635    }
636
637    async fn next_event(
638        &mut self,
639        two_phase_handler: &mut TwoPhaseCommitHandler,
640    ) -> anyhow::Result<CoordinatorWorkerEvent> {
641        if let CoordinatorWorkerState::WaitingForFlushed(pending_handle_ids) = &self.curr_state
642            && two_phase_handler.is_empty()
643        {
644            let pending_handle_ids = pending_handle_ids.clone();
645            self.handle_init_requests_impl(pending_handle_ids).await?;
646            self.curr_state = CoordinatorWorkerState::Running;
647        }
648
649        select! {
650            next_handle_event = self.handle_manager.next_event() => {
651                let (handle_id, event) = next_handle_event?;
652                Ok(CoordinatorWorkerEvent::HandleManagerEvent(handle_id, event))
653            }
654
655            next_item_to_commit = two_phase_handler.next_to_commit() => {
656                let (epoch, metadata, schema_change) = next_item_to_commit?;
657                Ok(CoordinatorWorkerEvent::ReadyToCommit(epoch, metadata, schema_change))
658            }
659        }
660    }
661
662    async fn run_coordination(
663        &mut self,
664        db: DatabaseConnection,
665        mut coordinator: SinkCommitCoordinator,
666        subscriber: SinkCommittedEpochSubscriber,
667    ) -> anyhow::Result<()> {
668        let sink_id = self.handle_manager.param.sink_id;
669
670        let mut two_phase_handler = self
671            .init_state_from_store(&db, sink_id, subscriber, &mut coordinator)
672            .await?;
673        match &mut coordinator {
674            SinkCommitCoordinator::SinglePhase(coordinator) => coordinator.init().await?,
675            SinkCommitCoordinator::TwoPhase(coordinator) => coordinator.init().await?,
676        }
677
678        let mut running_handles = self.handle_manager.wait_init_handles().await?;
679        self.try_handle_init_requests(&running_handles, &mut two_phase_handler)
680            .await?;
681
682        let mut pending_epochs: BTreeMap<u64, AligningRequests<_>> = BTreeMap::new();
683        let mut pending_new_handles = vec![];
684        loop {
685            let event = self.next_event(&mut two_phase_handler).await?;
686            let (handle_id, epoch, commit_request) = match event {
687                CoordinatorWorkerEvent::HandleManagerEvent(handle_id, event) => match event {
688                    CoordinationHandleManagerEvent::NewHandle => {
689                        pending_new_handles.push(handle_id);
690                        continue;
691                    }
692                    CoordinationHandleManagerEvent::UpdateVnodeBitmap => {
693                        running_handles = self
694                            .handle_manager
695                            .alter_parallelisms(pending_new_handles.drain(..).chain([handle_id]))
696                            .await?;
697                        self.try_handle_init_requests(&running_handles, &mut two_phase_handler)
698                            .await?;
699                        continue;
700                    }
701                    CoordinationHandleManagerEvent::Stop => {
702                        self.handle_manager.stop_handle(handle_id)?;
703                        running_handles = self
704                            .handle_manager
705                            .alter_parallelisms(pending_new_handles.drain(..))
706                            .await?;
707                        self.try_handle_init_requests(&running_handles, &mut two_phase_handler)
708                            .await?;
709
710                        continue;
711                    }
712                    CoordinationHandleManagerEvent::CommitRequest {
713                        epoch,
714                        metadata,
715                        schema_change,
716                    } => (handle_id, epoch, (metadata, schema_change)),
717                    CoordinationHandleManagerEvent::AlignInitialEpoch(_) => {
718                        bail!("receive AlignInitialEpoch after initialization")
719                    }
720                },
721                CoordinatorWorkerEvent::ReadyToCommit(epoch, metadata, schema_change) => {
722                    let start_time = Instant::now();
723                    let commit_fut = async {
724                        match &mut coordinator {
725                            SinkCommitCoordinator::SinglePhase(coordinator) => {
726                                assert!(metadata.is_none());
727                                if let Some(schema_change) = schema_change {
728                                    coordinator
729                                        .commit_schema_change(epoch, schema_change)
730                                        .instrument_await(Self::commit_span(
731                                            "single_phase_schema_change",
732                                            sink_id,
733                                            epoch,
734                                        ))
735                                        .await?;
736                                }
737                            }
738                            SinkCommitCoordinator::TwoPhase(coordinator) => {
739                                if let Some(metadata) = metadata {
740                                    coordinator
741                                        .commit_data(epoch, metadata)
742                                        .instrument_await(Self::commit_span(
743                                            "two_phase_commit_data",
744                                            sink_id,
745                                            epoch,
746                                        ))
747                                        .await?;
748                                }
749                                if let Some(schema_change) = schema_change {
750                                    coordinator
751                                        .commit_schema_change(epoch, schema_change)
752                                        .instrument_await(Self::commit_span(
753                                            "two_phase_commit_schema_change",
754                                            sink_id,
755                                            epoch,
756                                        ))
757                                        .await?;
758                                }
759                            }
760                        }
761                        Ok(())
762                    };
763                    let commit_res =
764                        run_future_with_periodic_fn(commit_fut, Duration::from_secs(5), || {
765                            warn!(
766                                elapsed = ?start_time.elapsed(),
767                                %sink_id,
768                                "committing"
769                            );
770                        })
771                        .await;
772
773                    match commit_res {
774                        Ok(_) => {
775                            two_phase_handler.ack_committed(epoch).await?;
776                        }
777                        Err(e) => {
778                            two_phase_handler.failed_committed(epoch, e);
779                        }
780                    }
781
782                    continue;
783                }
784            };
785            if !running_handles.contains(&handle_id) {
786                bail!(
787                    "receiving commit request from non-running handle {}, running handles: {:?}",
788                    handle_id,
789                    running_handles
790                );
791            }
792            pending_epochs.entry(epoch).or_default().add_new_request(
793                handle_id,
794                commit_request,
795                self.handle_manager.vnode_bitmap(handle_id),
796            )?;
797            if pending_epochs
798                .first_key_value()
799                .expect("non-empty")
800                .1
801                .aligned()
802            {
803                let (epoch, commit_requests) = pending_epochs.pop_first().expect("non-empty");
804                let mut metadatas = Vec::with_capacity(commit_requests.requests.len());
805                let mut requests = commit_requests.requests.into_iter();
806                let (first_metadata, first_schema_change) = requests.next().expect("non-empty");
807                metadatas.push(first_metadata);
808                for (metadata, schema_change) in requests {
809                    if first_schema_change != schema_change {
810                        return Err(anyhow!(
811                            "got different schema change {:?} to prev schema change {:?}",
812                            schema_change,
813                            first_schema_change
814                        ));
815                    }
816                    metadatas.push(metadata);
817                }
818
819                match &mut coordinator {
820                    SinkCommitCoordinator::SinglePhase(coordinator) => {
821                        if !metadatas.is_empty() {
822                            let start_time = Instant::now();
823                            run_future_with_periodic_fn(
824                                coordinator.commit_data(epoch, metadatas).instrument_await(
825                                    Self::commit_span("single_phase_commit_data", sink_id, epoch),
826                                ),
827                                Duration::from_secs(5),
828                                || {
829                                    warn!(
830                                        elapsed = ?start_time.elapsed(),
831                                        %sink_id,
832                                        "committing"
833                                    );
834                                },
835                            )
836                            .await
837                            .map_err(|e| anyhow!(e))?;
838                        }
839                        if first_schema_change.is_some() {
840                            persist_pre_commit_metadata(
841                                &db,
842                                sink_id as _,
843                                epoch,
844                                None,
845                                first_schema_change.as_ref(),
846                            )
847                            .await?;
848                            two_phase_handler.push_new_item(epoch, None, first_schema_change);
849                        }
850                    }
851                    SinkCommitCoordinator::TwoPhase(coordinator) => {
852                        let commit_metadata = coordinator
853                            .pre_commit(epoch, metadatas, first_schema_change.clone())
854                            .instrument_await(Self::commit_span(
855                                "two_phase_pre_commit",
856                                sink_id,
857                                epoch,
858                            ))
859                            .await?;
860                        // Persist every acknowledged epoch, even when there is no metadata or
861                        // schema change. Writers may truncate the epoch as soon as they receive
862                        // the acknowledgement, so recovery must retain the same progress. The
863                        // commit handler treats a `None`/`None` item as an external no-op before
864                        // marking it committed, and recovery re-enqueues it through the same path.
865                        persist_pre_commit_metadata(
866                            &db,
867                            sink_id as _,
868                            epoch,
869                            commit_metadata.clone(),
870                            first_schema_change.as_ref(),
871                        )
872                        .await?;
873                        two_phase_handler.push_new_item(
874                            epoch,
875                            commit_metadata,
876                            first_schema_change,
877                        );
878                    }
879                }
880
881                self.handle_manager
882                    .ack_commit(epoch, commit_requests.handle_ids)?;
883                self.last_writer_acked_epoch = Some(epoch);
884            }
885        }
886    }
887
888    /// Return `TwoPhaseCommitHandler` initialized from the persisted state in the meta store.
889    async fn init_state_from_store(
890        &mut self,
891        db: &DatabaseConnection,
892        sink_id: SinkId,
893        subscriber: SinkCommittedEpochSubscriber,
894        coordinator: &mut SinkCommitCoordinator,
895    ) -> anyhow::Result<TwoPhaseCommitHandler> {
896        let ordered_metadata = list_sink_states_ordered_by_epoch(db, sink_id as _).await?;
897
898        let mut metadata_iter = ordered_metadata.into_iter().peekable();
899        let last_committed_epoch = metadata_iter
900            .next_if(|(_, state, _, _)| matches!(state, SinkState::Committed))
901            .map(|(epoch, _, _, _)| epoch);
902
903        let pending_items = metadata_iter
904            .peeking_take_while(|(_, state, _, _)| matches!(state, SinkState::Pending))
905            .map(|(epoch, _, metadata, schema_change)| (epoch, metadata, schema_change))
906            .collect_vec();
907        self.last_writer_acked_epoch = pending_items
908            .last()
909            .map(|(epoch, _, _)| *epoch)
910            .or(last_committed_epoch);
911
912        let mut aborted_epochs = vec![];
913
914        for (epoch, state, metadata, _) in metadata_iter {
915            match state {
916                SinkState::Aborted => {
917                    if let Some(metadata) = metadata
918                        && let SinkCommitCoordinator::TwoPhase(coordinator) = coordinator
919                    {
920                        coordinator.abort(epoch, metadata).await;
921                    }
922                    aborted_epochs.push(epoch);
923                }
924                other => {
925                    unreachable!(
926                        "unexpected state {:?} after pending items at epoch {}",
927                        other, epoch
928                    );
929                }
930            }
931        }
932
933        // Records for all aborted epochs and previously committed epochs are no longer needed.
934        clean_aborted_records(db, sink_id, aborted_epochs).await?;
935
936        let (initial_hummock_committed_epoch, job_committed_epoch_rx) = subscriber(sink_id).await?;
937        let mut two_phase_handler = TwoPhaseCommitHandler::new(
938            db.clone(),
939            sink_id,
940            initial_hummock_committed_epoch,
941            job_committed_epoch_rx,
942            last_committed_epoch,
943        );
944
945        for (epoch, metadata, schema_change) in pending_items {
946            two_phase_handler.push_new_item(epoch, metadata, schema_change);
947        }
948
949        Ok(two_phase_handler)
950    }
951
952    fn commit_span(stage: &str, sink_id: SinkId, epoch: u64) -> await_tree::Span {
953        await_tree::span!("sink_coord_{stage} (sink_id {sink_id}, epoch {epoch})").long_running()
954    }
955}
risingwave_meta/manager/sink_coordination/coordinator_worker.rs

risingwave_meta/manager/sink_coordination/
coordinator_worker.rs
