krabka_remote_storage_topic/manager/
rlmm.rs

1//! The [`RemoteLogMetadataManager`] implementation for the topic-backed
2//! manager, with the publish-and-wait bridge its mutation methods use.
3//!
4//! Every mutation serializes an event, publishes it to the metadata log, and
5//! blocks until the consumer pump returns the result of applying the published
6//! offset. The bridge lives in the sibling `mutations` module. Each gated read
7//! consults [`ReadGate`] before it delegates to the inner cache.
8
9use {
10    super::{TopicBasedRemoteLogMetadataManager, assignment::ReadGate},
11    crate::{error::MetadataLogError, partitioning::metadata_partition_for, serde::MetadataEvent},
12    krabka_ids::LeaderEpoch,
13    krabka_remote_storage::{
14        InmemoryRemoteLogMetadataManager, RemoteLogMetadataManager, RemoteLogSegmentMetadata,
15        RemoteLogSegmentMetadataUpdate, RemoteLogSegmentState, RemotePartitionDeleteMetadata,
16        RemoteStorageError, TopicIdPartition,
17    },
18    std::sync::PoisonError,
19};
20
21impl TopicBasedRemoteLogMetadataManager {
22    /// Apply the assignment/readiness gate before consulting cached metadata.
23    fn read_if_ready<T: Default>(
24        &self,
25        partition: &TopicIdPartition,
26        read: impl FnOnce(&InmemoryRemoteLogMetadataManager) -> Result<T, RemoteStorageError>,
27    ) -> Result<T, RemoteStorageError> {
28        let mp = metadata_partition_for(partition, self.log.partition_count())
29            .map_err(MetadataLogError::into_storage)?;
30        match self.metadata_partition_gate(mp) {
31            // An unassigned partition is a genuine miss, even if the cache
32            // retains entries consumed under an earlier assignment.
33            ReadGate::Unassigned => Ok(T::default()),
34            // Catching up is retryable and must remain distinct from a miss.
35            ReadGate::NotReady => Err(RemoteStorageError::NotReady { partition: mp }),
36            ReadGate::Ready => read(&self.inner),
37        }
38    }
39
40    /// Whether metadata partition `mp` has applied everything below the
41    /// log's current high-water mark. The caller has already passed the
42    /// read gate, so the assignment-time target is reached; this check asks
43    /// the log again so it also covers events committed after assignment.
44    fn applied_through_current_hwm(&self, mp: i32) -> Result<bool, RemoteStorageError> {
45        let hwms = self.runtime.block_on(async {
46            tokio::select! {
47                biased;
48                () = self.shutdown.cancelled() => Err(MetadataLogError::Closed),
49                hwms = self.log.high_water_marks() => hwms,
50            }
51        });
52        let hwms = hwms.map_err(MetadataLogError::into_storage)?;
53        let idx = usize::try_from(mp).ok();
54        let Some(hwm) = idx.and_then(|i| hwms.get(i).copied()) else {
55            return Err(MetadataLogError::Other(format!(
56                "no high-water mark for metadata partition {mp}"
57            ))
58            .into_storage());
59        };
60        let applied = idx.and_then(|i| {
61            self.applied
62                .lock()
63                .unwrap_or_else(PoisonError::into_inner)
64                .get(i)
65                .copied()
66        });
67        Ok(applied.is_some_and(|applied| applied >= hwm - 1))
68    }
69}
70
71impl RemoteLogMetadataManager for TopicBasedRemoteLogMetadataManager {
72    fn add_remote_log_segment_metadata(
73        &self,
74        metadata: RemoteLogSegmentMetadata,
75    ) -> Result<(), RemoteStorageError> {
76        // Mirror the in-memory manager's eager precondition: fail
77        // fast before paying a round trip through Kafka.
78        if metadata.state() != RemoteLogSegmentState::CopySegmentStarted {
79            return Err(RemoteStorageError::InvalidAdd {
80                id: metadata.remote_log_segment_id().clone(),
81                reason: format!(
82                    "starting state must be CopySegmentStarted, got {:?}",
83                    metadata.state()
84                ),
85            });
86        }
87        let tp = metadata.remote_log_segment_id().topic_id_partition.clone();
88        let event = MetadataEvent::AddSegment(metadata)
89            .encode()
90            .map_err(|error| RemoteStorageError::InvalidArgument(error.to_string()))?;
91        self.publish_and_wait(&tp, event)
92    }
93
94    fn update_remote_log_segment_metadata(
95        &self,
96        update: RemoteLogSegmentMetadataUpdate,
97    ) -> Result<(), RemoteStorageError> {
98        let tp = update.remote_log_segment_id.topic_id_partition.clone();
99        let event = MetadataEvent::UpdateSegment(update)
100            .encode()
101            .map_err(|error| RemoteStorageError::InvalidArgument(error.to_string()))?;
102        self.publish_and_wait(&tp, event)
103    }
104
105    fn remote_log_segment_metadata(
106        &self,
107        topic_id_partition: &TopicIdPartition,
108        leader_epoch: LeaderEpoch,
109        offset: i64,
110    ) -> Result<Option<RemoteLogSegmentMetadata>, RemoteStorageError> {
111        self.read_if_ready(topic_id_partition, |inner| {
112            inner.remote_log_segment_metadata(topic_id_partition, leader_epoch, offset)
113        })
114    }
115
116    fn highest_offset_for_epoch(
117        &self,
118        topic_id_partition: &TopicIdPartition,
119        leader_epoch: LeaderEpoch,
120    ) -> Result<Option<i64>, RemoteStorageError> {
121        self.read_if_ready(topic_id_partition, |inner| {
122            inner.highest_offset_for_epoch(topic_id_partition, leader_epoch)
123        })
124    }
125
126    fn list_remote_log_segments(
127        &self,
128        topic_id_partition: &TopicIdPartition,
129    ) -> Result<Vec<RemoteLogSegmentMetadata>, RemoteStorageError> {
130        self.read_if_ready(topic_id_partition, |inner| {
131            inner.list_remote_log_segments(topic_id_partition)
132        })
133    }
134
135    fn list_remote_log_segments_by_epoch(
136        &self,
137        topic_id_partition: &TopicIdPartition,
138        leader_epoch: LeaderEpoch,
139    ) -> Result<Vec<RemoteLogSegmentMetadata>, RemoteStorageError> {
140        self.read_if_ready(topic_id_partition, |inner| {
141            inner.list_remote_log_segments_by_epoch(topic_id_partition, leader_epoch)
142        })
143    }
144
145    fn put_remote_partition_delete_metadata(
146        &self,
147        metadata: RemotePartitionDeleteMetadata,
148    ) -> Result<(), RemoteStorageError> {
149        let tp = metadata.topic_id_partition.clone();
150        let event = MetadataEvent::PartitionDelete(metadata)
151            .encode()
152            .map_err(|error| RemoteStorageError::InvalidArgument(error.to_string()))?;
153        self.publish_and_wait(&tp, event)
154    }
155
156    // Like every sync method here, this blocks on the manager's runtime, so it
157    // is called only from `spawn_blocking` threads.
158    fn caught_up(&self, topic_id_partition: &TopicIdPartition) -> Result<bool, RemoteStorageError> {
159        let mp = metadata_partition_for(topic_id_partition, self.log.partition_count())
160            .map_err(MetadataLogError::into_storage)?;
161        match self.metadata_partition_gate(mp) {
162            // This broker does not consume `mp`, so it cannot vouch for it.
163            ReadGate::Unassigned => Ok(false),
164            ReadGate::NotReady => Err(RemoteStorageError::NotReady { partition: mp }),
165            ReadGate::Ready => self.applied_through_current_hwm(mp),
166        }
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use {
173        super::*,
174        crate::{
175            log::{InProcessMetadataEventLog, MetadataEventLog},
176            manager::test_support::{
177                HwmFlakyLog, finish, on_blocking, seed, seed_finished, start_manager,
178                start_manager_all, started, tp,
179            },
180        },
181        assert2::{assert, check},
182        krabka_remote_storage::RemotePartitionDeleteState,
183        std::{
184            sync::Arc,
185            time::{Duration, Instant},
186        },
187        tokio::task,
188    };
189
190    krabka_macros::remote_segment_check!(check_finished_segment);
191
192    #[tokio::test(flavor = "multi_thread")]
193    async fn add_finish_query_round_trip() {
194        let log: Arc<dyn MetadataEventLog> = InProcessMetadataEventLog::new(4);
195        let m = start_manager_all(log).await;
196        seed_finished(&m, &[seed(10, 0, 99)]).await;
197
198        check_finished_segment(m.as_ref(), &tp());
199        m.shutdown();
200    }
201
202    #[tokio::test(flavor = "multi_thread")]
203    async fn add_with_wrong_state_is_rejected_eagerly() {
204        let log: Arc<dyn MetadataEventLog> = InProcessMetadataEventLog::new(2);
205        let m = start_manager(Arc::clone(&log));
206        // Force a non-Started state via the lifecycle helper.
207        let bad = started(10, 0, 9).with_update(&finish(10)).unwrap();
208        let m2 = Arc::clone(&m);
209        let err = on_blocking(move || m2.add_remote_log_segment_metadata(bad).unwrap_err()).await;
210        assert!(matches!(err, RemoteStorageError::InvalidAdd { .. }));
211        // Eager rejection means nothing was published.
212        assert!(log.high_water_marks().await.unwrap() == vec![0; 2]);
213        m.shutdown();
214    }
215
216    #[tokio::test(flavor = "multi_thread")]
217    async fn partition_delete_lifecycle_round_trip() {
218        let log: Arc<dyn MetadataEventLog> = InProcessMetadataEventLog::new(2);
219        let m = start_manager_all(log).await;
220        for state in [
221            RemotePartitionDeleteState::DeletePartitionMarked,
222            RemotePartitionDeleteState::DeletePartitionStarted,
223            RemotePartitionDeleteState::DeletePartitionFinished,
224        ] {
225            let m2 = Arc::clone(&m);
226            on_blocking(move || {
227                m2.put_remote_partition_delete_metadata(RemotePartitionDeleteMetadata {
228                    topic_id_partition: tp(),
229                    state,
230                    event_timestamp_ms: 500,
231                    broker_id: 1,
232                })
233                .unwrap();
234            })
235            .await;
236        }
237        m.shutdown();
238    }
239
240    /// How a `caught_up` call ended, in a form a test can compare whole.
241    #[derive(Debug, PartialEq, Eq)]
242    enum CaughtUp {
243        Answer(bool),
244        NotReady(i32),
245        Io,
246        /// Any other error, by its debug form.
247        Other(String),
248    }
249
250    impl From<Result<bool, RemoteStorageError>> for CaughtUp {
251        fn from(result: Result<bool, RemoteStorageError>) -> Self {
252            match result {
253                Ok(answer) => Self::Answer(answer),
254                Err(RemoteStorageError::NotReady { partition }) => Self::NotReady(partition),
255                Err(RemoteStorageError::Io(_)) => Self::Io,
256                Err(other) => Self::Other(format!("{other:?}")),
257            }
258        }
259    }
260
261    /// The state the manager is in when the test asks `caught_up`.
262    #[derive(Clone, Copy)]
263    enum Scenario {
264        /// Assigned, and every published event applied.
265        Seeded,
266        /// Assigned and past the read gate, but another writer committed an
267        /// event that the pump has not applied yet.
268        Behind,
269        /// The manager does not consume the metadata partition.
270        Unassigned,
271        /// Assigned, but the assignment-time HWM is still unknown.
272        NotReady,
273        /// Assigned and ready, but the current HWM cannot be fetched.
274        HwmFailure,
275        /// The manager has shut down.
276        Shutdown,
277    }
278
279    async fn ask(m: &Arc<TopicBasedRemoteLogMetadataManager>) -> CaughtUp {
280        let m = Arc::clone(m);
281        on_blocking(move || m.caught_up(&tp())).await.into()
282    }
283
284    #[yare::parameterized(
285        seeded = { Scenario::Seeded, |_| CaughtUp::Answer(true) },
286        behind = { Scenario::Behind, |_| CaughtUp::Answer(false) },
287        unassigned = { Scenario::Unassigned, |_| CaughtUp::Answer(false) },
288        not_ready = { Scenario::NotReady, CaughtUp::NotReady },
289        hwm_failure = { Scenario::HwmFailure, |_| CaughtUp::Io },
290        shutdown = { Scenario::Shutdown, |_| CaughtUp::Io },
291    )]
292    #[test_macro(tokio::test(flavor = "multi_thread"))]
293    async fn caught_up_answers_from_the_current_high_water_mark(
294        scenario: Scenario,
295        expected: fn(i32) -> CaughtUp,
296    ) {
297        let log = HwmFlakyLog::new(4);
298        let mp = metadata_partition_for(&tp(), log.partition_count()).unwrap();
299        // The expectation takes the metadata partition, which a NotReady
300        // answer names.
301        let expected = expected(mp);
302        let (m, got) = match scenario {
303            Scenario::Seeded => {
304                let m = start_manager_all(Arc::<HwmFlakyLog>::clone(&log)).await;
305                seed_finished(&m, &[seed(10, 0, 99)]).await;
306                let got = ask(&m).await;
307                (m, got)
308            }
309            Scenario::Behind => {
310                let m = start_manager_all(Arc::<HwmFlakyLog>::clone(&log)).await;
311                seed_finished(&m, &[seed(10, 0, 99)]).await;
312                // The pump takes this lock before it applies a record, so
313                // holding it keeps the next record committed but unapplied.
314                let pump_stall = m.mutation_waiters.stall_pump().await;
315                let event = MetadataEvent::AddSegment(started(11, 100, 199))
316                    .encode()
317                    .unwrap();
318                log.publish(mp, event).await.unwrap();
319                let got = ask(&m).await;
320                drop(pump_stall);
321                (m, got)
322            }
323            Scenario::Unassigned => {
324                let m = start_manager(Arc::<HwmFlakyLog>::clone(&log));
325                let got = ask(&m).await;
326                (m, got)
327            }
328            Scenario::NotReady => {
329                let m = start_manager(Arc::<HwmFlakyLog>::clone(&log));
330                log.set_fail_hwm(true);
331                m.reconcile_assignment(&[mp]).await;
332                let got = ask(&m).await;
333                (m, got)
334            }
335            Scenario::HwmFailure => {
336                let m = start_manager_all(Arc::<HwmFlakyLog>::clone(&log)).await;
337                log.set_fail_hwm(true);
338                let got = ask(&m).await;
339                (m, got)
340            }
341            Scenario::Shutdown => {
342                let m = start_manager_all(Arc::<HwmFlakyLog>::clone(&log)).await;
343                m.shutdown();
344                let got = ask(&m).await;
345                (m, got)
346            }
347        };
348        check!(got == expected);
349        m.shutdown();
350    }
351
352    #[tokio::test(flavor = "multi_thread")]
353    async fn caught_up_turns_true_once_the_pump_applies_a_peer_write() {
354        let log: Arc<dyn MetadataEventLog> = InProcessMetadataEventLog::new(4);
355        let reader = start_manager_all(Arc::clone(&log)).await;
356        let writer = start_manager_all(log).await;
357        let stall = reader.mutation_waiters.stall_pump().await;
358        let w = Arc::clone(&writer);
359        on_blocking(move || {
360            w.add_remote_log_segment_metadata(started(10, 0, 99))
361                .unwrap();
362        })
363        .await;
364        assert!(ask(&reader).await == CaughtUp::Answer(false));
365        drop(stall);
366        let deadline = Instant::now() + Duration::from_secs(2);
367        while ask(&reader).await != CaughtUp::Answer(true) {
368            assert!(Instant::now() < deadline, "reader never caught up");
369            task::yield_now().await;
370        }
371        assert!(reader.list_remote_log_segments(&tp()).unwrap().len() == 1);
372        reader.shutdown();
373        writer.shutdown();
374    }
375}
krabka_remote_storage_topic/manager/rlmm.rs

krabka_remote_storage_topic/manager/
rlmm.rs
