reddb_file/
append_only_segment.rs

1use serde::{Deserialize, Serialize};
2
3use crate::{RdbFileError, RdbFileResult};
4
5pub const APPEND_ONLY_SEGMENT_VERSION: u32 = 1;
6pub const APPEND_ONLY_SEGMENT_CHUNK_BYTES: u32 = 512 * 1024;
7
8const MAGIC: &[u8; 8] = b"RDBAOS01";
9const FIXED_HEADER_LEN: usize = 8 + 4 + 8 + 4 + 4;
10
11#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
12#[serde(rename_all = "kebab-case")]
13pub enum AppendOnlySegmentCodec {
14    #[default]
15    Zstd,
16    None,
17}
18
19impl AppendOnlySegmentCodec {
20    pub fn as_str(self) -> &'static str {
21        match self {
22            Self::Zstd => "zstd",
23            Self::None => "none",
24        }
25    }
26}
27
28#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29pub struct AppendOnlySegmentRow {
30    pub primary_key: Vec<u8>,
31    pub payload: Vec<u8>,
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
35pub struct AppendOnlySegmentChunkChecksum {
36    pub offset: u64,
37    pub len: u32,
38    pub checksum: String,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
42pub struct AppendOnlySegmentBloom {
43    pub num_hashes: u8,
44    pub bit_size: u32,
45    pub inserted: u32,
46    pub bits_hex: String,
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub struct AppendOnlySegment {
51    pub version: u32,
52    pub codec: AppendOnlySegmentCodec,
53    pub chunk_size: u32,
54    pub rows: Vec<AppendOnlySegmentRow>,
55    pub primary_min: Option<Vec<u8>>,
56    pub primary_max: Option<Vec<u8>>,
57    pub primary_bloom: Option<AppendOnlySegmentBloom>,
58    pub chunk_checksums: Vec<AppendOnlySegmentChunkChecksum>,
59}
60
61#[derive(Debug, Serialize, Deserialize)]
62struct SegmentHeader {
63    version: u32,
64    codec: AppendOnlySegmentCodec,
65    chunk_size: u32,
66    row_count: u64,
67    uncompressed_len: u64,
68    primary_min_hex: Option<String>,
69    primary_max_hex: Option<String>,
70    primary_bloom: Option<AppendOnlySegmentBloom>,
71    chunk_checksums: Vec<AppendOnlySegmentChunkChecksum>,
72}
73
74#[derive(Debug, Serialize, Deserialize)]
75struct RowJson {
76    primary_key_hex: String,
77    payload_hex: String,
78}
79
80pub fn encode_append_only_segment(
81    codec: AppendOnlySegmentCodec,
82    rows: &[AppendOnlySegmentRow],
83) -> RdbFileResult<Vec<u8>> {
84    let row_json = rows
85        .iter()
86        .map(|row| RowJson {
87            primary_key_hex: hex::encode(&row.primary_key),
88            payload_hex: hex::encode(&row.payload),
89        })
90        .collect::<Vec<_>>();
91    let plain = serde_json::to_vec(&row_json).map_err(invalid_operation)?;
92    let body = match codec {
93        AppendOnlySegmentCodec::Zstd => {
94            zstd::bulk::compress(&plain, 3).map_err(invalid_operation)?
95        }
96        AppendOnlySegmentCodec::None => plain.clone(),
97    };
98    let chunk_checksums = chunk_checksums(&body);
99    let header = SegmentHeader {
100        version: APPEND_ONLY_SEGMENT_VERSION,
101        codec,
102        chunk_size: APPEND_ONLY_SEGMENT_CHUNK_BYTES,
103        row_count: rows.len() as u64,
104        uncompressed_len: plain.len() as u64,
105        primary_min_hex: rows
106            .iter()
107            .map(|row| row.primary_key.as_slice())
108            .min()
109            .map(hex::encode),
110        primary_max_hex: rows
111            .iter()
112            .map(|row| row.primary_key.as_slice())
113            .max()
114            .map(hex::encode),
115        primary_bloom: build_primary_bloom(rows),
116        chunk_checksums,
117    };
118    let header_bytes = serde_json::to_vec(&header).map_err(invalid_operation)?;
119    let header_len = u32::try_from(header_bytes.len()).map_err(|_| {
120        RdbFileError::InvalidOperation("append-only segment header too large".into())
121    })?;
122    let body_len = u64::try_from(body.len())
123        .map_err(|_| RdbFileError::InvalidOperation("append-only segment body too large".into()))?;
124
125    let mut out = Vec::with_capacity(FIXED_HEADER_LEN + header_bytes.len() + body.len());
126    out.extend_from_slice(MAGIC);
127    out.extend_from_slice(&header_len.to_le_bytes());
128    out.extend_from_slice(&body_len.to_le_bytes());
129    out.extend_from_slice(&crc32(&header_bytes).to_le_bytes());
130    out.extend_from_slice(&crc32(&body).to_le_bytes());
131    out.extend_from_slice(&header_bytes);
132    out.extend_from_slice(&body);
133    Ok(out)
134}
135
136pub fn decode_append_only_segment(bytes: &[u8]) -> RdbFileResult<AppendOnlySegment> {
137    if bytes.len() < FIXED_HEADER_LEN {
138        return Err(RdbFileError::InvalidOperation(
139            "append-only segment frame too short".into(),
140        ));
141    }
142    if &bytes[..8] != MAGIC {
143        return Err(RdbFileError::InvalidOperation(
144            "invalid append-only segment magic".into(),
145        ));
146    }
147    let header_len = read_u32(bytes, 8)? as usize;
148    let body_len = read_u64(bytes, 12)? as usize;
149    let header_crc = read_u32(bytes, 20)?;
150    let body_crc = read_u32(bytes, 24)?;
151    let header_start = FIXED_HEADER_LEN;
152    let body_start = header_start.checked_add(header_len).ok_or_else(|| {
153        RdbFileError::InvalidOperation("append-only segment size overflow".into())
154    })?;
155    let end = body_start.checked_add(body_len).ok_or_else(|| {
156        RdbFileError::InvalidOperation("append-only segment size overflow".into())
157    })?;
158    if bytes.len() != end {
159        return Err(RdbFileError::InvalidOperation(
160            "append-only segment frame length mismatch".into(),
161        ));
162    }
163    let header_bytes = &bytes[header_start..body_start];
164    let body = &bytes[body_start..end];
165    if crc32(header_bytes) != header_crc {
166        return Err(RdbFileError::InvalidOperation(
167            "append-only segment header checksum mismatch".into(),
168        ));
169    }
170    if crc32(body) != body_crc {
171        return Err(RdbFileError::InvalidOperation(
172            "append-only segment body checksum mismatch".into(),
173        ));
174    }
175    let header: SegmentHeader = serde_json::from_slice(header_bytes).map_err(invalid_operation)?;
176    if header.version != APPEND_ONLY_SEGMENT_VERSION {
177        return Err(RdbFileError::InvalidOperation(format!(
178            "unsupported append-only segment version: {}",
179            header.version
180        )));
181    }
182    if header.chunk_size != APPEND_ONLY_SEGMENT_CHUNK_BYTES {
183        return Err(RdbFileError::InvalidOperation(
184            "append-only segment chunk size mismatch".into(),
185        ));
186    }
187    if header.chunk_checksums != chunk_checksums(body) {
188        return Err(RdbFileError::InvalidOperation(
189            "append-only segment chunk checksum mismatch".into(),
190        ));
191    }
192
193    let plain = match header.codec {
194        AppendOnlySegmentCodec::Zstd => {
195            let max_len = usize::try_from(header.uncompressed_len).map_err(|_| {
196                RdbFileError::InvalidOperation("append-only segment plain length too large".into())
197            })?;
198            zstd::bulk::decompress(body, max_len).map_err(invalid_operation)?
199        }
200        AppendOnlySegmentCodec::None => body.to_vec(),
201    };
202    if plain.len() as u64 != header.uncompressed_len {
203        return Err(RdbFileError::InvalidOperation(
204            "append-only segment plain length mismatch".into(),
205        ));
206    }
207    let row_json: Vec<RowJson> = serde_json::from_slice(&plain).map_err(invalid_operation)?;
208    if row_json.len() as u64 != header.row_count {
209        return Err(RdbFileError::InvalidOperation(
210            "append-only segment row count mismatch".into(),
211        ));
212    }
213    let rows = row_json
214        .into_iter()
215        .map(|row| {
216            Ok(AppendOnlySegmentRow {
217                primary_key: hex::decode(row.primary_key_hex).map_err(invalid_operation)?,
218                payload: hex::decode(row.payload_hex).map_err(invalid_operation)?,
219            })
220        })
221        .collect::<RdbFileResult<Vec<_>>>()?;
222    let primary_min = decode_optional_hex(header.primary_min_hex)?;
223    let primary_max = decode_optional_hex(header.primary_max_hex)?;
224    Ok(AppendOnlySegment {
225        version: header.version,
226        codec: header.codec,
227        chunk_size: header.chunk_size,
228        rows,
229        primary_min,
230        primary_max,
231        primary_bloom: header.primary_bloom,
232        chunk_checksums: header.chunk_checksums,
233    })
234}
235
236pub fn append_only_segment_chunk_checksums(bytes: &[u8]) -> Vec<AppendOnlySegmentChunkChecksum> {
237    chunk_checksums(bytes)
238}
239
240pub fn append_only_segment_primary_bloom_might_contain(
241    bloom: &AppendOnlySegmentBloom,
242    key: &[u8],
243) -> bool {
244    if bloom.bit_size == 0 || bloom.num_hashes == 0 {
245        return true;
246    }
247    let Ok(bits) = hex::decode(&bloom.bits_hex) else {
248        return true;
249    };
250    for idx in primary_bloom_indexes(key, bloom.bit_size, bloom.num_hashes) {
251        let byte = (idx / 8) as usize;
252        let mask = 1u8 << (idx % 8);
253        if bits
254            .get(byte)
255            .map(|value| value & mask == 0)
256            .unwrap_or(true)
257        {
258            return false;
259        }
260    }
261    true
262}
263
264fn build_primary_bloom(rows: &[AppendOnlySegmentRow]) -> Option<AppendOnlySegmentBloom> {
265    if rows.is_empty() {
266        return None;
267    }
268    let inserted = u32::try_from(rows.len()).ok()?;
269    let bit_size = inserted.saturating_mul(10).max(64);
270    let num_hashes = 3;
271    let mut bits = vec![0u8; (bit_size as usize).div_ceil(8)];
272    for row in rows {
273        for idx in primary_bloom_indexes(&row.primary_key, bit_size, num_hashes) {
274            bits[(idx / 8) as usize] |= 1u8 << (idx % 8);
275        }
276    }
277    Some(AppendOnlySegmentBloom {
278        num_hashes,
279        bit_size,
280        inserted,
281        bits_hex: hex::encode(bits),
282    })
283}
284
285fn primary_bloom_indexes(key: &[u8], bit_size: u32, num_hashes: u8) -> impl Iterator<Item = u32> {
286    let h1 = fnv1a64(key);
287    let h2 = djb2_64(key).max(1);
288    (0..num_hashes)
289        .map(move |idx| h1.wrapping_add(u64::from(idx).wrapping_mul(h2)) as u32 % bit_size)
290}
291
292fn fnv1a64(bytes: &[u8]) -> u64 {
293    let mut hash = 0xcbf2_9ce4_8422_2325u64;
294    for byte in bytes {
295        hash ^= u64::from(*byte);
296        hash = hash.wrapping_mul(0x100_0000_01b3);
297    }
298    hash
299}
300
301fn djb2_64(bytes: &[u8]) -> u64 {
302    let mut hash = 5381u64;
303    for byte in bytes {
304        hash = hash.wrapping_mul(33).wrapping_add(u64::from(*byte));
305    }
306    hash
307}
308
309fn chunk_checksums(bytes: &[u8]) -> Vec<AppendOnlySegmentChunkChecksum> {
310    bytes
311        .chunks(APPEND_ONLY_SEGMENT_CHUNK_BYTES as usize)
312        .enumerate()
313        .map(|(idx, chunk)| AppendOnlySegmentChunkChecksum {
314            offset: (idx as u64) * u64::from(APPEND_ONLY_SEGMENT_CHUNK_BYTES),
315            len: chunk.len() as u32,
316            checksum: format!("crc32:{:08x}", crc32(chunk)),
317        })
318        .collect()
319}
320
321fn read_u32(bytes: &[u8], offset: usize) -> RdbFileResult<u32> {
322    let slice = bytes
323        .get(offset..offset + 4)
324        .ok_or_else(|| RdbFileError::InvalidOperation("append-only segment u32 missing".into()))?;
325    Ok(u32::from_le_bytes(slice.try_into().expect("slice length")))
326}
327
328fn read_u64(bytes: &[u8], offset: usize) -> RdbFileResult<u64> {
329    let slice = bytes
330        .get(offset..offset + 8)
331        .ok_or_else(|| RdbFileError::InvalidOperation("append-only segment u64 missing".into()))?;
332    Ok(u64::from_le_bytes(slice.try_into().expect("slice length")))
333}
334
335fn decode_optional_hex(value: Option<String>) -> RdbFileResult<Option<Vec<u8>>> {
336    value
337        .map(hex::decode)
338        .transpose()
339        .map_err(invalid_operation)
340}
341
342fn crc32(data: &[u8]) -> u32 {
343    let mut hasher = crc32fast::Hasher::new();
344    hasher.update(data);
345    hasher.finalize()
346}
347
348fn invalid_operation(err: impl ToString) -> RdbFileError {
349    RdbFileError::InvalidOperation(err.to_string())
350}
reddb_file/append_only_segment.rs

reddb_file/
append_only_segment.rs
