krabka_remote_storage_topic/log/

in_process.rs

//! [`InProcessMetadataEventLog`], the in-memory broadcast-channel fixture that
//! implements [`MetadataEventLog`] without a Kafka cluster behind it.
//!
//! This module holds the fixture's shared state, which is the per-partition
//! record vectors and the per-subscription cursors, together with the
//! [`MetadataEventLog`] implementation over them. The two child modules act on
//! that state: [`self::handle`] carries the [`AssignmentHandle`] the fixture
//! hands back, and [`self::live`] carries the broadcast filter that decides
//! which live record a subscription sees. They are descendants of this module
//! so the state can stay private to this file.

use {
    self::{handle::InProcessAssignmentHandle, live::filtered_broadcast},
    super::{
        AssignmentHandle, MetadataEventLog, MetadataEventRecord, MetadataEventStream,
        PartitionStart, RangeVisitor,
    },
    crate::error::MetadataLogError,
    async_trait::async_trait,
    bytes::Bytes,
    futures_util::{
        StreamExt,
        stream::{self, unfold},
    },
    std::{
        collections::HashMap,
        sync::{Arc, Mutex, MutexGuard, PoisonError, atomic::AtomicU64},
    },
    tokio::sync::{broadcast, mpsc},
};

mod handle;
mod live;

/// Single-process [`MetadataEventLog`] for unit tests and for the
/// multi-broker fixture. Multiple manager instances that clone the same `Arc`
/// observe each other's writes.
pub struct InProcessMetadataEventLog {
    inner: Arc<InProcessInner>,
}

/// Live-assignment cursor for one partition within a subscription.
#[derive(Debug, Clone, Copy)]
struct PartitionCursor {
    /// Next offset that the backlog or live path has NOT yet delivered. The
    /// log filters out records below this offset.
    next: i64,
    /// When set, the log forwards live records for this partition through
    /// the `inject` FIFO rather than emits them directly on the broadcast
    /// path. A partition added mid-stream sets this, so its live appends
    /// queue *behind* its already-injected backlog. Without it,
    /// `stream::select` could interleave a live record ahead of undrained
    /// backlog and violate per-partition publish order. Initially-assigned
    /// partitions leave this `false`. Their backlog goes through the chained
    /// snapshot stream, which fully drains before any live record.
    via_inject: bool,
}

/// Per-subscription live assignment, plus a sender to inject backlog when a
/// partition is added mid-stream. A monotonically-increasing subscription id
/// keys it, so multiple subscribers stay independent.
struct SubscriptionState {
    /// partition -> cursor. Presence in the map means the partition is
    /// assigned.
    assigned: Mutex<HashMap<i32, PartitionCursor>>,
    /// Inject backlog records in FIFO publish order. For `add`-ed partitions
    /// this also injects live records.
    inject: mpsc::UnboundedSender<MetadataEventRecord>,
}

struct InProcessInner {
    /// `log[partition][offset] = event`.
    log: Mutex<Vec<Vec<MetadataEventRecord>>>,
    /// Notify subscribers of new writes.
    tx: broadcast::Sender<MetadataEventRecord>,
    /// Constant for the life of the log.
    partition_count: i32,
    /// Live subscriptions, keyed by id, for assignment filtering and
    /// mid-stream backlog injection.
    subscriptions: Mutex<HashMap<u64, Arc<SubscriptionState>>>,
    /// Allocates subscription ids.
    next_sub_id: AtomicU64,
}

impl InProcessMetadataEventLog {
    /// Construct an empty log with `partition_count` partitions.
    ///
    /// # Panics
    ///
    /// Panics when `partition_count <= 0`.
    #[must_use]
    pub fn new(partition_count: i32) -> Arc<Self> {
        assert2::assert!(
            let Some(cap) = usize::try_from(partition_count).ok().filter(|&cap| cap > 0),
            "partition_count must be positive"
        );
        let (tx, _rx) = broadcast::channel(1024);
        Arc::new(Self {
            inner: Arc::new(InProcessInner {
                log: Mutex::new(vec![Vec::new(); cap]),
                tx,
                partition_count,
                subscriptions: Mutex::new(HashMap::new()),
                next_sub_id: AtomicU64::new(0),
            }),
        })
    }

    /// The records of every partition, whether or not a thread panicked
    /// while it held the lock.
    fn log(&self) -> MutexGuard<'_, Vec<Vec<MetadataEventRecord>>> {
        self.inner
            .log
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
    }

    fn append(
        &self,
        partition: i32,
        key: Option<Bytes>,
        event: Option<Bytes>,
    ) -> Result<i64, MetadataLogError> {
        let out_of_range = || MetadataLogError::PartitionOutOfRange {
            partition,
            count: self.inner.partition_count,
        };
        if partition < 0 || partition >= self.inner.partition_count {
            return Err(out_of_range());
        }
        let idx = usize::try_from(partition).map_err(|_| out_of_range())?;
        // Hold the partition lock across the broadcast.send so that any
        // concurrent subscribe() observes either the appended record in
        // its snapshot or as a forwarded broadcast — never both.
        let mut guard = self.log();
        let log_for_p = guard.get_mut(idx).ok_or_else(out_of_range)?;
        let offset = i64::try_from(log_for_p.len()).map_err(|_| {
            MetadataLogError::Other(format!(
                "partition {partition} has more than i64::MAX records"
            ))
        })?;
        let tombstone = event.is_none();
        let record = MetadataEventRecord {
            partition,
            offset,
            key,
            payload: event.unwrap_or_default(),
            tombstone,
        };
        log_for_p.push(record.clone());
        // `send` only errors when there are no active receivers; that
        // is fine — the record is still durable in the in-memory log
        // and any later subscriber's snapshot will see it.
        let _ = self.inner.tx.send(record);
        drop(guard);
        Ok(offset)
    }
}

#[async_trait]
impl MetadataEventLog for InProcessMetadataEventLog {
    fn partition_count(&self) -> i32 { self.inner.partition_count }

    async fn publish(&self, partition: i32, event: Bytes) -> Result<i64, MetadataLogError> {
        self.append(partition, None, Some(event))
    }

    async fn publish_keyed(
        &self,
        partition: i32,
        key: Bytes,
        event: Option<Bytes>,
    ) -> Result<i64, MetadataLogError> {
        self.append(partition, Some(key), event)
    }

    fn subscribe(
        &self,
        assignment: Vec<PartitionStart>,
    ) -> (MetadataEventStream, Arc<dyn AssignmentHandle>) {
        use std::sync::atomic::Ordering;

        // Bracket snapshot + broadcast subscribe under the log lock so each
        // published record is seen exactly once (snapshot xor live).
        let guard = self.log();
        let rx = self.inner.tx.subscribe();

        // Initial assigned set: partition -> next live offset (= current
        // len), so the broadcast path forwards only records published after
        // subscribe; everything earlier comes from the snapshot below.
        let mut assigned: HashMap<i32, PartitionCursor> = HashMap::new();
        let mut snapshot: Vec<MetadataEventRecord> = Vec::new();
        for ps in &assignment {
            let Ok(idx) = usize::try_from(ps.partition) else {
                continue;
            };
            let Some(records) = guard.get(idx) else {
                continue;
            };
            let begin = usize::try_from(ps.start_offset.max(0)).unwrap_or(usize::MAX);
            for record in records.iter().skip(begin) {
                snapshot.push(record.clone());
            }
            assigned.insert(ps.partition, PartitionCursor {
                next: i64::try_from(records.len()).unwrap_or(i64::MAX),
                // Initially-assigned: backlog rides the chained
                // snapshot stream, so live records can go direct.
                via_inject: false,
            });
        }

        let (inject_tx, inject_rx) = mpsc::unbounded_channel::<MetadataEventRecord>();
        let state = Arc::new(SubscriptionState {
            assigned: Mutex::new(assigned),
            inject: inject_tx,
        });
        let sub_id = self.inner.next_sub_id.fetch_add(1, Ordering::Relaxed);
        self.inner
            .subscriptions
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .insert(sub_id, Arc::clone(&state));
        drop(guard);

        let snapshot_stream = stream::iter(snapshot);
        let inject_stream = unfold(inject_rx, |mut rx| async move {
            rx.recv().await.map(|r| (r, rx))
        });
        let live = filtered_broadcast(rx, Arc::clone(&state), Arc::clone(&self.inner));
        // Snapshot first (subscribe-time backlog), then a merge of injected
        // backlog (from `add`) and assignment-filtered live records.
        let merged = stream::select(inject_stream, live);
        let stream = snapshot_stream.chain(merged).boxed();

        let handle: Arc<dyn AssignmentHandle> = Arc::new(InProcessAssignmentHandle {
            inner: Arc::clone(&self.inner),
            sub_id,
        });
        (stream, handle)
    }

    async fn high_water_marks(&self) -> Result<Vec<i64>, MetadataLogError> {
        self.log()
            .iter()
            .map(|v| {
                i64::try_from(v.len()).map_err(|_| {
                    MetadataLogError::Other("a partition has more than i64::MAX records".into())
                })
            })
            .collect()
    }

    async fn visit_range(
        &self,
        partition: i32,
        start: i64,
        end: i64,
        visit: &mut RangeVisitor<'_>,
    ) -> Result<(), MetadataLogError> {
        let out_of_range = || MetadataLogError::PartitionOutOfRange {
            partition,
            count: self.inner.partition_count,
        };
        let idx = usize::try_from(partition).map_err(|_| out_of_range())?;
        // Copy the range out, so `visit` never runs under the log lock.
        let copied = self.log().get(idx).map(|records| {
            let range: Vec<MetadataEventRecord> = records
                .iter()
                .filter(|record| start <= record.offset && record.offset < end)
                .cloned()
                .collect();
            (records.len(), range)
        });
        let (len, range) = copied.ok_or_else(out_of_range)?;
        let high_water_mark = i64::try_from(len).map_err(|_| {
            MetadataLogError::Other(format!(
                "partition {partition} has more than i64::MAX records"
            ))
        })?;
        if end > high_water_mark {
            return Err(MetadataLogError::Other(format!(
                "partition {partition} ends at {high_water_mark}, before {end}"
            )));
        }
        range.into_iter().try_for_each(visit)
    }
}

#[cfg(test)]
mod tests {
    use {
        super::*,
        assert2::{assert, check},
        futures_util::StreamExt,
        tokio::time,
    };

    #[tokio::test]
    async fn publish_assigns_monotonic_offsets() {
        let log = InProcessMetadataEventLog::new(2);
        check!(log.publish(0, Bytes::from_static(b"a")).await.unwrap() == 0);
        check!(log.publish(0, Bytes::from_static(b"b")).await.unwrap() == 1);
        check!(log.publish(1, Bytes::from_static(b"c")).await.unwrap() == 0);
        let hwms = log.high_water_marks().await.unwrap();
        assert!(hwms == vec![2, 1]);
    }

    /// Every record `visit_range` hands over, in the order it hands them.
    async fn visited(
        log: &InProcessMetadataEventLog,
        partition: i32,
        start: i64,
        end: i64,
    ) -> Result<Vec<MetadataEventRecord>, MetadataLogError> {
        let mut records = Vec::new();
        log.visit_range(partition, start, end, &mut |record| {
            records.push(record);
            Ok(())
        })
        .await?;
        Ok(records)
    }

    #[tokio::test]
    async fn visit_range_hands_over_the_half_open_range_without_subscribing() {
        let log = InProcessMetadataEventLog::new(2);
        for value in [b"a", b"b", b"c"] {
            log.publish(0, Bytes::from_static(value)).await.unwrap();
        }

        let record = |offset, payload| MetadataEventRecord {
            partition: 0,
            offset,
            key: None,
            payload: Bytes::from_static(payload),
            tombstone: false,
        };
        check!(visited(&log, 0, 1, 3).await.unwrap() == vec![record(1, b"b"), record(2, b"c")]);
        check!(visited(&log, 1, 0, 0).await.unwrap().is_empty());
        check!(
            visited(&log, 0, 0, 4).await.unwrap_err().to_string()
                == "metadata log error: partition 0 ends at 3, before 4"
        );
        check!(log.inner.subscriptions.lock().unwrap().is_empty());
    }

    #[yare::parameterized(
        negative = { -1 },
        upper_bound = { 2 },
    )]
    #[test_macro(tokio::test)]
    async fn visit_range_rejects_out_of_range_partitions(partition: i32) {
        let log = InProcessMetadataEventLog::new(2);

        check!(matches!(
            visited(&log, partition, 0, 0).await,
            Err(MetadataLogError::PartitionOutOfRange { partition: got, count: 2 })
                if got == partition
        ));
        check!(log.inner.subscriptions.lock().unwrap().is_empty());
    }

    #[tokio::test]
    async fn visit_range_stops_at_the_first_visitor_error() {
        let log = InProcessMetadataEventLog::new(1);
        for value in [b"a", b"b", b"c"] {
            log.publish(0, Bytes::from_static(value)).await.unwrap();
        }

        let mut seen = Vec::new();
        let error = log
            .visit_range(0, 0, 3, &mut |record| {
                seen.push(record.offset);
                if record.offset == 1 {
                    Err(MetadataLogError::Closed)
                } else {
                    Ok(())
                }
            })
            .await
            .unwrap_err();

        check!(matches!(error, MetadataLogError::Closed));
        check!(seen == vec![0, 1]);
    }

    #[tokio::test]
    async fn keyed_tombstones_survive_replay() {
        let log = InProcessMetadataEventLog::new(1);
        log.publish_keyed(
            0,
            Bytes::from_static(b"range"),
            Some(Bytes::from_static(b"object")),
        )
        .await
        .unwrap();
        log.publish_keyed(0, Bytes::from_static(b"range"), None)
            .await
            .unwrap();
        let (mut stream, _) = log.subscribe(vec![PartitionStart {
            partition: 0,
            start_offset: 0,
        }]);

        let value = stream.next().await.unwrap();
        let tombstone = stream.next().await.unwrap();
        assert!(value.key.as_deref() == Some(b"range".as_slice()));
        assert!(!value.tombstone);
        assert!(tombstone.key.as_deref() == Some(b"range".as_slice()));
        assert!(tombstone.tombstone);
    }

    #[tokio::test]
    async fn subscribe_replays_history_then_forwards_new_writes() {
        let log = InProcessMetadataEventLog::new(1);
        log.publish(0, Bytes::from_static(b"a")).await.unwrap();
        log.publish(0, Bytes::from_static(b"b")).await.unwrap();
        let (mut stream, _h) = log.subscribe(vec![PartitionStart {
            partition: 0,
            start_offset: 0,
        }]);
        let a = stream.next().await.unwrap();
        let b = stream.next().await.unwrap();
        assert!(a.payload.as_ref() == b"a");
        assert!(b.payload.as_ref() == b"b");
        log.publish(0, Bytes::from_static(b"c")).await.unwrap();
        let c = stream.next().await.unwrap();
        check!(c.payload.as_ref() == b"c");
        check!(c.partition == 0);
        check!(c.offset == 2);
    }

    #[tokio::test]
    async fn subscribe_attached_after_history_still_sees_history() {
        let log = InProcessMetadataEventLog::new(1);
        for i in 0..5 {
            log.publish(0, Bytes::copy_from_slice(&[i])).await.unwrap();
        }
        let (mut stream, _h) = log.subscribe(vec![PartitionStart {
            partition: 0,
            start_offset: 0,
        }]);
        for i in 0..5 {
            let r = stream.next().await.unwrap();
            assert!(r.payload.as_ref() == &[i]);
            assert!(r.offset == i64::from(i));
        }
    }

    #[tokio::test]
    async fn lag_replays_exact_and_larger_bursts_once_in_order() {
        use std::time::Duration;

        for count in [1025_i64, 4097] {
            let log = InProcessMetadataEventLog::new(1);
            let (mut stream, _handle) = log.subscribe(vec![PartitionStart::new(0, 0)]);
            for offset in 0..count {
                let payload = Bytes::copy_from_slice(&offset.to_be_bytes());
                assert!(log.publish(0, payload).await.unwrap() == offset);
            }

            time::timeout(Duration::from_secs(3), async {
                for expected in 0..count {
                    let record = stream.next().await.expect("retained record is replayed");
                    assert!(record.offset == expected);
                    assert!(record.payload.as_ref() == expected.to_be_bytes());
                }
            })
            .await
            .expect("burst replay completes");

            // A new subscriber remains an independent retained-log control.
            let (mut fresh, _fresh_handle) = log.subscribe(vec![PartitionStart::new(0, 0)]);
            assert!(fresh.next().await.unwrap().offset == 0);
        }
    }

    #[tokio::test]
    async fn snapshot_then_live_delivery_has_no_overlap_or_gap() {
        let log = InProcessMetadataEventLog::new(1);
        for offset in 0..4_i64 {
            log.publish(0, Bytes::copy_from_slice(&offset.to_be_bytes()))
                .await
                .unwrap();
        }
        let (mut stream, _handle) = log.subscribe(vec![PartitionStart::new(0, 2)]);
        for offset in 4..8_i64 {
            log.publish(0, Bytes::copy_from_slice(&offset.to_be_bytes()))
                .await
                .unwrap();
        }

        for expected in 2..8_i64 {
            let record = stream.next().await.unwrap();
            assert!(record.offset == expected);
            assert!(record.payload.as_ref() == expected.to_be_bytes());
        }
    }

    #[tokio::test]
    async fn assignment_added_before_overflow_replays_its_backlog_and_live_events() {
        let log = InProcessMetadataEventLog::new(2);
        let (mut stream, handle) = log.subscribe(vec![PartitionStart::new(0, 0)]);
        for offset in 0..4_i64 {
            log.publish(1, Bytes::copy_from_slice(&offset.to_be_bytes()))
                .await
                .unwrap();
        }
        handle.add(PartitionStart::new(1, 1));
        log.publish(1, Bytes::from_static(b"live")).await.unwrap();
        for _ in 0..1025 {
            log.publish(0, Bytes::from_static(b"overflow"))
                .await
                .unwrap();
        }

        let mut partition_zero = Vec::new();
        let mut partition_one = Vec::new();
        for _ in 0..1029 {
            let record = stream.next().await.unwrap();
            match record.partition {
                0 => partition_zero.push(record.offset),
                1 => partition_one.push((record.offset, record.payload)),
                _ => panic!("unexpected partition"),
            }
        }
        assert!(partition_zero == (0..1025).collect::<Vec<_>>());
        assert!(
            partition_one
                == vec![
                    (1, Bytes::copy_from_slice(&1_i64.to_be_bytes())),
                    (2, Bytes::copy_from_slice(&2_i64.to_be_bytes())),
                    (3, Bytes::copy_from_slice(&3_i64.to_be_bytes())),
                    (4, Bytes::from_static(b"live")),
                ]
        );
    }

    #[tokio::test]
    async fn publish_out_of_range_is_rejected() {
        let log = InProcessMetadataEventLog::new(2);
        let err = log.publish(5, Bytes::from_static(b"x")).await.unwrap_err();
        assert!(matches!(err, MetadataLogError::PartitionOutOfRange { .. }));
    }

    #[tokio::test]
    async fn two_subscribers_see_the_same_history() {
        let log = InProcessMetadataEventLog::new(1);
        log.publish(0, Bytes::from_static(b"a")).await.unwrap();
        let (mut s1, _h1) = log.subscribe(vec![PartitionStart {
            partition: 0,
            start_offset: 0,
        }]);
        let (mut s2, _h2) = log.subscribe(vec![PartitionStart {
            partition: 0,
            start_offset: 0,
        }]);
        log.publish(0, Bytes::from_static(b"b")).await.unwrap();
        for s in [&mut s1, &mut s2] {
            assert!(s.next().await.unwrap().payload.as_ref() == b"a");
            assert!(s.next().await.unwrap().payload.as_ref() == b"b");
        }
    }

    #[tokio::test]
    async fn subscribe_delivers_only_assigned_partitions_from_start_offset() {
        #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
        struct Delivered {
            partition: i32,
            offset: i64,
            payload: Vec<u8>,
        }
        let log = InProcessMetadataEventLog::new(3);
        // partition 0: a,b,c ; partition 1: x,y ; partition 2: z
        for p0 in [b"a".as_slice(), b"b", b"c"] {
            log.publish(0, Bytes::copy_from_slice(p0)).await.unwrap();
        }
        for p1 in [b"x".as_slice(), b"y"] {
            log.publish(1, Bytes::copy_from_slice(p1)).await.unwrap();
        }
        log.publish(2, Bytes::from_static(b"z")).await.unwrap();

        // Assign partition 0 from offset 1 and partition 1 from offset 0;
        // partition 2 is NOT assigned.
        let (mut stream, _h) =
            log.subscribe(vec![PartitionStart::new(0, 1), PartitionStart::new(1, 0)]);

        let delivered = |partition, offset, payload: &[u8]| Delivered {
            partition,
            offset,
            payload: payload.to_vec(),
        };
        let mut got: Vec<Delivered> = Vec::new();
        for _ in 0..3 {
            let r = stream.next().await.unwrap();
            got.push(delivered(r.partition, r.offset, &r.payload));
        }
        got.sort();
        assert2::assert!(
            got == vec![
                delivered(0, 1, b"b"),
                delivered(0, 2, b"c"),
                delivered(1, 0, b"x"),
            ]
        );
        // partition 1 offset 1 ("y") is the only remaining assigned record.
        let r = stream.next().await.unwrap();
        check!(r.partition == 1);
        check!(r.offset == 1);
        check!(r.payload.as_ref() == b"y");
    }

    #[tokio::test]
    async fn live_appends_only_for_assigned_partitions() {
        let log = InProcessMetadataEventLog::new(2);
        let (mut stream, _h) = log.subscribe(vec![PartitionStart {
            partition: 0,
            start_offset: 0,
        }]);
        // Unassigned partition write must not appear.
        log.publish(1, Bytes::from_static(b"skip")).await.unwrap();
        log.publish(0, Bytes::from_static(b"keep")).await.unwrap();
        let r = stream.next().await.unwrap();
        check!(r.partition == 0);
        check!(r.payload.as_ref() == b"keep");
    }
}