krabka_remote_storage_topic/manager/assignment.rs
1//! The metadata-partition assignment reconciler and the read gate it feeds.
2//!
3//! The broker drives which `__remote_log_metadata` partitions this manager
4//! consumes, and every gated read consults the readiness state that the same
5//! reconciler records. Both halves share the `ready_targets` map and the
6//! HWM-unknown sentinel, so they belong in one module.
7
8use {
9 super::TopicBasedRemoteLogMetadataManager,
10 crate::log::PartitionStart,
11 std::{collections::HashMap, sync::PoisonError},
12 tracing::warn,
13};
14
15#[cfg(test)]
16mod tests;
17
18/// Sentinel target HWM that means "this partition is assigned but its real
19/// high-water mark is not yet known". That happens when the
20/// `high_water_marks` RPC failed, or when the partition had no entry in the
21/// returned index. The gate treats the sentinel as `NotReady`, because a real
22/// applied offset can never reach `i64::MAX`. The next `reconcile_assignment`
23/// retries the HWM fetch and replaces the sentinel with the real target.
24pub(super) const HWM_UNKNOWN: i64 = i64::MAX;
25
26/// Outcome of the per-metadata-partition readiness check that gates the
27/// gated [`RemoteLogMetadataManager`](krabka_remote_storage::RemoteLogMetadataManager)
28/// read methods
29/// ([`RemoteLogMetadataManager::remote_log_segment_metadata`](krabka_remote_storage::RemoteLogMetadataManager::remote_log_segment_metadata),
30/// [`RemoteLogMetadataManager::list_remote_log_segments`](krabka_remote_storage::RemoteLogMetadataManager::list_remote_log_segments),
31/// [`RemoteLogMetadataManager::list_remote_log_segments_by_epoch`](krabka_remote_storage::RemoteLogMetadataManager::list_remote_log_segments_by_epoch),
32/// [`RemoteLogMetadataManager::highest_offset_for_epoch`](krabka_remote_storage::RemoteLogMetadataManager::highest_offset_for_epoch)).
33pub enum ReadGate {
34 /// This broker does not consume the metadata partition, because it
35 /// neither leads nor follows any covered user-partition. Answer
36 /// a miss (`Ok(None)` or an empty list).
37 Unassigned,
38 /// Assigned, but the consumer pump has not reached the assignment-time
39 /// HWM. Answer with the retryable `Err(NotReady)`.
40 NotReady,
41 /// Assigned and caught up. Delegate to the inner cache.
42 Ready,
43}
44
45/// The read decision for an assigned metadata partition whose
46/// assignment-time high-water mark is `target`, given the offset of the last
47/// event applied on it. `applied` is `None` for a partition index outside
48/// the metadata log.
49pub(super) const fn read_gate(target: i64, applied: Option<i64>) -> ReadGate {
50 if target == 0 {
51 return ReadGate::Ready; // empty partition: nothing to catch up to
52 }
53 // A sentinel target means the real HWM is not yet known (the
54 // assignment-time fetch failed); the partition is assigned but the
55 // answer is unknown → retryable, never a false `Ok(None)`.
56 if target == HWM_UNKNOWN {
57 return ReadGate::NotReady;
58 }
59 // Defensive: a partition the metadata log does not have must NOT fail
60 // open into `Ready`, which would serve a possibly-stale or false-miss
61 // answer. Treat it as still catching up.
62 match applied {
63 Some(applied) if applied >= target - 1 => ReadGate::Ready,
64 _ => ReadGate::NotReady,
65 }
66}
67
68impl TopicBasedRemoteLogMetadataManager {
69 /// The read decision for metadata partition `mp`, used to gate
70 /// [`RemoteLogMetadataManager::remote_log_segment_metadata`](krabka_remote_storage::RemoteLogMetadataManager::remote_log_segment_metadata).
71 pub(super) fn metadata_partition_gate(&self, mp: i32) -> ReadGate {
72 let target = self
73 .ready_targets
74 .lock()
75 .unwrap_or_else(PoisonError::into_inner)
76 .get(&mp)
77 .copied();
78 // Not assigned: this broker neither leads nor follows any
79 // user-partition in `mp`, so it must not answer from any stale
80 // cache it happened to consume earlier. A genuine miss — `Ok(None)`,
81 // NOT `NotReady`.
82 let Some(target) = target else {
83 return ReadGate::Unassigned;
84 };
85 read_gate(target, self.applied_offset(mp))
86 }
87
88 /// The offset of the last event the pump applied on metadata partition
89 /// `mp`, or `None` when `mp` is not a partition of the metadata log.
90 pub(super) fn applied_offset(&self, mp: i32) -> Option<i64> {
91 let index = usize::try_from(mp).ok()?;
92 self.applied
93 .lock()
94 .unwrap_or_else(PoisonError::into_inner)
95 .get(index)
96 .copied()
97 }
98
99 /// `true` when metadata partition `mp` is assigned and caught up to its
100 /// assignment-time HWM. Tests use this to poll for catch-up.
101 #[cfg(test)]
102 pub(super) fn metadata_partition_ready(&self, mp: i32) -> bool {
103 matches!(self.metadata_partition_gate(mp), ReadGate::Ready)
104 }
105
106 /// The metadata partitions this manager is currently assigned (tracked
107 /// for readiness). Sorted ascending.
108 #[must_use]
109 pub fn assigned_metadata_partitions(&self) -> Vec<i32> {
110 let mut v: Vec<i32> = self
111 .ready_targets
112 .lock()
113 .unwrap_or_else(PoisonError::into_inner)
114 .keys()
115 .copied()
116 .collect();
117 v.sort_unstable();
118 v
119 }
120
121 /// Diff `desired` against the current assignment and drive the
122 /// [`AssignmentHandle`](crate::log::AssignmentHandle).
123 ///
124 /// This method adds newly-needed partitions, seeded from the snapshot
125 /// committed offset + 1, or from 0 when there is no committed event. It
126 /// removes partitions that are no longer needed. It records each added
127 /// partition's assignment-time HWM, so reads gate on `NotReady` until the
128 /// pump catches up.
129 ///
130 /// An HWM-fetch failure fails CLOSED. This method records a partition
131 /// whose real high-water mark it could not get with the `HWM_UNKNOWN`
132 /// sentinel target, so the gate returns the retryable `NotReady` and
133 /// never a false `Ok(None)`. Every subsequent reconcile retries such
134 /// partitions, and the broker drives a reconcile on each image change and
135 /// each reconciler tick. A transient `high_water_marks` failure therefore
136 /// self-heals: the real target replaces the sentinel as soon as the fetch
137 /// succeeds.
138 ///
139 /// A SINGLE task MUST drive this method. It is not internally serialized.
140 /// It interleaves `.await` points with reads and writes of the
141 /// `ready_targets` map under short, non-overlapping locks, so two
142 /// concurrent callers could race the add, remove, and refresh logic.
143 /// Correctness depends on the broker invoking it from exactly one
144 /// reconciler task.
145 ///
146 /// It is async because it reads the log's high-water marks. The broker
147 /// calls it from its reconciler task on the runtime, never from a
148 /// `spawn_blocking` thread.
149 pub async fn reconcile_assignment(&self, desired: &[i32]) {
150 use std::collections::HashSet;
151 let want: HashSet<i32> = desired.iter().copied().collect();
152 // Snapshot the current per-partition targets so we can both diff the
153 // assigned set and find partitions still carrying the HWM-unknown
154 // sentinel (which need a refresh). Lock released before the `.await`.
155 let current: HashMap<i32, i64> = self
156 .ready_targets
157 .lock()
158 .unwrap_or_else(PoisonError::into_inner)
159 .clone();
160 let have: HashSet<i32> = current.keys().copied().collect();
161
162 let needs_add = want.difference(&have).copied().collect::<Vec<_>>();
163 // Partitions still assigned (in want) whose recorded target is the
164 // sentinel: their HWM is still unknown, so re-attempt the fetch.
165 let needs_refresh = want
166 .iter()
167 .copied()
168 .filter(|mp| current.get(mp) == Some(&HWM_UNKNOWN))
169 .collect::<Vec<_>>();
170
171 // One HWM snapshot covers both additions and sentinel refreshes.
172 let needs_hwm = !needs_add.is_empty() || !needs_refresh.is_empty();
173 let hwms = if needs_hwm {
174 match self.log.high_water_marks().await {
175 Ok(h) => Some(h),
176 Err(e) => {
177 warn!(error = ?e, "topic-based RLMM: high_water_marks fetch failed; \
178 assigned partitions gate NotReady until a later reconcile refreshes");
179 None
180 }
181 }
182 } else {
183 None
184 };
185
186 // Resolve a partition's target HWM from the (maybe-missing) snapshot.
187 // A failed fetch (`None`) or a missing per-partition entry both yield
188 // the sentinel so the gate stays NotReady — never fail open to 0.
189 let target_for = |mp: i32| -> i64 {
190 hwms.as_ref()
191 .and_then(|h| {
192 let i = usize::try_from(mp).ok()?;
193 h.get(i).copied()
194 })
195 .unwrap_or(HWM_UNKNOWN)
196 };
197
198 for mp in needs_add {
199 // `committed_offset` is `-1` when there is no committed event
200 // (full replay), so `+ 1` lands on the resume start offset (0).
201 // Startup validated the snapshot's committed offsets, so this
202 // cannot fail; were one invalid, the partition replays in full
203 // from offset 0, as a startup with an invalid snapshot does.
204 let start_offset = match Self::resume_start_offset(self.committed_offset(mp)) {
205 Ok(start_offset) => start_offset,
206 Err(error) => {
207 warn!(
208 ?error,
209 mp, "topic-based RLMM: invalid committed offset; replaying in full"
210 );
211 0
212 }
213 };
214 self.assignment.add(PartitionStart {
215 partition: mp,
216 start_offset,
217 });
218 // Assign-but-NotReady when the HWM is unknown: the broker DOES
219 // own this partition, so leaving it Unassigned would wrongly
220 // return Ok(None). The sentinel makes the gate return NotReady.
221 self.ready_targets
222 .lock()
223 .unwrap_or_else(PoisonError::into_inner)
224 .insert(mp, target_for(mp));
225 }
226 // Replace the sentinel for already-assigned partitions whose HWM is
227 // now known (the partition stays assigned; only its target changes).
228 for mp in needs_refresh {
229 let target = target_for(mp);
230 if target != HWM_UNKNOWN {
231 let mut guard = self
232 .ready_targets
233 .lock()
234 .unwrap_or_else(PoisonError::into_inner);
235 // Only refresh if still assigned with the sentinel (a
236 // concurrent remove would have dropped it — see the
237 // single-task contract above).
238 if guard.get(&mp) == Some(&HWM_UNKNOWN) {
239 guard.insert(mp, target);
240 }
241 }
242 }
243 for mp in have.difference(&want).copied() {
244 // Unassign under the waiter lock, so a publish either sees the
245 // partition unassigned or learns of the removal when it
246 // registers its offset.
247 self.mutation_waiters.remove_partition(mp, || {
248 self.assignment.remove(mp);
249 self.ready_targets
250 .lock()
251 .unwrap_or_else(PoisonError::into_inner)
252 .remove(&mp);
253 });
254 }
255 }
256}