krabka_remote_storage_topic/manager.rs
1//! [`TopicBasedRemoteLogMetadataManager`]: the production
2//! [`RemoteLogMetadataManager`](krabka_remote_storage::RemoteLogMetadataManager)
3//! implementation backed by a publish and subscribe [`MetadataEventLog`].
4//!
5//! The manager keeps the canonical in-memory view in an
6//! [`InmemoryRemoteLogMetadataManager`], so the lifecycle state machine is the
7//! single source of truth for cache mutation. It uses the
8//! [`MetadataEventLog`] as the durable event log.
9//!
10//! Lifecycle:
11//!
12//! - [`TopicBasedRemoteLogMetadataManager::start`][]: load any on-disk
13//! snapshot into the cache and spawn the consumer pump subscribed to
14//! NOTHING. The broker then drives the consumed set with
15//! [`TopicBasedRemoteLogMetadataManager::reconcile_assignment`]. That call
16//! adds only the `__remote_log_metadata` partitions that cover the
17//! user-partitions this broker leads or follows. `NotReady` gates a
18//! newly-added partition until the pump reaches the HWM observed at
19//! assignment time. A partition this broker does not consume is a genuine
20//! `Ok(None)`, and the manager never serves it from any stale cache.
21//! - Mutation calls, which are `add`, `update`, and `put_partition_delete`:
22//! serialize, publish, and wait until the consumer pump has applied
23//! the published offset to the inner cache. The sync return means
24//! "the event has been recorded and is visible to local reads".
25//! - Read calls: pure local lookups against the inner cache.
26//! - Drop / [`TopicBasedRemoteLogMetadataManager::shutdown`]: cancel the consumer pump.
27
28use {
29 krabka_remote_storage::{InmemoryRemoteLogMetadataManager, RemoteStorageError},
30 krabka_units::prelude::{StdDurationExt as _, TimeExt as _},
31 std::{
32 collections::HashMap,
33 path::PathBuf,
34 sync::{Arc, Mutex, PoisonError},
35 time::Duration,
36 },
37 tokio::{runtime::Handle, task::JoinHandle, time},
38 tokio_util::sync::CancellationToken,
39 tracing::warn,
40};
41
42mod assignment;
43mod mutations;
44mod persistence;
45mod pump;
46mod readiness;
47mod rlmm;
48#[cfg(test)]
49pub(crate) mod test_support;
50
51pub use self::readiness::MetadataPartitionReadiness;
52use {
53 self::{mutations::MutationWaiters, persistence::ResumePlan, pump::pump_loop},
54 crate::{
55 error::MetadataLogError,
56 log::{AssignmentHandle, MetadataEventLog},
57 snapshot::{SNAPSHOT_FILE_NAME, Snapshot},
58 },
59};
60
61/// Production [`RemoteLogMetadataManager`](krabka_remote_storage::RemoteLogMetadataManager)
62/// backed by the `__remote_log_metadata` topic through a [`MetadataEventLog`]
63/// adapter.
64///
65/// Construct it with [`Self::start`]. That call loads any on-disk snapshot,
66/// but it consumes no metadata partitions until [`Self::reconcile_assignment`]
67/// adds the broker's leader-derived and follower-derived set.
68pub struct TopicBasedRemoteLogMetadataManager {
69 log: Arc<dyn MetadataEventLog>,
70 inner: Arc<InmemoryRemoteLogMetadataManager>,
71 applied: Arc<Mutex<Vec<i64>>>,
72 mutation_waiters: Arc<MutationWaiters>,
73 runtime: Handle,
74 shutdown: CancellationToken,
75 pump: Mutex<Option<JoinHandle<()>>>,
76 /// Directory the on-disk RLMM cache snapshot is written to (one
77 /// [`SNAPSHOT_FILE_NAME`](crate::snapshot::SNAPSHOT_FILE_NAME) file).
78 snapshot_dir: PathBuf,
79 /// Handle of the background snapshotter task. `Drop` aborts it, and
80 /// [`Self::shutdown_and_flush`] joins it.
81 snapshotter: Mutex<Option<JoinHandle<()>>>,
82 /// Live assignment handle for the metadata-log subscription. The manager
83 /// holds it so the resume-from-snapshot logic and the per-broker
84 /// partition-assignment logic can mutate the consumed set at runtime.
85 /// [`Self::reconcile_assignment`] drives it.
86 assignment: Arc<dyn AssignmentHandle>,
87 /// Per-metadata-partition committed offsets loaded from the snapshot at
88 /// `start()`, indexed by metadata partition. `-1` means there is no
89 /// committed event for that partition, so it needs a full replay. This
90 /// map is the single canonical source for resume-offset lookups. The
91 /// assignment reconciler reads it with [`Self::committed_offset`] when it
92 /// dynamically adds a partition, so that partition starts at
93 /// `committed + 1`.
94 committed_offsets: Vec<i64>,
95 /// Metadata partition → target HWM observed at assignment time.
96 /// Presence means this manager is currently assigned that partition.
97 /// Reads for a user-partition that hashes into it return
98 /// [`RemoteStorageError::NotReady`] until `applied[mp] >= target - 1`.
99 /// An empty map means no partition is assigned and every read is a miss.
100 ready_targets: Arc<Mutex<HashMap<i32, i64>>>,
101}
102
103impl TopicBasedRemoteLogMetadataManager {
104 /// Actual partition count of the opened metadata topic. An existing
105 /// topic's count takes precedence over its provisioning configuration.
106 #[must_use]
107 pub fn metadata_partition_count(&self) -> i32 { self.log.partition_count() }
108
109 /// Load any on-disk snapshot into the cache and spawn the consumer
110 /// pump with an empty assignment. The manager consumes nothing until the
111 /// broker drives [`Self::reconcile_assignment`].
112 ///
113 /// `runtime` must be a Tokio runtime handle that lives at least as long
114 /// as the returned manager. The synchronous
115 /// [`RemoteLogMetadataManager`](krabka_remote_storage::RemoteLogMetadataManager)
116 /// methods bridge to this handle with `block_on`, so a caller must NOT
117 /// call them from a task that runs on this same runtime. The broker
118 /// invokes them only through `spawn_blocking`, which is the only
119 /// supported call pattern.
120 ///
121 /// # Errors
122 ///
123 /// Returns an error when the event log reports a negative partition
124 /// count. Nothing else fails, because the consumed set starts empty.
125 pub fn start(
126 log: Arc<dyn MetadataEventLog>,
127 runtime: Handle,
128 snapshot_dir: PathBuf,
129 snapshot_interval: Duration,
130 ) -> Result<Arc<Self>, RemoteStorageError> {
131 // The parameter is a `Duration` because the broker's config layer hands
132 // one over; the cadence is a time extent everywhere below this line.
133 let snapshot_interval = snapshot_interval.as_time();
134 let count = log.partition_count();
135 let n = usize::try_from(count).map_err(|_| {
136 MetadataLogError::Other(format!("negative metadata partition count {count}"))
137 .into_storage()
138 })?;
139 let mutation_waiters = Arc::new(MutationWaiters::default());
140 let inner = Arc::new(InmemoryRemoteLogMetadataManager::new());
141 let shutdown = CancellationToken::new();
142
143 // Load the snapshot (if any) ONCE. Validate its resume cursors before
144 // seeding the cache from its dump; `resume_from_snapshot` turns the
145 // loaded snapshot into the per-partition committed offsets through
146 // the same proved adapter that dynamic assignment calls.
147 // On absence/corruption, committed[] is all -1 (full replay) and the
148 // cache stays empty — never fatal.
149 let snapshot = match Snapshot::load(&snapshot_dir.join(SNAPSHOT_FILE_NAME)) {
150 Ok(snap) => snap,
151 Err(e) => {
152 warn!(error = ?e, "topic-based RLMM: snapshot corrupt; starting from empty cache");
153 None
154 }
155 };
156 // A freshly-started manager consumes NOTHING. The broker drives
157 // the consumed set via [`Self::reconcile_assignment`], adding only the
158 // metadata partitions covering user-partitions this broker leads or
159 // follows (each resumed at its snapshot `committed + 1`). This is what
160 // makes an unassigned partition a genuine `Ok(None)` rather than a
161 // false hit from globally-replayed state.
162 let committed = match Self::resume_from_snapshot(snapshot.as_ref(), n) {
163 Ok(ResumePlan { committed, .. }) => {
164 if let Some(snap) = &snapshot {
165 inner.import(snap.dump.clone());
166 }
167 committed
168 }
169 Err(e) => {
170 warn!(error = ?e, "topic-based RLMM: snapshot cursor invalid; starting from empty cache");
171 Self::resume_from_snapshot(None, n)
172 .map_err(|error| MetadataLogError::Other(error.to_string()).into_storage())?
173 .committed
174 }
175 };
176
177 // Pre-seed `applied` to the committed offsets so readiness checks for
178 // a later-added partition only block on the delta from committed+1 to
179 // the assignment-time HWM.
180 let applied = Arc::new(Mutex::new(committed.clone()));
181
182 let (stream, assignment_handle) = log.subscribe(Vec::new());
183 let pump = runtime.spawn(pump_loop(
184 stream,
185 Arc::clone(&inner),
186 Arc::clone(&applied),
187 Arc::clone(&mutation_waiters),
188 shutdown.clone(),
189 ));
190
191 let manager = Arc::new(Self {
192 log,
193 inner,
194 applied,
195 mutation_waiters,
196 runtime,
197 shutdown,
198 pump: Mutex::new(Some(pump)),
199 snapshot_dir,
200 snapshotter: Mutex::new(None),
201 assignment: assignment_handle,
202 committed_offsets: committed,
203 ready_targets: Arc::new(Mutex::new(HashMap::new())),
204 });
205
206 // Spawn the periodic snapshotter: flush whenever the cache
207 // advanced since the last write, plus a final flush on shutdown.
208 let snapshotter = {
209 let weak = Arc::downgrade(&manager);
210 let shutdown = manager.shutdown.clone();
211 manager.runtime.spawn(async move {
212 let mut last_written: i64 = -1;
213 loop {
214 tokio::select! {
215 biased;
216 () = shutdown.cancelled() => return,
217 () = time::sleep(snapshot_interval.to_std()) => {}
218 }
219 let Some(m) = weak.upgrade() else { return };
220 // Only write when the cache advanced since the last snapshot.
221 let highest = {
222 let applied = m.applied.lock().unwrap_or_else(PoisonError::into_inner);
223 applied.iter().copied().max().unwrap_or(-1)
224 };
225 if highest > last_written {
226 match m.write_snapshot() {
227 Ok(written) => last_written = written,
228 Err(e) => {
229 warn!(error = ?e, "topic-based RLMM: periodic snapshot failed");
230 }
231 }
232 }
233 }
234 })
235 };
236 *manager
237 .snapshotter
238 .lock()
239 .unwrap_or_else(PoisonError::into_inner) = Some(snapshotter);
240
241 // Nothing is consumed at bootstrap (empty assignment), so the manager
242 // is immediately ready. Per-partition catch-up after a later
243 // `reconcile_assignment` is governed by `metadata_partition_ready`,
244 // which gates reads with `NotReady` until the pump reaches the
245 // assignment-time HWM.
246 Ok(manager)
247 }
248
249 /// Cancel the consumer pump. Read methods continue to work against
250 /// whatever the pump applied before shutdown. Pending and subsequent
251 /// mutation methods return a closed-log error.
252 pub fn shutdown(&self) { self.shutdown.cancel(); }
253
254 /// Cancel the pump and the snapshotter, then write a final snapshot that
255 /// captures everything applied so far. It is safe to call this once on a
256 /// graceful shutdown.
257 pub async fn shutdown_and_flush(&self) {
258 self.shutdown.cancel();
259 // Take the handle out of the lock BEFORE awaiting it, so the
260 // (sync) mutex is not held across the await point.
261 let handle = self
262 .snapshotter
263 .lock()
264 .unwrap_or_else(PoisonError::into_inner)
265 .take();
266 // Let the snapshotter observe cancellation and stop touching
267 // `applied` before we take the final consistent capture.
268 if let Some(h) = handle {
269 let _ = h.await;
270 }
271 if let Err(e) = self.write_snapshot() {
272 warn!(error = ?e, "topic-based RLMM: final snapshot flush failed");
273 }
274 }
275}
276
277impl Drop for TopicBasedRemoteLogMetadataManager {
278 fn drop(&mut self) {
279 self.shutdown.cancel();
280 let pump = self
281 .pump
282 .lock()
283 .unwrap_or_else(PoisonError::into_inner)
284 .take();
285 if let Some(handle) = pump {
286 handle.abort();
287 }
288 let snapshotter = self
289 .snapshotter
290 .lock()
291 .unwrap_or_else(PoisonError::into_inner)
292 .take();
293 if let Some(handle) = snapshotter {
294 handle.abort();
295 }
296 }
297}