ruststream-rdkafka 0.7.1

//! Topics, partitions and the records they retain.
//!
//! A partition is an append-only log and a record's offset is its index in it. Offsets are never
//! reused, and a transaction's commit or abort takes an offset of its own for the control marker
//! a consumer never sees, which is why a consumer's offsets can have gaps here as on a cluster.

use std::sync::Arc;

use bytes::Bytes;

use crate::convert;

/// One wire header of a stored record: a name and a value, which Kafka lets be null.
pub(crate) type WireHeader = (String, Option<Bytes>);

/// Who produced a record inside a transaction: the transactional id and the producer epoch the
/// pairing was fenced with.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct ProducerId {
    pub(crate) transactional_id: String,
    pub(crate) epoch: u64,
}

/// Whether a data record is visible, and to which readers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum Visibility {
    /// Produced outside a transaction, or by one that committed.
    Committed,
    /// Produced inside a transaction that is still open: a `read_committed` reader stops in front
    /// of it, a `read_uncommitted` one reads it.
    Pending(ProducerId),
    /// Produced inside a transaction that aborted: only a `read_uncommitted` reader sees it.
    Aborted,
}

/// What occupies one offset.
#[derive(Debug, Clone)]
pub(crate) enum Entry {
    /// A record a producer wrote.
    Data(Arc<Stored>, Visibility),
    /// A transaction marker: it takes an offset and is delivered to nobody.
    Control,
}

/// A record as the log retains it.
#[derive(Debug)]
pub(crate) struct Stored {
    /// Cluster-wide publish order, so the publish log reads back in the order records were
    /// produced across a topic's partitions.
    pub(crate) sequence: u64,
    /// `CreateTime`, in milliseconds since the epoch, stamped when the record was produced.
    pub(crate) timestamp: i64,
    pub(crate) payload: Bytes,
    pub(crate) key: Option<Bytes>,
    pub(crate) headers: Vec<WireHeader>,
}

/// How a reader treats records of transactions (librdkafka's `isolation.level`).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Isolation {
    /// Stops at the first record of an open transaction and skips aborted ones.
    ReadCommitted,
    /// Reads every data record as it is written.
    ReadUncommitted,
}

impl Isolation {
    /// The next offset at or after `from` this reader is handed a record from, with that record,
    /// or `None` when the reader has to wait.
    pub(crate) fn next(self, entries: &[Entry], from: i64) -> Option<(i64, &Arc<Stored>)> {
        let start = usize::try_from(from).unwrap_or(0);
        for (index, entry) in entries.iter().enumerate().skip(start) {
            match entry {
                Entry::Data(stored, Visibility::Committed) => {
                    return Some((offset_of(index), stored));
                }
                Entry::Data(stored, Visibility::Pending(_) | Visibility::Aborted)
                    if self == Self::ReadUncommitted =>
                {
                    return Some((offset_of(index), stored));
                }
                // The last stable offset: a committed reader reads nothing past an open
                // transaction until it resolves.
                Entry::Data(_, Visibility::Pending(_)) => return None,
                Entry::Control | Entry::Data(_, Visibility::Aborted) => {}
            }
        }
        None
    }

    /// How many records a reader at `from` is still to be handed right now.
    pub(crate) fn backlog(self, entries: &[Entry], from: i64) -> usize {
        let mut count = 0;
        let mut at = from;
        while let Some((offset, _)) = self.next(entries, at) {
            count += 1;
            at = offset + 1;
        }
        count
    }
}

/// A log index as the offset it is.
pub(crate) fn offset_of(index: usize) -> i64 {
    i64::try_from(index).unwrap_or(i64::MAX)
}

/// One partition's log.
#[derive(Debug, Default)]
pub(crate) struct Partition {
    pub(crate) entries: Vec<Entry>,
}

impl Partition {
    /// The high watermark: the offset the next record takes.
    pub(crate) fn end(&self) -> i64 {
        offset_of(self.entries.len())
    }

    /// The offset of the first record stamped at or after `when`, or the end of the log.
    pub(crate) fn offset_for_time(&self, when: i64) -> i64 {
        self.entries
            .iter()
            .position(|entry| matches!(entry, Entry::Data(stored, _) if stored.timestamp >= when))
            .map_or_else(|| self.end(), offset_of)
    }
}

/// A topic: its partitions, and the cursor of its keyless placement.
#[derive(Debug)]
pub(crate) struct Topic {
    pub(crate) partitions: Vec<Partition>,
    keyless: usize,
}

/// How many partitions a topic has when it comes into being: the broker default
/// (`num.partitions`), which the test stand runs with too.
pub(crate) const DEFAULT_PARTITIONS: usize = 1;

impl Topic {
    pub(crate) fn new(partitions: usize) -> Self {
        Self {
            partitions: (0..partitions.max(1))
                .map(|_| Partition::default())
                .collect(),
            keyless: 0,
        }
    }

    /// The partition a record lands on: the one named, the one its key hashes to
    /// (librdkafka's default `consistent_random` partitioner, a CRC-32 of the key), or the next
    /// one in turn for a record with neither.
    pub(crate) fn place(&mut self, named: Option<i32>, key: Option<&[u8]>) -> Option<usize> {
        let count = self.partitions.len();
        match (named, key) {
            (Some(partition), _) => usize::try_from(partition)
                .ok()
                .filter(|index| *index < count),
            (None, Some(key)) if !key.is_empty() => {
                let hash = usize::try_from(crc32(key)).unwrap_or(0);
                Some(hash % count)
            }
            (None, _) => {
                let index = self.keyless % count;
                self.keyless = self.keyless.wrapping_add(1);
                Some(index)
            }
        }
    }
}

/// The most framing librdkafka reserves for one record of the v2 message format: its length,
/// attributes, timestamp and offset deltas, the key and value lengths and the header count, each
/// counted at the widest its varint can be (`RD_KAFKAP_MESSAGE_V2_MAX_OVERHEAD`).
const RECORD_OVERHEAD: usize = 5 + 1 + 10 + 5 + 5 + 5 + 5;

/// The size librdkafka holds a record to against `message.max.bytes` before it takes it: the
/// record's framing, its key and value, and each header as it is serialized, a varint length
/// before its name and before its value.
pub(crate) fn record_size(payload: &[u8], key: Option<&[u8]>, headers: &[WireHeader]) -> usize {
    let headers: usize = headers
        .iter()
        .map(|(name, value)| {
            let value = value.as_ref().map_or(0, Bytes::len);
            name.len() + value + length_varint(name.len()) + length_varint(value)
        })
        .sum();
    RECORD_OVERHEAD + key.map_or(0, <[u8]>::len) + payload.len() + headers
}

/// How many bytes the varint librdkafka writes a length as takes: the length is zigzag encoded,
/// which doubles it, then written seven bits to a byte.
fn length_varint(len: usize) -> usize {
    // Doubling a non-zero length adds one bit to it.
    let bits = match len {
        0 => 0,
        _ => (usize::BITS - len.leading_zeros() + 1) as usize,
    };
    bits.div_ceil(7).max(1)
}

/// The wire headers a publish's header map becomes: every header except the ones this crate maps
/// onto the record itself (the key) or strips (the exactly-once source coordinates).
pub(crate) fn wire_headers(headers: &ruststream::HeaderMap) -> Vec<WireHeader> {
    headers
        .iter()
        .filter(|(name, _)| convert::rides_the_wire(name))
        .map(|(name, value)| (name.to_owned(), Some(Bytes::copy_from_slice(value))))
        .collect()
}

/// The CRC-32 librdkafka's consistent partitioner hashes a key with (the IEEE polynomial, as in
/// zlib).
fn crc32(bytes: &[u8]) -> u32 {
    let mut crc = !0u32;
    for byte in bytes {
        crc ^= u32::from(*byte);
        for _ in 0..8 {
            let mask = (crc & 1).wrapping_neg();
            crc = (crc >> 1) ^ (0xEDB8_8320 & mask);
        }
    }
    !crc
}

#[cfg(test)]
mod tests {
    use super::*;

    fn data(visibility: Visibility) -> Entry {
        Entry::Data(
            Arc::new(Stored {
                sequence: 0,
                timestamp: 0,
                payload: Bytes::new(),
                key: None,
                headers: Vec::new(),
            }),
            visibility,
        )
    }

    fn pending() -> Visibility {
        Visibility::Pending(ProducerId {
            transactional_id: "tx".to_owned(),
            epoch: 1,
        })
    }

    #[test]
    fn the_crc_is_the_ieee_one() {
        // The check value of CRC-32/ISO-HDLC, the one zlib and librdkafka compute.
        assert_eq!(crc32(b"123456789"), 0xCBF4_3926);
    }

    #[test]
    fn a_committed_reader_stops_at_an_open_transaction_and_skips_an_aborted_one() {
        let entries = vec![
            data(Visibility::Committed),
            data(Visibility::Aborted),
            Entry::Control,
            data(Visibility::Committed),
            data(pending()),
            data(Visibility::Committed),
        ];
        let committed = Isolation::ReadCommitted;
        assert_eq!(committed.next(&entries, 0).map(|(at, _)| at), Some(0));
        assert_eq!(committed.next(&entries, 1).map(|(at, _)| at), Some(3));
        assert_eq!(committed.next(&entries, 4).map(|(at, _)| at), None);
        assert_eq!(committed.backlog(&entries, 0), 2);

        let uncommitted = Isolation::ReadUncommitted;
        assert_eq!(uncommitted.next(&entries, 1).map(|(at, _)| at), Some(1));
        assert_eq!(uncommitted.backlog(&entries, 0), 5);
    }

    #[test]
    fn a_record_counts_as_librdkafka_counts_it() {
        assert_eq!(record_size(b"", None, &[]), 36);
        assert_eq!(record_size(&vec![0; 999_964], None, &[]), 1_000_000);
        // A 63-byte length doubles into one varint byte, a 64-byte one into two.
        let header =
            |name: usize, value: usize| ("n".repeat(name), Some(Bytes::from(vec![0; value])));
        assert_eq!(
            record_size(b"", Some(b"key"), &[header(1, 63)]),
            36 + 3 + 1 + 63 + 1 + 1
        );
        assert_eq!(
            record_size(b"", None, &[header(1, 64)]),
            36 + 1 + 64 + 1 + 2
        );
    }

    #[test]
    fn a_key_always_lands_on_one_partition_and_a_named_partition_must_exist() {
        let mut topic = Topic::new(4);
        let first = topic.place(None, Some(b"tenant-1"));
        assert_eq!(first, topic.place(None, Some(b"tenant-1")));
        assert_eq!(topic.place(Some(3), Some(b"tenant-1")), Some(3));
        assert_eq!(topic.place(Some(4), None), None);
        assert_eq!(topic.place(Some(-1), None), None);
        let keyless: Vec<_> = (0..4).map(|_| topic.place(None, None)).collect();
        assert_eq!(keyless, [Some(0), Some(1), Some(2), Some(3)]);
    }
}