krabka_bench_driver/
payload.rs

1//! Record-payload generator.
2//!
3//! The first 24 bytes of every produced record hold `(magic_be, scenario_id_be,
4//! send_unix_nanos_be)`, so a consumer can compute the end-to-end latency when
5//! it re-reads the embedded `send_unix_nanos`. The remaining bytes are a
6//! deterministic filler, so the wire size is exactly the scenario's message
7//! size.
8//!
9//! The header is 24 bytes and not the 16 the plan sketched, because it needs a
10//! magic to detect "this is one of ours". Kafka's own producers leave their own
11//! headers in there, and this driver must not misread their bytes.
12
13use {
14    crate::numeric::saturating_u128_to_u64,
15    bytes::{Bytes, BytesMut},
16    krabka_units::prelude::*,
17    std::time::{SystemTime, UNIX_EPOCH},
18};
19
20/// Magic prefix on every record, so a consumer can confirm that this driver
21/// produced the record and that it is not data already in the topic.
22pub const MAGIC: [u8; 8] = *b"KRABKA_B";
23pub const HEADER_LEN: usize = MAGIC.len() + 8 + 8; // magic + scenario_id + send_nanos = 24
24
25/// Builds a reusable filler template of exactly `msg_size` bytes, or of the
26/// header length if that is larger.
27///
28/// The first 24 bytes are zero, and `stamp_into` overwrites them at send time.
29/// The remaining bytes are a repeating pattern.
30#[must_use]
31pub fn template(msg_size: ByteSize) -> BytesMut {
32    let len = msg_size.bytes_usize().max(HEADER_LEN);
33    let mut b = BytesMut::with_capacity(len);
34    b.resize(len, 0u8);
35    // Fill the body with a repeating ramp so compression has *some* work.
36    // All-zeros compresses too well; all-random compresses too poorly.
37    for (i, byte) in b.iter_mut().enumerate().skip(HEADER_LEN) {
38        *byte = u8::try_from(i & 0xff).unwrap_or_default();
39    }
40    b
41}
42
43/// Builds a record key of exactly `key_size` bytes, distinct per `sequence`.
44///
45/// A scenario that sets `key_size` asks for a keyed workload, and a keyed
46/// workload is not a keyless one: the partitioner hashes the key instead of
47/// round-robining, and the key adds its own bytes to the wire. The key has to
48/// vary, or every record hashes to one partition and the run measures a single
49/// partition rather than the scenario's partition count.
50///
51/// The first eight bytes are the big-endian sequence, so keys are distinct and
52/// spread. The rest is the same repeating ramp the value filler uses, so a
53/// compressed batch is not dominated by a run of zeros. A `key_size` below
54/// eight bytes is honoured as given and takes a prefix of the sequence.
55#[must_use]
56pub fn key(key_size: ByteSize, sequence: u64) -> Option<Bytes> {
57    let len = key_size.bytes_usize();
58    if len == 0 {
59        return None;
60    }
61    let mut buf = BytesMut::with_capacity(len);
62    buf.resize(len, 0u8);
63    let stamped = len.min(8);
64    for (slot, byte) in buf.iter_mut().zip(sequence.to_be_bytes()) {
65        *slot = byte;
66    }
67    for (index, byte) in buf.iter_mut().enumerate().skip(stamped) {
68        *byte = u8::try_from(index & 0xff).unwrap_or_default();
69    }
70    Some(buf.freeze())
71}
72
73/// Stamps the magic, the `scenario_id`, and the current `unix_nanos` into the
74/// first 24 bytes of `buf`. Returns the value as a `Bytes`, which the caller
75/// clones cheaply with a `BytesMut::freeze`-style copy.
76pub fn stamp_into(buf: &mut BytesMut, scenario_id: u64) -> Bytes {
77    assert2::debug_assert!(buf.len() >= HEADER_LEN, "buf too short for header");
78    let nanos = SystemTime::now()
79        .duration_since(UNIX_EPOCH)
80        .map_or(0, |d| saturating_u128_to_u64(d.as_nanos()));
81    let header = MAGIC
82        .into_iter()
83        .chain(scenario_id.to_be_bytes())
84        .chain(nanos.to_be_bytes());
85    for (slot, byte) in buf.iter_mut().zip(header) {
86        *slot = byte;
87    }
88    // Freeze a copy into a Bytes the producer can keep.
89    Bytes::copy_from_slice(buf)
90}
91
92/// Reads the embedded `send_unix_nanos` if the record is one of ours.
93/// Returns `None` if the magic does not match, and also if the record is too
94/// short. The consumer skips both cases silently.
95#[must_use]
96pub fn read_send_nanos(value: &[u8], scenario_id: u64) -> Option<u64> {
97    let (magic, rest) = value.split_first_chunk::<8>()?;
98    let (sid, rest) = rest.split_first_chunk::<8>()?;
99    let (nanos, _) = rest.split_first_chunk::<8>()?;
100    if *magic != MAGIC || u64::from_be_bytes(*sid) != scenario_id {
101        return None;
102    }
103    Some(u64::from_be_bytes(*nanos))
104}
105
106#[cfg(test)]
107mod tests {
108
109    use super::*;
110
111    #[test]
112    fn a_zero_key_size_means_a_keyless_record() {
113        assert2::assert!(key(ByteSize::ZERO, 7).is_none());
114    }
115
116    #[test]
117    fn a_key_has_the_configured_size_and_varies_per_record() {
118        let first = key(bytes(16), 1).expect("a sized key");
119        let second = key(bytes(16), 2).expect("a sized key");
120        assert2::assert!(first.len() == 16);
121        assert2::assert!(second.len() == 16);
122        assert2::assert!(first != second);
123        assert2::assert!(first[..8] == 1u64.to_be_bytes());
124    }
125
126    #[test]
127    fn a_key_shorter_than_the_sequence_is_still_the_configured_size() {
128        let short = key(bytes(4), 0x0102_0304_0506_0708).expect("a sized key");
129        assert2::assert!(short.len() == 4);
130        assert2::assert!(short[..] == [0x01, 0x02, 0x03, 0x04]);
131    }
132
133    #[test]
134    fn round_trip_send_nanos() {
135        let mut t = template(bytes(64));
136        let b = stamp_into(&mut t, 0xdead_beef);
137        let n = read_send_nanos(&b, 0xdead_beef).expect("magic+sid match");
138        assert2::assert!(n > 0);
139    }
140
141    #[test]
142    fn rejects_wrong_scenario_id() {
143        let mut t = template(bytes(64));
144        let b = stamp_into(&mut t, 42);
145        assert2::assert!(read_send_nanos(&b, 7).is_none());
146    }
147
148    #[test]
149    fn rejects_short() {
150        assert2::assert!(read_send_nanos(&[0u8; 8], 0).is_none());
151    }
152
153    #[test]
154    fn rejects_bad_magic() {
155        let mut b = vec![0u8; HEADER_LEN];
156        b[16..24].copy_from_slice(&123u64.to_be_bytes());
157        assert2::assert!(read_send_nanos(&b, 0).is_none());
158    }
159
160    #[test]
161    fn template_size_honoured_above_header() {
162        assert2::assert!(template(kibibytes(1)).len() == 1024);
163        assert2::assert!(template(bytes(512)).len() == 512);
164    }
165
166    #[test]
167    fn template_min_size_is_header() {
168        assert2::assert!(template(ByteSize::ZERO).len() == HEADER_LEN);
169    }
170}
krabka_bench_driver/payload.rs

krabka_bench_driver/
payload.rs
