use std::borrow::Cow;
use std::collections::BTreeSet;
use bytes::Bytes;
use crate::storage_adapter::{StorageSpace, StorageSpaceId, ValueSemantics};
use crate::tracked_state::codec::decode_value;
use crate::tracked_state::types::TrackedStateIndexValue;
use crate::{LixError, storage_codec};
pub(crate) const CURRENT_STATE_DATA_PART_SPACE: StorageSpace = StorageSpace::declare(
StorageSpaceId(0x0004_002f),
"tracked_state.current_state_data_part.v1",
ValueSemantics::Immutable,
);
pub(crate) const CURRENT_STATE_DATA_PART_MAX_ROWS: usize = 512;
pub(crate) const CURRENT_STATE_DATA_PART_TARGET_BYTES: usize = 64 * 1024;
const CURRENT_STATE_DATA_PART_MAX_BYTES: usize = 4 * 1024 * 1024;
const CURRENT_STATE_DATA_PART_MAX_DECODED_BYTES: usize = 16 * 1024 * 1024;
const RAW_MAGIC: &[u8; 7] = b"LXCSP03";
const ZSTD_MAGIC: &[u8; 7] = b"LXCSPZ3";
const DIGEST_CONTEXT: &str = "lix native current-state data part v3";
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct CurrentStateDataRow {
pub(crate) encoded_key: Vec<u8>,
pub(crate) value: TrackedStateIndexValue,
pub(crate) metadata: Option<lix_schema::Jsonb>,
pub(crate) snapshot: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, musli::Encode, musli::Decode)]
#[musli(packed)]
struct StoredCurrentStateDataRow {
#[musli(bytes)]
encoded_key: Vec<u8>,
#[musli(bytes)]
encoded_value: Vec<u8>,
#[musli(with = crate::changelog::jsonb_option_storage)]
metadata: Option<lix_schema::Jsonb>,
#[musli(bytes)]
snapshot: Vec<u8>,
}
#[derive(Debug, Clone)]
pub(crate) struct EncodedCurrentStateDataPart {
pub(crate) digest: [u8; 32],
pub(crate) bytes: Bytes,
pub(crate) first_key: Vec<u8>,
pub(crate) last_key: Vec<u8>,
pub(crate) row_count: u16,
}
pub(crate) fn encode_current_state_data_part(
rows: &[CurrentStateDataRow],
compressor: &mut Option<crate::compression::ZstdLevel1Compressor>,
) -> Result<EncodedCurrentStateDataPart, LixError> {
validate_rows(rows)?;
let stored = rows
.iter()
.map(|row| StoredCurrentStateDataRow {
encoded_key: row.encoded_key.clone(),
encoded_value: super::codec::encode_value_ref(
super::types::TrackedStateIndexValueRef {
change_id: row.value.change_id,
commit_id: row.value.commit_id,
deleted: row.value.deleted,
created_at: row.value.created_at,
updated_at: row.value.updated_at,
},
),
metadata: row.metadata.clone(),
snapshot: row.snapshot.clone(),
})
.collect::<Vec<_>>();
let payload = storage_codec::encode("native current-state data part", &stored)?;
if payload.len() > CURRENT_STATE_DATA_PART_MAX_DECODED_BYTES {
return Err(part_error("decoded payload exceeds its bound"));
}
let mut encoded = Vec::with_capacity(RAW_MAGIC.len() + payload.len());
encoded.extend_from_slice(RAW_MAGIC);
encoded.extend_from_slice(&payload);
if payload.len() >= 512 {
if compressor.is_none() {
*compressor = Some(
crate::compression::ZstdLevel1Compressor::new()
.map_err(|error| part_error(format!("compressor init failed: {error}")))?,
);
}
let compressed = compressor
.as_mut()
.expect("compressor was initialized")
.compress(&payload)
.map_err(|error| part_error(format!("compression failed: {error}")))?;
let mut physical = Vec::with_capacity(ZSTD_MAGIC.len() + 4 + compressed.len());
physical.extend_from_slice(ZSTD_MAGIC);
physical.extend_from_slice(
&u32::try_from(payload.len())
.map_err(|_| part_error("decoded length exceeds u32"))?
.to_be_bytes(),
);
physical.extend_from_slice(&compressed);
if physical.len() < encoded.len() {
encoded = physical;
}
}
if encoded.len() > CURRENT_STATE_DATA_PART_MAX_BYTES {
return Err(part_error("physical payload exceeds its bound"));
}
let digest = digest(&encoded);
Ok(EncodedCurrentStateDataPart {
digest,
bytes: Bytes::from(encoded),
first_key: rows.first().expect("validated rows").encoded_key.clone(),
last_key: rows.last().expect("validated rows").encoded_key.clone(),
row_count: u16::try_from(rows.len()).expect("native part row count is bounded"),
})
}
pub(crate) fn encode_bounded_current_state_data_parts(
rows: &[CurrentStateDataRow],
) -> Result<Vec<EncodedCurrentStateDataPart>, LixError> {
let mut compressor = None;
let mut encoded = Vec::new();
let mut offset = 0usize;
while offset < rows.len() {
let mut count = (rows.len() - offset).min(CURRENT_STATE_DATA_PART_MAX_ROWS);
let part = loop {
let part =
encode_current_state_data_part(&rows[offset..offset + count], &mut compressor)?;
if part.bytes.len() <= CURRENT_STATE_DATA_PART_TARGET_BYTES || count == 1 {
break part;
}
count = count.div_ceil(2);
};
encoded.push(part);
offset += count;
}
Ok(encoded)
}
pub(crate) fn decode_current_state_data_part(
expected_digest: &[u8; 32],
encoded: &[u8],
) -> Result<Vec<CurrentStateDataRow>, LixError> {
if encoded.len() > CURRENT_STATE_DATA_PART_MAX_BYTES || &digest(encoded) != expected_digest {
return Err(part_error("content digest or physical bound is invalid"));
}
let payload: Cow<'_, [u8]> = if let Some(payload) = encoded.strip_prefix(RAW_MAGIC) {
Cow::Borrowed(payload)
} else if let Some(body) = encoded.strip_prefix(ZSTD_MAGIC) {
let (decoded_len, compressed) = body
.split_at_checked(4)
.ok_or_else(|| part_error("compressed payload is truncated"))?;
let decoded_len = usize::try_from(u32::from_be_bytes(
decoded_len.try_into().expect("fixed decoded length"),
))
.expect("u32 fits usize");
if decoded_len > CURRENT_STATE_DATA_PART_MAX_DECODED_BYTES {
return Err(part_error("compressed payload exceeds its decoded bound"));
}
Cow::Owned(
crate::compression::decompress_zstd(compressed, decoded_len)
.map_err(|error| part_error(format!("decompression failed: {error}")))?,
)
} else {
return Err(part_error("unsupported format; recreate the repository"));
};
let stored: Vec<StoredCurrentStateDataRow> =
storage_codec::decode("native current-state data part", &payload)?;
let rows = stored
.into_iter()
.map(|row| {
Ok(CurrentStateDataRow {
encoded_key: row.encoded_key,
value: decode_value(&row.encoded_value)?,
metadata: row.metadata,
snapshot: row.snapshot,
})
})
.collect::<Result<Vec<_>, LixError>>()?;
validate_rows(&rows)?;
Ok(rows)
}
pub(crate) fn decode_current_state_data_part_commit_ids(
expected_digest: &[u8; 32],
encoded: &[u8],
) -> Result<BTreeSet<crate::changelog::CommitId>, LixError> {
Ok(decode_current_state_data_part(expected_digest, encoded)?
.into_iter()
.map(|row| row.value.commit_id)
.collect())
}
fn validate_rows(rows: &[CurrentStateDataRow]) -> Result<(), LixError> {
if rows.is_empty()
|| rows.len() > CURRENT_STATE_DATA_PART_MAX_ROWS
|| rows
.iter()
.any(|row| row.encoded_key.is_empty() || row.value.deleted || row.snapshot.is_empty())
|| rows
.windows(2)
.any(|pair| pair[0].encoded_key >= pair[1].encoded_key)
{
return Err(part_error(
"rows are empty, deleted, oversized, or unordered",
));
}
Ok(())
}
fn digest(encoded: &[u8]) -> [u8; 32] {
*blake3::Hasher::new_derive_key(DIGEST_CONTEXT)
.update(encoded)
.finalize()
.as_bytes()
}
fn part_error(message: impl std::fmt::Display) -> LixError {
LixError::new(
LixError::CODE_INTERNAL_ERROR,
format!("tracked_state native current-state data part {message}"),
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::changelog::{ChangeId, CommitId};
use crate::common::LixTimestamp;
fn row(index: usize) -> CurrentStateDataRow {
CurrentStateDataRow {
encoded_key: format!("key-{index:04}").into_bytes(),
value: TrackedStateIndexValue {
change_id: ChangeId::for_test_label(&format!("native-change-{index}")),
commit_id: CommitId::for_test_label(&format!("native-commit-{index}")),
deleted: false,
created_at: LixTimestamp::from_unix_millis_utc_lossy(index as i64),
updated_at: LixTimestamp::from_unix_millis_utc_lossy(index as i64 + 1),
},
metadata: None,
snapshot: format!("typed-{index}").into_bytes(),
}
}
#[test]
fn native_parts_round_trip_exact_provenance_and_reject_corruption() {
let mut rows = (0..513).map(row).collect::<Vec<_>>();
rows[0].metadata = Some(lix_schema::Jsonb::from_value(
serde_json::json!({"source": "native"}),
));
let parts = encode_bounded_current_state_data_parts(&rows).expect("parts should encode");
assert_eq!(parts.len(), 2, "row bound must split the post-image");
let decoded = parts
.iter()
.flat_map(|part| {
decode_current_state_data_part(&part.digest, &part.bytes)
.expect("part should decode")
})
.collect::<Vec<_>>();
assert_eq!(decoded, rows);
let mut corrupt = parts[0].bytes.to_vec();
*corrupt.last_mut().expect("encoded part is non-empty") ^= 1;
assert!(decode_current_state_data_part(&parts[0].digest, &corrupt).is_err());
}
#[test]
fn native_parts_reject_tombstones_and_unordered_rows() {
let mut rows = vec![row(1), row(0)];
assert!(encode_bounded_current_state_data_parts(&rows).is_err());
rows.sort_by(|left, right| left.encoded_key.cmp(&right.encoded_key));
rows[0].value.deleted = true;
assert!(encode_bounded_current_state_data_parts(&rows).is_err());
}
#[test]
fn native_parts_round_trip_snapshot_without_json() {
let typed = crate::row_payload::TypedRow {
schema_fingerprint: [8; 32],
row_pk: vec![lix_schema::Value::Text("typed-row".to_owned())].into(),
row: lix_schema::Row::from([("value".to_owned(), lix_schema::Value::Int8(7))]),
native_payload: std::sync::OnceLock::new(),
boundary_create_validation: std::sync::OnceLock::new(),
};
let row = CurrentStateDataRow {
encoded_key: b"typed-row".to_vec(),
value: TrackedStateIndexValue {
change_id: ChangeId::for_test_label("typed-change"),
commit_id: CommitId::for_test_label("typed-commit"),
deleted: false,
created_at: LixTimestamp::from_unix_millis_utc_lossy(1),
updated_at: LixTimestamp::from_unix_millis_utc_lossy(2),
},
metadata: None,
snapshot: crate::hot_state::encode_snapshot(&typed).expect("typed payload encodes"),
};
let part = encode_bounded_current_state_data_parts(&[row.clone()])
.expect("typed current-state part should encode")
.pop()
.expect("one part");
let decoded = decode_current_state_data_part(&part.digest, &part.bytes)
.expect("typed current-state part should decode");
assert_eq!(decoded[0], row);
assert_eq!(
crate::row_payload::TypedRow::decode_durable_payload(
decoded[0].snapshot.clone().into(),
"schema",
&crate::row_pk::RowPk::single("typed-row"),
)
.expect("typed payload decodes against its current-state envelope"),
typed
);
}
}