risingwave_meta/manager/iceberg_pk_index_sink/committed_epoch.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
15//! Per-`PartialGraphId` committed-epoch tracking for the iceberg pk-index sink.
16//!
17//! Tracks, per partial graph, the latest checkpoint epoch whose iceberg commit has completed. A
18//! cursor entry exists exactly while a partial graph has a registered pk-index sink:
19//! [`PartialGraphCommittedEpochs::ensure`] is called from `register_sink` and
20//! [`PartialGraphCommittedEpochs::remove`] from unregister. Both [`PartialGraphCommittedEpochs::advance_all`]
21//! and [`PartialGraphCommittedEpochs::wait`] are non-creating, so a partial graph without a sink never
22//! accumulates an entry, and a `wait` that races an unregister resolves immediately (the caller then
23//! observes the coordinator is gone and errors out) rather than blocking on a resurrected entry that
24//! would never advance.
25//!
26//! The key is a `PartialGraphId` rather than a `DatabaseId` so that a sink running in an independent
27//! partial graph (e.g. a batch-refresh job, `to_partial_graph_id(database_id, Some(job_id))`) is tracked
28//! by the very same mechanism: the barrier-completion path advances every completed partial graph, so no
29//! change here is needed to support such jobs. For a normal sink the key is its database's main graph
30//! (`to_partial_graph_id(database_id, None)`).
31//!
32//! The barrier-completion path advances the cursor on **every** checkpoint completion (even epochs where
33//! a sink reported nothing), so a merger's seed wait converges even on idle tables. Waiters block on a
34//! `tokio::sync::watch` receiver and never hold any other lock.
35//!
36//! Note: this is unrelated to RisingWave's stream watermark abstraction; it is purely a per-partial-graph
37//! "iceberg commit has progressed to epoch N" cursor.
38
39use std::collections::HashMap;
40use std::future::Future;
41
42use parking_lot::Mutex;
43use risingwave_common::id::PartialGraphId;
44use tokio::sync::watch;
45
46/// Per-`PartialGraphId` monotonic committed-epoch cursor, backed by one `watch` channel per tracked
47/// partial graph.
48#[derive(Default)]
49pub struct PartialGraphCommittedEpochs {
50 inner: Mutex<HashMap<PartialGraphId, watch::Sender<u64>>>,
51}
52
53impl PartialGraphCommittedEpochs {
54 /// Start tracking `partial_graph_id` (initialized to 0) if not already tracked.
55 pub fn ensure(&self, partial_graph_id: PartialGraphId) {
56 self.inner
57 .lock()
58 .entry(partial_graph_id)
59 .or_insert_with(|| watch::channel(0).0);
60 }
61
62 /// Advance each `(partial_graph_id, epoch)`'s committed epoch to `max(current, epoch)`
63 pub fn advance_all(&self, epochs: impl IntoIterator<Item = (PartialGraphId, u64)>) {
64 let map = self.inner.lock();
65 for (partial_graph_id, epoch) in epochs {
66 if let Some(sender) = map.get(&partial_graph_id) {
67 sender.send_if_modified(|cur| {
68 if epoch > *cur {
69 *cur = epoch;
70 true
71 } else {
72 false
73 }
74 });
75 }
76 }
77 }
78
79 /// Advance a single partial graph's committed epoch (test-only convenience over `advance_all`).
80 #[cfg(test)]
81 fn advance(&self, partial_graph_id: PartialGraphId, epoch: u64) {
82 self.advance_all([(partial_graph_id, epoch)]);
83 }
84
85 /// Resolve once `partial_graph_id`'s committed epoch is `>= target`. If the partial graph is not
86 /// tracked (never registered, or unregistered while this call raced), resolve immediately — the caller
87 /// then observes the coordinator is gone and errors out, instead of blocking forever on an entry that
88 /// nothing would advance.
89 ///
90 /// Returns an owned (`'static`) future rather than an `async fn` borrowing `&self`, so callers can
91 /// `tokio::spawn` the wait without holding a borrow of `PartialGraphCommittedEpochs` for the lifetime
92 /// of the wait. The subscription is taken eagerly (before returning); the wait holds no lock.
93 pub fn wait(
94 &self,
95 partial_graph_id: PartialGraphId,
96 target: u64,
97 ) -> impl Future<Output = ()> + 'static {
98 let rx = self
99 .inner
100 .lock()
101 .get(&partial_graph_id)
102 .map(|s| s.subscribe());
103 async move {
104 let Some(mut rx) = rx else {
105 // Not tracked (unregistered / never registered): do not block.
106 return;
107 };
108 while *rx.borrow_and_update() < target {
109 if rx.changed().await.is_err() {
110 // Sender dropped (partial graph removed). Nothing more will advance it; stop waiting.
111 return;
112 }
113 }
114 }
115 }
116
117 /// Stop tracking one partial graph (unregister of its last sink).
118 pub fn remove(&self, partial_graph_id: PartialGraphId) {
119 self.inner.lock().remove(&partial_graph_id);
120 }
121
122 /// Stop tracking every partial graph (global recovery reset).
123 pub fn clear(&self) {
124 self.inner.lock().clear();
125 }
126}
127
128#[cfg(test)]
129mod tests {
130 use std::time::Duration;
131
132 use super::*;
133
134 fn pg(id: u64) -> PartialGraphId {
135 PartialGraphId::new(id)
136 }
137
138 #[tokio::test]
139 async fn wait_returns_immediately_when_already_reached() {
140 let epochs = PartialGraphCommittedEpochs::default();
141 epochs.ensure(pg(1));
142 epochs.advance(pg(1), 100);
143 // Already >= target: must resolve without an advance.
144 tokio::time::timeout(Duration::from_secs(1), epochs.wait(pg(1), 50))
145 .await
146 .expect("wait should resolve immediately");
147 }
148
149 #[tokio::test]
150 async fn wait_unblocks_on_later_advance() {
151 let epochs = PartialGraphCommittedEpochs::default();
152 epochs.ensure(pg(1));
153 let handle = tokio::spawn(epochs.wait(pg(1), 100));
154 // Not yet reached.
155 assert!(!handle.is_finished());
156 epochs.advance(pg(1), 100);
157 tokio::time::timeout(Duration::from_secs(1), handle)
158 .await
159 .expect("wait should unblock after advance")
160 .expect("task ok");
161 }
162
163 #[tokio::test]
164 async fn advance_is_monotonic_max() {
165 let epochs = PartialGraphCommittedEpochs::default();
166 epochs.ensure(pg(1));
167 epochs.advance(pg(1), 100);
168 epochs.advance(pg(1), 50); // must not regress
169 tokio::time::timeout(Duration::from_secs(1), epochs.wait(pg(1), 100))
170 .await
171 .expect("committed epoch must not regress below 100");
172 }
173
174 #[tokio::test]
175 async fn partial_graphs_are_independent() {
176 let epochs = PartialGraphCommittedEpochs::default();
177 epochs.ensure(pg(1));
178 epochs.ensure(pg(2));
179 epochs.advance(pg(1), 100);
180 // pg(2) is tracked but has not advanced; a target on pg(2) must still be pending.
181 let pending =
182 tokio::time::timeout(Duration::from_millis(200), epochs.wait(pg(2), 10)).await;
183 assert!(
184 pending.is_err(),
185 "pg(2) committed epoch should not be satisfied by pg(1)"
186 );
187 }
188
189 #[tokio::test]
190 async fn wait_on_untracked_graph_resolves_immediately() {
191 let epochs = PartialGraphCommittedEpochs::default();
192 // No `ensure`: a wait racing an unregister (or on a never-registered graph) must not block.
193 tokio::time::timeout(Duration::from_secs(1), epochs.wait(pg(9), 100))
194 .await
195 .expect("wait on an untracked partial graph must not block");
196 }
197
198 #[tokio::test]
199 async fn advance_before_ensure_is_noop() {
200 let epochs = PartialGraphCommittedEpochs::default();
201 epochs.advance(pg(9), 100); // untracked -> no-op, must not create an entry at 100
202 epochs.ensure(pg(9)); // now tracked, freshly at 0
203 let pending =
204 tokio::time::timeout(Duration::from_millis(200), epochs.wait(pg(9), 50)).await;
205 assert!(
206 pending.is_err(),
207 "advance on an untracked partial graph must not create/set an entry"
208 );
209 }
210}