use crate::meta::{ColumnId, PredicateValue};
use super::{DocId, DocumentVersion, IngestDocument, Revision};
pub const UPSERT_V1: u16 = 1;
pub const DELETE_V1: u16 = 2;
pub const METADATA_EDIT_V1: u16 = 3;
pub const UPSERT_WITH_TIMESTAMP_V1: u16 = 4;
pub const UPSERT_WITH_METADATA_V1: u16 = 5;
pub const UPSERT_WITH_TIMESTAMP_AND_METADATA_V1: u16 = 6;
pub const UPSERT_V2: u16 = 7;
pub const UPSERT_V2_BATCH_MEMBER: u16 = 8;
pub const GRAPH_COMMIT_V1: u16 = 10;
pub const MIXED_BATCH_MEMBER_V1: u16 = 11;
pub const MIXED_BATCH_MEMBER_FRAMED_OVERHEAD: usize = 10 + crate::wal::record::MIN_RECORD_LEN;
pub fn mixed_batch_group_bytes(
header_bytes: usize,
inner_payload_bytes: &[usize],
max_group_bytes: usize,
) -> Result<usize, crate::wal::WalWriteError> {
let encoded_bytes = inner_payload_bytes
.iter()
.try_fold(header_bytes, |sum, length| {
sum.checked_add(*length)?
.checked_add(MIXED_BATCH_MEMBER_FRAMED_OVERHEAD)
});
match encoded_bytes {
Some(encoded_bytes) if encoded_bytes <= max_group_bytes => Ok(encoded_bytes),
encoded_bytes => Err(crate::wal::WalWriteError::GroupTooLarge {
encoded_bytes: encoded_bytes.unwrap_or(usize::MAX),
max_group_bytes,
}),
}
}
pub const UPSERT_V2_VECTOR: u32 = 1 << 0;
pub const UPSERT_V2_TEXT: u32 = 1 << 1;
pub const UPSERT_V2_TIMESTAMP: u32 = 1 << 2;
pub const UPSERT_V2_METADATA: u32 = 1 << 3;
pub const UPSERT_V2_TYPED_COLUMNS: u32 = 1 << 4;
pub const UPSERT_V2_EPOCH_TAG_RESERVED: u32 = 1 << 5;
const UPSERT_V2_KNOWN_FIELDS: u32 = UPSERT_V2_VECTOR
| UPSERT_V2_TEXT
| UPSERT_V2_TIMESTAMP
| UPSERT_V2_METADATA
| UPSERT_V2_TYPED_COLUMNS;
const PAYLOAD_VERSION: u16 = 1;
const COMMON_HEADER_LEN: usize = 4;
const DOCUMENT_VERSION_LEN: usize = 24;
const UPSERT_PREFIX_LEN: usize = COMMON_HEADER_LEN + DOCUMENT_VERSION_LEN + 4;
const TIMESTAMPED_UPSERT_PREFIX_LEN: usize = COMMON_HEADER_LEN + DOCUMENT_VERSION_LEN + 8 + 4;
const METADATA_UPSERT_PREFIX_LEN: usize = COMMON_HEADER_LEN + DOCUMENT_VERSION_LEN + 4 + 4;
const TIMESTAMPED_METADATA_UPSERT_PREFIX_LEN: usize =
COMMON_HEADER_LEN + DOCUMENT_VERSION_LEN + 8 + 4 + 4;
const DELETE_PREFIX_LEN: usize = COMMON_HEADER_LEN + 4;
const METADATA_PREFIX_LEN: usize = COMMON_HEADER_LEN + DOCUMENT_VERSION_LEN + 4 + 1 + 3 + 4;
const VALUE_NULL: u8 = 0;
const VALUE_ID128: u8 = 6;
const VALUE_U64: u8 = 1;
const VALUE_I64: u8 = 2;
const VALUE_F64: u8 = 3;
const VALUE_BOOL: u8 = 4;
const VALUE_STRING: u8 = 5;
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PayloadError {
#[cfg(feature = "graph-cypher")]
GraphEnvelope(crate::property_graph::wal::WalError),
UnknownOperation(u16),
TransactionBinding,
LengthOverflow,
Truncated,
Version(u16),
Flags(u16),
Reserved(u32),
EmptyVector,
EmptyDelete,
NonFiniteVector {
index: usize,
},
ValueKind(u8),
ValueLength {
kind: u8,
expected: usize,
actual: usize,
},
Boolean(u8),
Utf8,
TrailingBytes(usize),
FieldBitmap(u32),
Lexical(String),
Columns(String),
BatchPosition {
index: u32,
count: u32,
},
OrphanBatchMember {
index: u32,
count: u32,
},
}
impl std::fmt::Display for PayloadError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
#[cfg(feature = "graph-cypher")]
Self::GraphEnvelope(error) => write!(formatter, "WAL graph commit: {error}"),
Self::TransactionBinding => write!(formatter, "invalid transaction binding"),
Self::UnknownOperation(op) => {
write!(formatter, "WAL mutation operation {op} is unknown")
}
Self::LengthOverflow => formatter.write_str("WAL mutation payload length overflow"),
Self::Truncated => formatter.write_str("WAL mutation payload is truncated"),
Self::Version(version) => write!(
formatter,
"WAL mutation payload version {version} is unsupported"
),
Self::Flags(flags) => write!(
formatter,
"WAL mutation payload flags {flags:#06x} are non-zero"
),
Self::Reserved(value) => write!(
formatter,
"WAL metadata payload reserved bytes are {value:#08x}, expected zero"
),
Self::EmptyVector => formatter.write_str("WAL upsert vector is empty"),
Self::EmptyDelete => formatter.write_str("WAL delete document list is empty"),
Self::NonFiniteVector { index } => write!(
formatter,
"WAL upsert vector coordinate {index} is non-finite"
),
Self::ValueKind(kind) => {
write!(formatter, "WAL metadata value kind {kind} is unknown")
}
Self::ValueLength {
kind,
expected,
actual,
} => write!(
formatter,
"WAL metadata value kind {kind} needs {expected} bytes, got {actual}"
),
Self::Boolean(value) => write!(
formatter,
"WAL metadata Boolean byte {value} is not canonical zero or one"
),
Self::Utf8 => formatter.write_str("WAL metadata string is not valid UTF-8"),
Self::TrailingBytes(bytes) => {
write!(formatter, "WAL mutation payload has {bytes} trailing bytes")
}
Self::FieldBitmap(bitmap) => write!(
formatter,
"WAL upsert-v2 field bitmap {bitmap:#010x} names reserved fields"
),
Self::Lexical(detail) => write!(formatter, "WAL lexical replay failed: {detail}"),
Self::Columns(detail) => write!(formatter, "WAL column replay failed: {detail}"),
Self::BatchPosition { index, count } => write!(
formatter,
"WAL batch member {index} of {count} is not a valid position"
),
Self::OrphanBatchMember { index, count } => write!(
formatter,
"WAL batch member {index} of {count} continues no earlier member"
),
}
}
}
impl std::error::Error for PayloadError {}
#[derive(Clone, Debug, PartialEq)]
pub enum MetadataValue {
Null,
U64(u64),
Id128(DocId),
I64(i64),
F64(f64),
Bool(bool),
String(String),
}
#[derive(Clone, Debug, PartialEq)]
pub struct MetadataEdit {
version: DocumentVersion,
column: ColumnId,
value: MetadataValue,
}
impl MetadataEdit {
#[must_use]
pub const fn new(version: DocumentVersion, column: ColumnId, value: MetadataValue) -> Self {
Self {
version,
column,
value,
}
}
#[must_use]
pub const fn version(&self) -> DocumentVersion {
self.version
}
#[must_use]
pub const fn column(&self) -> ColumnId {
self.column
}
#[must_use]
pub const fn value(&self) -> &MetadataValue {
&self.value
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum MutationPayload {
#[cfg(feature = "graph-cypher")]
GraphCommit(Vec<u8>),
MixedBatchMember {
index: u32,
count: u32,
op: u16,
mutation: Box<MutationPayload>,
},
Upsert(IngestDocument),
Delete(Vec<DocId>),
MetadataEdit(MetadataEdit),
BatchMember {
index: u32,
count: u32,
document: IngestDocument,
},
}
pub fn encode_upsert(document: &IngestDocument) -> Result<Vec<u8>, PayloadError> {
encode_upsert_body(document, None, false)
}
pub fn encode_upsert_with_timestamp(document: &IngestDocument) -> Result<Vec<u8>, PayloadError> {
encode_upsert_body(document, Some(document.timestamp()), false)
}
pub fn encode_upsert_with_metadata(document: &IngestDocument) -> Result<Vec<u8>, PayloadError> {
encode_upsert_body(document, None, true)
}
pub fn encode_upsert_with_timestamp_and_metadata(
document: &IngestDocument,
) -> Result<Vec<u8>, PayloadError> {
encode_upsert_body(document, Some(document.timestamp()), true)
}
pub fn encode_upsert_v2(document: &IngestDocument) -> Result<Vec<u8>, PayloadError> {
let mut bitmap = UPSERT_V2_VECTOR;
if document.text().is_some() {
bitmap |= UPSERT_V2_TEXT;
}
if document.has_timestamp() {
bitmap |= UPSERT_V2_TIMESTAMP;
}
if document.has_metadata() {
bitmap |= UPSERT_V2_METADATA;
}
if document.has_columns() {
bitmap |= UPSERT_V2_TYPED_COLUMNS;
}
validate_vector(document.vector())?;
let dims = u32::try_from(document.vector().len()).map_err(|_| PayloadError::LengthOverflow)?;
let mut payload = Vec::new();
payload.extend_from_slice(&bitmap.to_le_bytes());
append_version(&mut payload, document.version());
payload.extend_from_slice(&dims.to_le_bytes());
for value in document.vector() {
payload.extend_from_slice(&value.to_bits().to_le_bytes());
}
if let Some(text) = document.text() {
append_length_prefixed(text.as_bytes(), &mut payload)?;
}
if document.has_timestamp() {
payload.extend_from_slice(&document.timestamp().to_le_bytes());
}
if document.has_metadata() {
append_length_prefixed(document.metadata(), &mut payload)?;
}
if document.has_columns() {
encode_typed_columns(document.columns(), &mut payload)?;
}
Ok(payload)
}
fn validate_vector(vector: &[f32]) -> Result<(), PayloadError> {
if vector.is_empty() {
return Err(PayloadError::EmptyVector);
}
if let Some(index) = vector.iter().position(|value| !value.is_finite()) {
return Err(PayloadError::NonFiniteVector { index });
}
Ok(())
}
fn append_length_prefixed(bytes: &[u8], output: &mut Vec<u8>) -> Result<(), PayloadError> {
let length = u32::try_from(bytes.len()).map_err(|_| PayloadError::LengthOverflow)?;
output.extend_from_slice(&length.to_le_bytes());
output.extend_from_slice(bytes);
Ok(())
}
fn encode_typed_columns(
columns: &[(ColumnId, PredicateValue)],
output: &mut Vec<u8>,
) -> Result<(), PayloadError> {
let count = u32::try_from(columns.len()).map_err(|_| PayloadError::LengthOverflow)?;
output.extend_from_slice(&count.to_le_bytes());
for (column, value) in columns {
let metadata = predicate_to_metadata(value);
let (kind, body) = encode_metadata_value(&metadata)?;
output.extend_from_slice(&column.get().to_le_bytes());
output.push(kind);
output.extend_from_slice(&[0_u8; 3]);
append_length_prefixed(&body, output)?;
}
Ok(())
}
pub(crate) fn encode_column_values(
columns: &[(ColumnId, PredicateValue)],
) -> Result<Vec<u8>, PayloadError> {
let mut encoded = Vec::new();
encode_typed_columns(columns, &mut encoded)?;
Ok(encoded)
}
fn predicate_to_metadata(value: &PredicateValue) -> MetadataValue {
match value {
PredicateValue::Id128(value) => MetadataValue::Id128(*value),
PredicateValue::U64(value) => MetadataValue::U64(*value),
PredicateValue::I64(value) => MetadataValue::I64(*value),
PredicateValue::F64(value) => MetadataValue::F64(*value),
PredicateValue::Bool(value) => MetadataValue::Bool(*value),
PredicateValue::String(value) => MetadataValue::String(value.clone()),
}
}
fn encode_upsert_body(
document: &IngestDocument,
timestamp: Option<i64>,
include_metadata: bool,
) -> Result<Vec<u8>, PayloadError> {
if document.vector().is_empty() {
return Err(PayloadError::EmptyVector);
}
let dims = u32::try_from(document.vector().len()).map_err(|_| PayloadError::LengthOverflow)?;
if let Some(index) = document
.vector()
.iter()
.position(|value| !value.is_finite())
{
return Err(PayloadError::NonFiniteVector { index });
}
let vector_bytes = document
.vector()
.len()
.checked_mul(std::mem::size_of::<f32>())
.ok_or(PayloadError::LengthOverflow)?;
let prefix = match (timestamp.is_some(), include_metadata) {
(false, false) => UPSERT_PREFIX_LEN,
(true, false) => TIMESTAMPED_UPSERT_PREFIX_LEN,
(false, true) => METADATA_UPSERT_PREFIX_LEN,
(true, true) => TIMESTAMPED_METADATA_UPSERT_PREFIX_LEN,
};
let capacity = prefix
.checked_add(if include_metadata {
document.metadata().len()
} else {
0
})
.and_then(|value| value.checked_add(vector_bytes))
.ok_or(PayloadError::LengthOverflow)?;
let mut payload = Vec::with_capacity(capacity);
append_header(&mut payload);
append_version(&mut payload, document.version());
if let Some(timestamp) = timestamp {
payload.extend_from_slice(×tamp.to_le_bytes());
}
if include_metadata {
let metadata_len =
u32::try_from(document.metadata().len()).map_err(|_| PayloadError::LengthOverflow)?;
payload.extend_from_slice(&metadata_len.to_le_bytes());
}
payload.extend_from_slice(&dims.to_le_bytes());
if include_metadata {
payload.extend_from_slice(document.metadata());
}
for value in document.vector() {
payload.extend_from_slice(&value.to_bits().to_le_bytes());
}
Ok(payload)
}
pub fn encode_delete(doc_ids: &[DocId]) -> Result<Vec<u8>, PayloadError> {
if doc_ids.is_empty() {
return Err(PayloadError::EmptyDelete);
}
let count = u32::try_from(doc_ids.len()).map_err(|_| PayloadError::LengthOverflow)?;
let body_bytes = doc_ids
.len()
.checked_mul(std::mem::size_of::<u128>())
.ok_or(PayloadError::LengthOverflow)?;
let capacity = DELETE_PREFIX_LEN
.checked_add(body_bytes)
.ok_or(PayloadError::LengthOverflow)?;
let mut payload = Vec::with_capacity(capacity);
append_header(&mut payload);
payload.extend_from_slice(&count.to_le_bytes());
for doc_id in doc_ids {
payload.extend_from_slice(&doc_id.get().to_le_bytes());
}
Ok(payload)
}
pub fn encode_metadata_edit(edit: &MetadataEdit) -> Result<Vec<u8>, PayloadError> {
let (kind, body) = encode_metadata_value(edit.value())?;
let value_len = u32::try_from(body.len()).map_err(|_| PayloadError::LengthOverflow)?;
let capacity = METADATA_PREFIX_LEN
.checked_add(body.len())
.ok_or(PayloadError::LengthOverflow)?;
let mut payload = Vec::with_capacity(capacity);
append_header(&mut payload);
append_version(&mut payload, edit.version());
payload.extend_from_slice(&edit.column().get().to_le_bytes());
payload.push(kind);
payload.extend_from_slice(&[0_u8; 3]);
payload.extend_from_slice(&value_len.to_le_bytes());
payload.extend_from_slice(&body);
Ok(payload)
}
pub fn decode_mutation(op: u16, payload: &[u8]) -> Result<MutationPayload, PayloadError> {
match op {
MIXED_BATCH_MEMBER_V1 => decode_mixed_batch_member(payload),
#[cfg(feature = "graph-cypher")]
GRAPH_COMMIT_V1 => decode_graph_commit(payload).map(MutationPayload::GraphCommit),
UPSERT_V1 => decode_upsert(payload).map(MutationPayload::Upsert),
DELETE_V1 => decode_delete(payload).map(MutationPayload::Delete),
METADATA_EDIT_V1 => decode_metadata_edit(payload).map(MutationPayload::MetadataEdit),
UPSERT_WITH_TIMESTAMP_V1 => {
decode_upsert_with_timestamp(payload).map(MutationPayload::Upsert)
}
UPSERT_WITH_METADATA_V1 => {
decode_upsert_with_metadata(payload).map(MutationPayload::Upsert)
}
UPSERT_WITH_TIMESTAMP_AND_METADATA_V1 => {
decode_upsert_with_timestamp_and_metadata(payload).map(MutationPayload::Upsert)
}
UPSERT_V2 => decode_upsert_v2(payload).map(MutationPayload::Upsert),
UPSERT_V2_BATCH_MEMBER => {
decode_upsert_v2_batch_member(payload).map(|(index, count, document)| {
MutationPayload::BatchMember {
index,
count,
document,
}
})
}
unknown => Err(PayloadError::UnknownOperation(unknown)),
}
}
pub fn encode_mixed_batch_member(
index: u32,
count: u32,
op: u16,
inner: &[u8],
) -> Result<Vec<u8>, PayloadError> {
validate_batch_position(index, count)?;
validate_mixed_operation(op)?;
let length = inner
.len()
.checked_add(10)
.ok_or(PayloadError::LengthOverflow)?;
let mut payload = Vec::new();
payload
.try_reserve_exact(length)
.map_err(|_| PayloadError::LengthOverflow)?;
payload.extend_from_slice(&index.to_le_bytes());
payload.extend_from_slice(&count.to_le_bytes());
payload.extend_from_slice(&op.to_le_bytes());
payload.extend_from_slice(inner);
Ok(payload)
}
pub fn decode_mixed_batch_member(payload: &[u8]) -> Result<MutationPayload, PayloadError> {
let mut cursor = Cursor::new(payload);
let index = cursor.read_u32()?;
let count = cursor.read_u32()?;
validate_batch_position(index, count)?;
let op = cursor.read_u16()?;
validate_mixed_operation(op)?;
let inner = payload.get(10..).ok_or(PayloadError::Truncated)?;
let mutation = Box::new(decode_mutation(op, inner)?);
Ok(MutationPayload::MixedBatchMember {
index,
count,
op,
mutation,
})
}
fn validate_mixed_operation(op: u16) -> Result<(), PayloadError> {
if !matches!(op, DELETE_V1 | UPSERT_V2 | GRAPH_COMMIT_V1) {
return Err(PayloadError::UnknownOperation(op));
}
Ok(())
}
#[cfg(feature = "graph-cypher")]
pub fn encode_graph_commit(envelope: &[u8]) -> Result<Vec<u8>, PayloadError> {
decode_graph_commit(envelope)
}
#[cfg(feature = "graph-cypher")]
pub fn decode_graph_commit(payload: &[u8]) -> Result<Vec<u8>, PayloadError> {
use crate::property_graph::wal::{
STACK_RESERVATION_BYTES, WalResources, validate_envelope_framing,
};
let mut cancelled = || false;
let mut resources = WalResources::new(
64 * crate::property_graph::wal::MAX_ENVELOPE_BYTES as u64,
STACK_RESERVATION_BYTES,
&mut cancelled,
)
.map_err(PayloadError::GraphEnvelope)?;
validate_envelope_framing(payload, &mut resources).map_err(PayloadError::GraphEnvelope)?;
let mut bytes = Vec::new();
bytes
.try_reserve_exact(payload.len())
.map_err(|_| PayloadError::LengthOverflow)?;
bytes.extend_from_slice(payload);
Ok(bytes)
}
pub fn encode_upsert_v2_batch_member(
index: u32,
count: u32,
upsert_v2: &[u8],
) -> Result<Vec<u8>, PayloadError> {
validate_batch_position(index, count)?;
let mut payload = Vec::new();
payload
.try_reserve_exact(upsert_v2.len().saturating_add(8))
.map_err(|_| PayloadError::LengthOverflow)?;
payload.extend_from_slice(&index.to_le_bytes());
payload.extend_from_slice(&count.to_le_bytes());
payload.extend_from_slice(upsert_v2);
Ok(payload)
}
pub fn decode_upsert_v2_batch_member(
payload: &[u8],
) -> Result<(u32, u32, IngestDocument), PayloadError> {
let mut cursor = Cursor::new(payload);
let index = cursor.read_u32()?;
let count = cursor.read_u32()?;
validate_batch_position(index, count)?;
let body = payload.get(8..).ok_or(PayloadError::Truncated)?;
decode_upsert_v2(body).map(|document| (index, count, document))
}
fn validate_batch_position(index: u32, count: u32) -> Result<(), PayloadError> {
if count < 2 || index >= count {
return Err(PayloadError::BatchPosition { index, count });
}
Ok(())
}
pub fn decode_upsert_v2(payload: &[u8]) -> Result<IngestDocument, PayloadError> {
let mut cursor = Cursor::new(payload);
let bitmap = cursor.read_u32()?;
let reserved = bitmap & !UPSERT_V2_KNOWN_FIELDS;
if reserved != 0 {
return Err(PayloadError::FieldBitmap(reserved));
}
let version = cursor.read_version()?;
let vector = if bitmap & UPSERT_V2_VECTOR != 0 {
let dims = usize::try_from(cursor.read_u32()?).map_err(|_| PayloadError::LengthOverflow)?;
if dims == 0 {
return Err(PayloadError::EmptyVector);
}
let mut vector = Vec::with_capacity(dims);
for index in 0..dims {
let value = f32::from_bits(cursor.read_u32()?);
if !value.is_finite() {
return Err(PayloadError::NonFiniteVector { index });
}
vector.push(value);
}
vector
} else {
Vec::new()
};
let text = if bitmap & UPSERT_V2_TEXT != 0 {
Some(cursor.read_string()?)
} else {
None
};
let timestamp = if bitmap & UPSERT_V2_TIMESTAMP != 0 {
Some(i64::from_le_bytes(cursor.take()?))
} else {
None
};
let metadata = if bitmap & UPSERT_V2_METADATA != 0 {
Some(cursor.read_length_prefixed()?.to_vec())
} else {
None
};
let columns = if bitmap & UPSERT_V2_TYPED_COLUMNS != 0 {
Some(decode_typed_columns(&mut cursor)?)
} else {
None
};
cursor.finish()?;
let mut document = IngestDocument::new(version, vector);
if let Some(text) = text {
document = document.with_text(text);
}
if let Some(timestamp) = timestamp {
document = document.with_timestamp(timestamp);
}
if let Some(metadata) = metadata {
document = document.with_metadata(metadata);
}
if let Some(columns) = columns {
document = document.with_columns(columns);
}
Ok(document)
}
fn decode_typed_columns(
cursor: &mut Cursor<'_>,
) -> Result<Vec<(ColumnId, PredicateValue)>, PayloadError> {
let count = usize::try_from(cursor.read_u32()?).map_err(|_| PayloadError::LengthOverflow)?;
let mut columns = Vec::with_capacity(count);
for _ in 0..count {
let column = ColumnId::new(cursor.read_u32()?);
let kind = cursor.read_u8()?;
let reserved = cursor.read_u24()?;
if reserved != 0 {
return Err(PayloadError::Reserved(reserved));
}
let body = cursor.read_length_prefixed()?;
let value = metadata_to_predicate(decode_metadata_value(kind, body)?)?;
columns.push((column, value));
}
Ok(columns)
}
pub(crate) fn decode_column_values(
encoded: &[u8],
) -> Result<Vec<(ColumnId, PredicateValue)>, PayloadError> {
let mut cursor = Cursor::new(encoded);
let columns = decode_typed_columns(&mut cursor)?;
cursor.finish()?;
Ok(columns)
}
fn metadata_to_predicate(value: MetadataValue) -> Result<PredicateValue, PayloadError> {
match value {
MetadataValue::Null => Err(PayloadError::ValueKind(VALUE_NULL)),
MetadataValue::Id128(value) => Ok(PredicateValue::Id128(value)),
MetadataValue::U64(value) => Ok(PredicateValue::U64(value)),
MetadataValue::I64(value) => Ok(PredicateValue::I64(value)),
MetadataValue::F64(value) => Ok(PredicateValue::F64(value)),
MetadataValue::Bool(value) => Ok(PredicateValue::Bool(value)),
MetadataValue::String(value) => Ok(PredicateValue::String(value)),
}
}
pub fn decode_upsert(payload: &[u8]) -> Result<IngestDocument, PayloadError> {
decode_upsert_body(payload, false, false)
}
pub fn decode_upsert_with_timestamp(payload: &[u8]) -> Result<IngestDocument, PayloadError> {
decode_upsert_body(payload, true, false)
}
pub fn decode_upsert_with_metadata(payload: &[u8]) -> Result<IngestDocument, PayloadError> {
decode_upsert_body(payload, false, true)
}
pub fn decode_upsert_with_timestamp_and_metadata(
payload: &[u8],
) -> Result<IngestDocument, PayloadError> {
decode_upsert_body(payload, true, true)
}
fn decode_upsert_body(
payload: &[u8],
has_timestamp: bool,
has_metadata: bool,
) -> Result<IngestDocument, PayloadError> {
let mut cursor = Cursor::new(payload);
cursor.require_header()?;
let version = cursor.read_version()?;
let timestamp = if has_timestamp {
i64::from_le_bytes(cursor.take()?)
} else {
0
};
let metadata_len = if has_metadata {
usize::try_from(cursor.read_u32()?).map_err(|_| PayloadError::LengthOverflow)?
} else {
0
};
let dims = usize::try_from(cursor.read_u32()?).map_err(|_| PayloadError::LengthOverflow)?;
if dims == 0 {
return Err(PayloadError::EmptyVector);
}
let body_bytes = dims
.checked_mul(std::mem::size_of::<f32>())
.ok_or(PayloadError::LengthOverflow)?;
let remaining = metadata_len
.checked_add(body_bytes)
.ok_or(PayloadError::LengthOverflow)?;
cursor.require_remaining(remaining)?;
let metadata = cursor.read_bytes(metadata_len)?.to_vec();
let mut vector = Vec::with_capacity(dims);
for index in 0..dims {
let value = f32::from_bits(cursor.read_u32()?);
if !value.is_finite() {
return Err(PayloadError::NonFiniteVector { index });
}
vector.push(value);
}
cursor.finish()?;
let mut document = IngestDocument::new(version, vector);
if has_timestamp {
document = document.with_timestamp(timestamp);
}
if has_metadata {
document = document.with_metadata(metadata);
}
Ok(document)
}
pub fn decode_delete(payload: &[u8]) -> Result<Vec<DocId>, PayloadError> {
let mut cursor = Cursor::new(payload);
cursor.require_header()?;
let count = usize::try_from(cursor.read_u32()?).map_err(|_| PayloadError::LengthOverflow)?;
if count == 0 {
return Err(PayloadError::EmptyDelete);
}
let body_bytes = count
.checked_mul(std::mem::size_of::<u128>())
.ok_or(PayloadError::LengthOverflow)?;
cursor.require_remaining(body_bytes)?;
let mut doc_ids = Vec::with_capacity(count);
for _ in 0..count {
doc_ids.push(DocId::new(cursor.read_u128()?));
}
cursor.finish()?;
Ok(doc_ids)
}
pub fn decode_metadata_edit(payload: &[u8]) -> Result<MetadataEdit, PayloadError> {
let mut cursor = Cursor::new(payload);
cursor.require_header()?;
let version = cursor.read_version()?;
let column = ColumnId::new(cursor.read_u32()?);
let kind = cursor.read_u8()?;
let reserved = cursor.read_u24()?;
if reserved != 0 {
return Err(PayloadError::Reserved(reserved));
}
let value_len =
usize::try_from(cursor.read_u32()?).map_err(|_| PayloadError::LengthOverflow)?;
let body = cursor.read_bytes(value_len)?;
let value = decode_metadata_value(kind, body)?;
cursor.finish()?;
Ok(MetadataEdit::new(version, column, value))
}
fn append_header(payload: &mut Vec<u8>) {
payload.extend_from_slice(&PAYLOAD_VERSION.to_le_bytes());
payload.extend_from_slice(&0_u16.to_le_bytes());
}
fn append_version(payload: &mut Vec<u8>, version: DocumentVersion) {
payload.extend_from_slice(&version.doc_id().get().to_le_bytes());
payload.extend_from_slice(&version.revision().get().to_le_bytes());
}
fn encode_metadata_value(value: &MetadataValue) -> Result<(u8, Vec<u8>), PayloadError> {
let encoded = match value {
MetadataValue::Null => (VALUE_NULL, Vec::new()),
MetadataValue::Id128(value) => (VALUE_ID128, value.get().to_le_bytes().to_vec()),
MetadataValue::U64(value) => (VALUE_U64, value.to_le_bytes().to_vec()),
MetadataValue::I64(value) => (VALUE_I64, value.to_le_bytes().to_vec()),
MetadataValue::F64(value) => (VALUE_F64, value.to_bits().to_le_bytes().to_vec()),
MetadataValue::Bool(value) => (VALUE_BOOL, vec![u8::from(*value)]),
MetadataValue::String(value) => {
u32::try_from(value.len()).map_err(|_| PayloadError::LengthOverflow)?;
(VALUE_STRING, value.as_bytes().to_vec())
}
};
Ok(encoded)
}
fn decode_metadata_value(kind: u8, body: &[u8]) -> Result<MetadataValue, PayloadError> {
match kind {
VALUE_NULL => {
require_value_len(kind, body, 0)?;
Ok(MetadataValue::Null)
}
VALUE_ID128 => {
require_value_len(kind, body, 16)?;
read_array(body)
.map(u128::from_le_bytes)
.map(DocId::new)
.map(MetadataValue::Id128)
}
VALUE_U64 => {
require_value_len(kind, body, 8)?;
read_array(body)
.map(u64::from_le_bytes)
.map(MetadataValue::U64)
}
VALUE_I64 => {
require_value_len(kind, body, 8)?;
read_array(body)
.map(i64::from_le_bytes)
.map(MetadataValue::I64)
}
VALUE_F64 => {
require_value_len(kind, body, 8)?;
read_array(body)
.map(u64::from_le_bytes)
.map(f64::from_bits)
.map(MetadataValue::F64)
}
VALUE_BOOL => {
require_value_len(kind, body, 1)?;
match body.first().copied().ok_or(PayloadError::Truncated)? {
0 => Ok(MetadataValue::Bool(false)),
1 => Ok(MetadataValue::Bool(true)),
value => Err(PayloadError::Boolean(value)),
}
}
VALUE_STRING => std::str::from_utf8(body)
.map(|value| MetadataValue::String(value.to_owned()))
.map_err(|_| PayloadError::Utf8),
unknown => Err(PayloadError::ValueKind(unknown)),
}
}
fn require_value_len(kind: u8, body: &[u8], expected: usize) -> Result<(), PayloadError> {
if body.len() == expected {
Ok(())
} else {
Err(PayloadError::ValueLength {
kind,
expected,
actual: body.len(),
})
}
}
fn read_array<const N: usize>(bytes: &[u8]) -> Result<[u8; N], PayloadError> {
<[u8; N]>::try_from(bytes).map_err(|_| PayloadError::Truncated)
}
struct Cursor<'a> {
bytes: &'a [u8],
offset: usize,
}
impl<'a> Cursor<'a> {
const fn new(bytes: &'a [u8]) -> Self {
Self { bytes, offset: 0 }
}
fn require_header(&mut self) -> Result<(), PayloadError> {
let version = self.read_u16()?;
if version != PAYLOAD_VERSION {
return Err(PayloadError::Version(version));
}
let flags = self.read_u16()?;
if flags != 0 {
return Err(PayloadError::Flags(flags));
}
Ok(())
}
fn read_version(&mut self) -> Result<DocumentVersion, PayloadError> {
Ok(DocumentVersion::new(
DocId::new(self.read_u128()?),
Revision::new(self.read_u64()?),
))
}
fn take<const N: usize>(&mut self) -> Result<[u8; N], PayloadError> {
read_array(self.read_bytes(N)?)
}
fn read_bytes(&mut self, length: usize) -> Result<&'a [u8], PayloadError> {
let end = self
.offset
.checked_add(length)
.ok_or(PayloadError::LengthOverflow)?;
let bytes = self
.bytes
.get(self.offset..end)
.ok_or(PayloadError::Truncated)?;
self.offset = end;
Ok(bytes)
}
fn require_remaining(&self, length: usize) -> Result<(), PayloadError> {
if self.remaining() < length {
Err(PayloadError::Truncated)
} else {
Ok(())
}
}
fn read_u8(&mut self) -> Result<u8, PayloadError> {
self.take::<1>().map(u8::from_le_bytes)
}
fn read_u16(&mut self) -> Result<u16, PayloadError> {
self.take().map(u16::from_le_bytes)
}
fn read_u24(&mut self) -> Result<u32, PayloadError> {
let [low, middle, high] = self.take::<3>()?;
Ok(u32::from(low) | (u32::from(middle) << 8) | (u32::from(high) << 16))
}
fn read_u32(&mut self) -> Result<u32, PayloadError> {
self.take().map(u32::from_le_bytes)
}
fn read_u64(&mut self) -> Result<u64, PayloadError> {
self.take().map(u64::from_le_bytes)
}
fn read_u128(&mut self) -> Result<u128, PayloadError> {
self.take().map(u128::from_le_bytes)
}
fn read_length_prefixed(&mut self) -> Result<&'a [u8], PayloadError> {
let length = usize::try_from(self.read_u32()?).map_err(|_| PayloadError::LengthOverflow)?;
self.read_bytes(length)
}
fn read_string(&mut self) -> Result<String, PayloadError> {
std::str::from_utf8(self.read_length_prefixed()?)
.map(str::to_owned)
.map_err(|_| PayloadError::Utf8)
}
fn remaining(&self) -> usize {
self.bytes.len().saturating_sub(self.offset)
}
fn finish(self) -> Result<(), PayloadError> {
let remaining = self.remaining();
if remaining == 0 {
Ok(())
} else {
Err(PayloadError::TrailingBytes(remaining))
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod id128_tests {
use super::*;
#[test]
fn id128_wal_layout_is_sixteen_little_endian_bytes() {
let value = MetadataValue::Id128(DocId::new(0xffeeddccbbaa99887766554433221100));
let bytes = vec![
0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd,
0xee, 0xff,
];
assert_eq!(encode_metadata_value(&value).unwrap(), (6, bytes.clone()));
assert_eq!(decode_metadata_value(6, &bytes).unwrap(), value);
for length in [0, 8, 15, 17] {
assert!(decode_metadata_value(6, &vec![0; length]).is_err());
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod transaction_format_tests {
use super::*;
fn binding() -> TransactionBinding {
TransactionBinding {
transaction: 1,
participant: 2,
first_seq: 3,
last_seq: 3,
manifest_digest: 4,
final_generation: 5,
}
}
#[test]
#[allow(clippy::indexing_slicing)]
fn prepared_encoding_golden_and_rejections() {
let body = encode_delete(&[DocId::new(6)]).unwrap();
let bytes = encode_prepared(binding(), 0, 1, DELETE_V1, &body).unwrap();
let hex = bytes.iter().map(|b| format!("{b:02x}")).collect::<String>();
assert_eq!(
hex,
"010000000100000000000000000000000000000002000000000000000000000000000000030000000000000003000000000000000400000000000000050000000000000000000000010000000200010000000100000006000000000000000000000000000000"
);
let decoded = decode_prepared(&bytes).unwrap();
assert_eq!(decoded.binding, binding());
assert_eq!(
decoded.mutation,
MutationPayload::Delete(vec![DocId::new(6)])
);
for end in 0..bytes.len() {
assert!(decode_prepared(&bytes[..end]).is_err(), "{end}");
}
for offset in [0, 2, 4, 20, 36, 60, 68, 72, 76] {
let mut bad = bytes.clone();
bad[offset] = 255;
if matches!(offset, 4 | 20 | 36 | 60) {
if let Ok(member) = decode_prepared(&bad) {
assert_eq!(
transaction_decision(member.binding, Some(binding())),
TransactionDecision::Mismatch
);
}
} else {
assert!(decode_prepared(&bad).is_err(), "{offset}");
}
}
assert!(encode_prepared(binding(), 0, 1, PREPARED_MUTATION_V1, &bytes).is_err());
let mut trailing = bytes.clone();
trailing.push(0);
assert!(decode_prepared(&trailing).is_err());
}
#[test]
fn transaction_decision_table() {
assert_eq!(
transaction_decision(binding(), None),
TransactionDecision::Undecided
);
assert_eq!(
transaction_decision(binding(), Some(binding())),
TransactionDecision::Committed
);
for field in 0..6 {
let mut changed = binding();
match field {
0 => changed.transaction += 1,
1 => changed.participant += 1,
2 => changed.first_seq += 1,
3 => changed.last_seq += 1,
4 => changed.manifest_digest += 1,
_ => changed.final_generation += 1,
}
assert_eq!(
transaction_decision(binding(), Some(changed)),
TransactionDecision::Mismatch
);
}
let mut invalid = binding();
invalid.transaction = 0;
assert_eq!(
transaction_decision(invalid, None),
TransactionDecision::Mismatch
);
}
#[test]
fn prepared_operation_is_recognized() {
let mut bytes = vec![1, 0, 0, 0];
bytes.extend_from_slice(&1_u128.to_le_bytes());
bytes.extend_from_slice(&2_u128.to_le_bytes());
for value in [3_u64, 3, 4, 5] {
bytes.extend_from_slice(&value.to_le_bytes());
}
bytes.extend_from_slice(&0_u32.to_le_bytes());
bytes.extend_from_slice(&1_u32.to_le_bytes());
bytes.extend_from_slice(&2_u16.to_le_bytes());
bytes.extend_from_slice(&super::encode_delete(&[super::DocId::new(6)]).unwrap());
assert!(super::decode_prepared(&bytes).is_ok());
}
}
pub const PREPARED_MUTATION_V1: u16 = 9;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct TransactionBinding {
pub transaction: u128,
pub participant: u128,
pub first_seq: u64,
pub last_seq: u64,
pub manifest_digest: u64,
pub final_generation: u64,
}
impl TransactionBinding {
pub fn encode(self) -> Result<Vec<u8>, PayloadError> {
self.member_count()?;
let mut bytes = Vec::with_capacity(64);
bytes.extend_from_slice(&self.transaction.to_le_bytes());
bytes.extend_from_slice(&self.participant.to_le_bytes());
for value in [
self.first_seq,
self.last_seq,
self.manifest_digest,
self.final_generation,
] {
bytes.extend_from_slice(&value.to_le_bytes());
}
Ok(bytes)
}
pub fn decode(bytes: &[u8]) -> Result<Self, PayloadError> {
let mut cursor = Cursor::new(bytes);
let binding = Self {
transaction: cursor.read_u128()?,
participant: cursor.read_u128()?,
first_seq: cursor.read_u64()?,
last_seq: cursor.read_u64()?,
manifest_digest: cursor.read_u64()?,
final_generation: cursor.read_u64()?,
};
cursor.finish()?;
binding.member_count()?;
Ok(binding)
}
fn member_count(self) -> Result<u32, PayloadError> {
if self.transaction == 0
|| self.participant == 0
|| self.first_seq == 0
|| self.final_generation == 0
{
return Err(PayloadError::TransactionBinding);
}
self.last_seq
.checked_sub(self.first_seq)
.and_then(|n| n.checked_add(1))
.and_then(|n| u32::try_from(n).ok())
.ok_or(PayloadError::TransactionBinding)
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct PreparedMutation {
pub binding: TransactionBinding,
pub index: u32,
pub count: u32,
pub op: u16,
pub mutation: MutationPayload,
}
pub fn encode_prepared(
binding: TransactionBinding,
index: u32,
count: u32,
op: u16,
body: &[u8],
) -> Result<Vec<u8>, PayloadError> {
if count != binding.member_count()? || index >= count {
return Err(PayloadError::BatchPosition { index, count });
}
standalone_mutation(op, body)?;
let mut bytes = Vec::new();
bytes
.try_reserve_exact(
body.len()
.checked_add(78)
.ok_or(PayloadError::LengthOverflow)?,
)
.map_err(|_| PayloadError::LengthOverflow)?;
append_header(&mut bytes);
bytes.extend_from_slice(&binding.encode()?);
bytes.extend_from_slice(&index.to_le_bytes());
bytes.extend_from_slice(&count.to_le_bytes());
bytes.extend_from_slice(&op.to_le_bytes());
bytes.extend_from_slice(body);
Ok(bytes)
}
pub fn decode_prepared(bytes: &[u8]) -> Result<PreparedMutation, PayloadError> {
let mut cursor = Cursor::new(bytes);
cursor.require_header()?;
let binding = TransactionBinding::decode(cursor.read_bytes(64)?)?;
let index = cursor.read_u32()?;
let count = cursor.read_u32()?;
if count != binding.member_count()? || index >= count {
return Err(PayloadError::BatchPosition { index, count });
}
let op = cursor.read_u16()?;
let mutation = standalone_mutation(op, cursor.read_bytes(cursor.remaining())?)?;
Ok(PreparedMutation {
binding,
index,
count,
op,
mutation,
})
}
fn standalone_mutation(op: u16, bytes: &[u8]) -> Result<MutationPayload, PayloadError> {
if !(UPSERT_V1..=UPSERT_V2).contains(&op) {
return Err(PayloadError::UnknownOperation(op));
}
decode_mutation(op, bytes)
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TransactionDecision {
Committed,
Undecided,
Mismatch,
}
pub fn transaction_decision(
prepared: TransactionBinding,
selected: Option<TransactionBinding>,
) -> TransactionDecision {
match selected {
None if prepared.member_count().is_ok() => TransactionDecision::Undecided,
Some(binding) if prepared.member_count().is_ok() && binding == prepared => {
TransactionDecision::Committed
}
_ => TransactionDecision::Mismatch,
}
}