risingwave_meta/manager/iceberg_pk_index_sink/
manager.rs

1// Copyright 2026 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::HashMap;
16use std::sync::Arc;
17
18use anyhow::anyhow;
19use parking_lot::RwLock;
20use risingwave_common::id::PartialGraphId;
21use risingwave_connector::sink::catalog::SinkId;
22use risingwave_connector::sink::iceberg::IcebergConfig;
23use sea_orm::DatabaseConnection;
24use tokio::sync::Mutex;
25use tracing::warn;
26
27use super::IcebergPkIndexPreCommitMetadata;
28use super::committed_epoch::PartialGraphCommittedEpochs;
29use super::coordinator::IcebergPkIndexSinkCoordinator;
30
31type CoordinatorRef = Arc<Mutex<IcebergPkIndexSinkCoordinator>>;
32
33/// Manager for the Iceberg pk-index sink path, cheap to clone.
34#[derive(Clone)]
35pub struct IcebergPkIndexSinkManager {
36    inner: Arc<ManagerInner>,
37}
38
39struct ManagerInner {
40    db: DatabaseConnection,
41    /// `sink_id -> (partial_graph_id, coordinator)`. Read on every pre-commit/commit/wait (only to
42    /// clone out the ref / read the partial graph id, never held across an await); written only by
43    /// register/unregister/reset, which are rare control-plane events. The `partial_graph_id` is the
44    /// graph whose committed epoch the sink's merger waits on (its database main graph today; an
45    /// independent job's graph in the future — see `committed_epoch`).
46    coordinators: RwLock<HashMap<SinkId, (PartialGraphId, CoordinatorRef)>>,
47    /// Per-partial-graph committed-epoch cursor. A cursor entry exists exactly while a partial graph has
48    /// a registered pk-index sink: created by `ensure` in `register_sink` and dropped by `remove` in
49    /// `unregister_sinks` (and `clear` in `reset`), all under the `coordinators` write lock. Advanced on
50    /// every checkpoint completion via `advance_committed_epochs` (a no-op for partial graphs with no
51    /// registered sink).
52    committed_epochs: PartialGraphCommittedEpochs,
53}
54
55impl IcebergPkIndexSinkManager {
56    pub fn new(db: DatabaseConnection) -> Self {
57        IcebergPkIndexSinkManager {
58            inner: Arc::new(ManagerInner {
59                db,
60                coordinators: RwLock::new(HashMap::new()),
61                committed_epochs: PartialGraphCommittedEpochs::default(),
62            }),
63        }
64    }
65
66    /// Register an Iceberg pk-index sink so its commit coordinator is ready to receive epoch reports. Builds and
67    /// fully initializes the coordinator (loading the iceberg table and draining any recovered pending
68    /// commits) BEFORE inserting it, so a successful return means the sink is ready to serve. Idempotent:
69    /// registering the same `sink_id` replaces the existing coordinator.
70    pub async fn register_sink(
71        &self,
72        sink_id: SinkId,
73        partial_graph_id: PartialGraphId,
74        iceberg_config: IcebergConfig,
75    ) -> anyhow::Result<()> {
76        // Initialize (load + recover + drain) outside the map lock; this is the slow, fallible part.
77        let coordinator =
78            IcebergPkIndexSinkCoordinator::init(sink_id, iceberg_config, self.inner.db.clone())
79                .await?;
80
81        let prev = {
82            let mut coordinators = self.inner.coordinators.write();
83            let prev = coordinators.insert(
84                sink_id,
85                (partial_graph_id, Arc::new(Mutex::new(coordinator))),
86            );
87            // Start tracking this partial graph's committed epoch under the same write lock as the
88            // coordinator insert, so a cursor entry exists exactly while a sink is registered (and a
89            // `wait` racing an unregister can never observe a resurrected, never-advanced entry).
90            self.inner.committed_epochs.ensure(partial_graph_id);
91            prev
92        };
93        if prev.is_some() {
94            // Replacing an existing coordinator. Any in-flight commit on the old one keeps it alive via its
95            // own `Arc` until it finishes; the snapshot_id idempotency check guards against double-commit.
96            warn!(%sink_id, "iceberg pk-index sink coordinator re-registered; replacing previous instance");
97        }
98        Ok(())
99    }
100
101    /// Pre-commit one epoch, with an optional compaction overwrite folded into the same pending row.
102    /// The barrier-complete path awaits this before issuing Hummock `commit_epoch`.
103    pub(crate) async fn pre_commit(
104        &self,
105        input: IcebergPkIndexPreCommitMetadata,
106    ) -> anyhow::Result<()> {
107        let coordinator = self.coordinator(input.sink_id)?;
108        coordinator
109            .lock()
110            .await
111            .pre_commit(input.prev_epoch, input.reports, input.compaction)
112            .await
113    }
114
115    /// Commit phase for one epoch: run an iceberg `overwrite_files` transaction and mark its pending row
116    /// committed. The barrier-complete path awaits this AFTER hummock `commit_epoch`.
117    pub async fn commit_epoch(&self, sink_id: SinkId) -> anyhow::Result<()> {
118        let coordinator = self.coordinator(sink_id)?;
119        coordinator.lock().await.commit().await
120    }
121
122    /// Advance the per-partial-graph committed epoch after a checkpoint completion.
123    pub fn advance_committed_epochs(
124        &self,
125        epochs: impl IntoIterator<Item = (PartialGraphId, u64)>,
126    ) {
127        self.inner.committed_epochs.advance_all(epochs);
128    }
129
130    /// Block until `sink_id`'s partial graph has committed through `target_epoch`, then return the
131    /// coordinator's committed iceberg snapshot id (a lower bound the caller must observe when it
132    /// reloads the table). Does NOT hold the coordinator lock while waiting.
133    pub async fn wait_epoch(
134        &self,
135        sink_id: SinkId,
136        target_epoch: u64,
137    ) -> anyhow::Result<Option<i64>> {
138        let partial_graph_id = self.partial_graph_of(sink_id)?;
139        self.inner
140            .committed_epochs
141            .wait(partial_graph_id, target_epoch)
142            .await;
143        let coordinator = self.coordinator(sink_id)?;
144        let snapshot_id = coordinator.lock().await.current_snapshot_id();
145        Ok(snapshot_id)
146    }
147
148    /// Unregister the given `sink_id`(s)' coordinator(s) (e.g. at DROP SINK time). Unregistering an unknown
149    /// `sink_id` is a no-op.
150    pub fn unregister_sinks(&self, sink_ids: Vec<SinkId>) {
151        let mut coordinators = self.inner.coordinators.write();
152        let mut touched_graphs = Vec::new();
153        for sink_id in sink_ids {
154            if let Some((partial_graph_id, _coord)) = coordinators.remove(&sink_id) {
155                touched_graphs.push(partial_graph_id);
156            }
157        }
158        // Drop a partial graph's committed-epoch cursor only once none of its sinks remain registered.
159        for partial_graph_id in touched_graphs {
160            if !coordinators.values().any(|(pg, _)| *pg == partial_graph_id) {
161                self.inner.committed_epochs.remove(partial_graph_id);
162            }
163        }
164    }
165
166    /// Drop every coordinator. Used at recovery time.
167    pub fn reset(&self) {
168        let mut coordinators = self.inner.coordinators.write();
169        coordinators.clear();
170        self.inner.committed_epochs.clear();
171    }
172
173    fn coordinator(&self, sink_id: SinkId) -> anyhow::Result<CoordinatorRef> {
174        self.inner
175            .coordinators
176            .read()
177            .get(&sink_id)
178            .map(|(_pg, coord)| coord.clone())
179            .ok_or_else(|| {
180                anyhow!(
181                    "iceberg pk-index sink coordinator for sink {} is not registered",
182                    sink_id
183                )
184            })
185    }
186
187    fn partial_graph_of(&self, sink_id: SinkId) -> anyhow::Result<PartialGraphId> {
188        self.inner
189            .coordinators
190            .read()
191            .get(&sink_id)
192            .map(|(pg, _coord)| *pg)
193            .ok_or_else(|| {
194                anyhow!(
195                    "iceberg pk-index sink coordinator for sink {} is not registered",
196                    sink_id
197                )
198            })
199    }
200}
risingwave_meta/manager/iceberg_pk_index_sink/manager.rs

risingwave_meta/manager/iceberg_pk_index_sink/
manager.rs
