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 risingwave_pb::stream_service::barrier_complete_response::IcebergPkIndexSinkMetadata as PbIcebergPkIndexSinkMetadata;
24use sea_orm::DatabaseConnection;
25use tokio::sync::Mutex;
26use tracing::warn;
27
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 phase for one epoch: persist the merged report under `pending_sink_state` (no iceberg IO).
102 /// The barrier-complete path awaits this BEFORE issuing hummock `commit_epoch`.
103 pub async fn pre_commit_epoch(
104 &self,
105 sink_id: SinkId,
106 prev_epoch: u64,
107 reports: Vec<PbIcebergPkIndexSinkMetadata>,
108 ) -> anyhow::Result<()> {
109 let coordinator = self.coordinator(sink_id)?;
110 coordinator
111 .lock()
112 .await
113 .pre_commit(prev_epoch, reports)
114 .await
115 }
116
117 /// Commit phase for one epoch: run an iceberg `overwrite_files` transaction and mark its pending row
118 /// committed. The barrier-complete path awaits this AFTER hummock `commit_epoch`.
119 pub async fn commit_epoch(&self, sink_id: SinkId) -> anyhow::Result<()> {
120 let coordinator = self.coordinator(sink_id)?;
121 coordinator.lock().await.commit().await
122 }
123
124 /// Advance the per-partial-graph committed epoch after a checkpoint completion.
125 pub fn advance_committed_epochs(
126 &self,
127 epochs: impl IntoIterator<Item = (PartialGraphId, u64)>,
128 ) {
129 self.inner.committed_epochs.advance_all(epochs);
130 }
131
132 /// Block until `sink_id`'s partial graph has committed through `target_epoch`, then return the
133 /// coordinator's committed iceberg snapshot id (a lower bound the caller must observe when it
134 /// reloads the table). Does NOT hold the coordinator lock while waiting.
135 pub async fn wait_epoch(
136 &self,
137 sink_id: SinkId,
138 target_epoch: u64,
139 ) -> anyhow::Result<Option<i64>> {
140 let partial_graph_id = self.partial_graph_of(sink_id)?;
141 self.inner
142 .committed_epochs
143 .wait(partial_graph_id, target_epoch)
144 .await;
145 let coordinator = self.coordinator(sink_id)?;
146 let snapshot_id = coordinator.lock().await.current_snapshot_id();
147 Ok(snapshot_id)
148 }
149
150 /// Unregister the given `sink_id`(s)' coordinator(s) (e.g. at DROP SINK time). Unregistering an unknown
151 /// `sink_id` is a no-op.
152 pub fn unregister_sinks(&self, sink_ids: Vec<SinkId>) {
153 let mut coordinators = self.inner.coordinators.write();
154 let mut touched_graphs = Vec::new();
155 for sink_id in sink_ids {
156 if let Some((partial_graph_id, _coord)) = coordinators.remove(&sink_id) {
157 touched_graphs.push(partial_graph_id);
158 }
159 }
160 // Drop a partial graph's committed-epoch cursor only once none of its sinks remain registered.
161 for partial_graph_id in touched_graphs {
162 if !coordinators.values().any(|(pg, _)| *pg == partial_graph_id) {
163 self.inner.committed_epochs.remove(partial_graph_id);
164 }
165 }
166 }
167
168 /// Drop every coordinator. Used at recovery time.
169 pub fn reset(&self) {
170 let mut coordinators = self.inner.coordinators.write();
171 coordinators.clear();
172 self.inner.committed_epochs.clear();
173 }
174
175 fn coordinator(&self, sink_id: SinkId) -> anyhow::Result<CoordinatorRef> {
176 self.inner
177 .coordinators
178 .read()
179 .get(&sink_id)
180 .map(|(_pg, coord)| coord.clone())
181 .ok_or_else(|| {
182 anyhow!(
183 "iceberg pk-index sink coordinator for sink {} is not registered",
184 sink_id
185 )
186 })
187 }
188
189 fn partial_graph_of(&self, sink_id: SinkId) -> anyhow::Result<PartialGraphId> {
190 self.inner
191 .coordinators
192 .read()
193 .get(&sink_id)
194 .map(|(pg, _coord)| *pg)
195 .ok_or_else(|| {
196 anyhow!(
197 "iceberg pk-index sink coordinator for sink {} is not registered",
198 sink_id
199 )
200 })
201 }
202}