use std::num::{NonZeroU32, NonZeroUsize};
use super::mutations::PreparedWrite;
use super::{Captured, Failure, Storage, document_key, identifier, reads, revision_key};
use crate::serve::protocol::{
Change, DocumentAppendPrune, DocumentWrite, RetentionResult, Version,
};
#[derive(Clone, Copy, Debug)]
pub struct AppendPruneLimits {
keep_ceiling: NonZeroU32,
removal_cap: NonZeroU32,
}
impl Default for AppendPruneLimits {
fn default() -> Self {
Self {
keep_ceiling: NonZeroU32::MIN.saturating_add(4_095),
removal_cap: NonZeroU32::MIN.saturating_add(63),
}
}
}
impl AppendPruneLimits {
pub fn new(keep_ceiling: u32, removal_cap: u32) -> Result<Self, Failure> {
if removal_cap == 1 {
return Err(Failure::new(
"append_prune_removal_cap_too_small",
"append_prune_removal_cap must be at least two",
));
}
Ok(Self {
keep_ceiling: NonZeroU32::new(keep_ceiling).ok_or_else(|| {
Failure::new(
"append_prune_keep_ceiling_zero",
"append_prune_keep_ceiling must be non-zero",
)
})?,
removal_cap: NonZeroU32::new(removal_cap).ok_or_else(|| {
Failure::new(
"append_prune_removal_cap_zero",
"append_prune_removal_cap must be non-zero",
)
})?,
})
}
}
pub(in crate::serve) fn append_prune_address(
scope: &[String],
suffix: &str,
) -> Result<(String, String), Failure> {
if !(1..=8).contains(&scope.len()) {
return Err(Failure::new(
"append_prune_scope_components_invalid",
"scope must have one to eight components",
));
}
for component in scope {
if component.is_empty() || component.contains('/') || component.contains('\0') {
return Err(Failure::new(
"append_prune_scope_component_invalid",
"scope components must be non-empty and contain no separator or NUL",
));
}
}
if suffix.is_empty() || suffix.contains('/') || suffix.contains('\0') {
return Err(Failure::new(
"append_prune_suffix_invalid",
"suffix must be non-empty and contain no separator or NUL",
));
}
let count = scope.len().to_string();
let length = scope
.iter()
.try_fold(count.len() + 1, |length, component| {
length
.checked_add(component.len())
.and_then(|n| n.checked_add(1))
})
.and_then(|length| length.checked_add(suffix.len()));
if length.is_none_or(|length| length > 512) {
return Err(Failure::new(
"append_prune_document_id_invalid",
"formed document id exceeds its byte bound",
));
}
let mut prefix = count;
prefix.push('/');
for component in scope {
prefix.push_str(component);
prefix.push('/');
}
let mut document_id = prefix.clone();
document_id.push_str(suffix);
identifier(&document_id)
.map_err(|failure| Failure::new("append_prune_document_id_invalid", failure.message))?;
Ok((prefix, document_id))
}
impl Storage {
pub(super) fn prepare_append_prune(
&self,
at: &Captured,
write: &DocumentAppendPrune,
next: Version,
limits: AppendPruneLimits,
) -> Result<(Vec<PreparedWrite>, RetentionResult), Failure> {
identifier(&write.collection)?;
let (scope_prefix, document_id) = append_prune_address(&write.scope, &write.suffix)?;
if write.keep == 0 {
return Err(Failure::new(
"append_prune_keep_zero",
"keep must be non-zero",
));
}
if write.keep > limits.keep_ceiling.get() {
return Err(Failure::new(
"append_prune_keep_limit",
"keep exceeds append_prune_keep_ceiling",
));
}
if reads::get_bytes(
&self.reader,
at,
&document_key(&write.collection, &document_id),
)?
.is_some()
|| reads::get_bytes(
&self.reader,
at,
&revision_key(&write.collection, &document_id),
)?
.is_some()
{
return Err(Failure::new(
"append_prune_key_exists",
"new key is already present",
));
}
let field_prefix = format!("w\0{}\0", write.collection).into_bytes();
if !reads::prefix_keys(&self.reader, at, &field_prefix, NonZeroUsize::MIN)?.is_empty() {
return Err(Failure::new(
"append_prune_vector_field",
"plain retention refuses vector fields",
));
}
let keep = usize::try_from(write.keep).map_err(Failure::storage)?;
let cap = usize::try_from(limits.removal_cap.get()).map_err(Failure::storage)?;
let scan = keep
.checked_add(cap)
.and_then(|n| n.checked_add(1))
.and_then(NonZeroUsize::new)
.ok_or_else(|| {
Failure::new(
"append_prune_limit_overflow",
"retention scan bound exceeds usize",
)
})?;
let prefix = document_key(&write.collection, &scope_prefix);
let mut keys = reads::prefix_keys(&self.reader, at, &prefix, scan)?;
let new_key = document_key(&write.collection, &document_id);
keys.push(new_key.clone());
keys.sort();
let excess = keys.len().saturating_sub(keep);
let removed = excess.min(cap);
if keys.iter().take(removed).any(|key| key == &new_key) {
return Err(Failure::new(
"append_prune_key_would_be_removed",
"new key would be removed by this append",
));
}
let mut writes = self.prepare_change(
at,
&Change::Put(DocumentWrite {
collection: write.collection.clone(),
document_id: document_id.clone(),
document: write.document.clone(),
}),
next,
)?;
self.prepare_retention_removals(
at,
&write.collection,
keys.into_iter().take(removed),
&mut writes,
)?;
Ok((
writes,
RetentionResult {
document_id,
removed: u32::try_from(removed).map_err(Failure::storage)?,
more: excess > removed,
},
))
}
fn prepare_retention_removals(
&self,
at: &Captured,
collection: &str,
keys: impl Iterator<Item = Vec<u8>>,
writes: &mut Vec<PreparedWrite>,
) -> Result<(), Failure> {
let document_prefix = document_key(collection, "");
for key in keys {
let id = key
.strip_prefix(document_prefix.as_slice())
.ok_or_else(|| {
Failure::new("corrupt_key", "retention record is outside its collection")
})?;
let id = std::str::from_utf8(id).map_err(Failure::storage)?;
identifier(id)?;
for (tag, code) in [
("e", "append_prune_outgoing_edge"),
("i", "append_prune_incoming_edge"),
] {
let prefix = format!("{tag}\0{collection}\0{id}\0").into_bytes();
if !reads::prefix_keys(&self.reader, at, &prefix, NonZeroUsize::MIN)?.is_empty() {
return Err(Failure::new(
code,
"plain retention refuses an incident edge",
));
}
}
let revision = revision_key(collection, id);
if reads::get_bytes(&self.reader, at, &revision)?.is_none() {
return Err(Failure::new(
"append_prune_revision_missing",
"retained record has no revision",
));
}
writes.push((revision, None));
writes.push((key, None));
}
Ok(())
}
}