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::{JobId, 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_jobs` (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 sink coordinators belonging to the given streaming jobs. Non-sink and unknown
149 /// job IDs are no-ops.
150 pub fn unregister_jobs(&self, job_ids: impl IntoIterator<Item = JobId>) {
151 let mut coordinators = self.inner.coordinators.write();
152 let mut touched_graphs = Vec::new();
153 for job_id in job_ids {
154 let sink_id = job_id.as_sink_id();
155 if let Some((partial_graph_id, _coord)) = coordinators.remove(&sink_id) {
156 touched_graphs.push(partial_graph_id);
157 }
158 }
159 // Drop a partial graph's committed-epoch cursor only once none of its sinks remain registered.
160 for partial_graph_id in touched_graphs {
161 if !coordinators.values().any(|(pg, _)| *pg == partial_graph_id) {
162 self.inner.committed_epochs.remove(partial_graph_id);
163 }
164 }
165 }
166
167 /// Drop every coordinator. Used at recovery time.
168 pub fn reset(&self) {
169 let mut coordinators = self.inner.coordinators.write();
170 coordinators.clear();
171 self.inner.committed_epochs.clear();
172 }
173
174 fn coordinator(&self, sink_id: SinkId) -> anyhow::Result<CoordinatorRef> {
175 self.inner
176 .coordinators
177 .read()
178 .get(&sink_id)
179 .map(|(_pg, coord)| coord.clone())
180 .ok_or_else(|| {
181 anyhow!(
182 "iceberg pk-index sink coordinator for sink {} is not registered",
183 sink_id
184 )
185 })
186 }
187
188 fn partial_graph_of(&self, sink_id: SinkId) -> anyhow::Result<PartialGraphId> {
189 self.inner
190 .coordinators
191 .read()
192 .get(&sink_id)
193 .map(|(pg, _coord)| *pg)
194 .ok_or_else(|| {
195 anyhow!(
196 "iceberg pk-index sink coordinator for sink {} is not registered",
197 sink_id
198 )
199 })
200 }
201}