use std::collections::{BTreeMap, btree_map::Entry};
use haematite::branch::{CommitDurability, CommitRequest, current_timestamp};
use haematite::db::{ReadOnlyDatabase, branches::BranchDatabase};
use haematite::tree::Hash;
use super::mutations::PreparedWrite;
use super::{
BranchRefRecord, BranchRefStore, BranchRegistry, Captured, Failure, MergePlan, MergeState,
Receipt, SERVICE, Storage, branch_name, collection_key, edges, identifier, open_branch, reads,
receipt_key, revision_key, version,
};
use crate::serve::access::Access;
use crate::serve::protocol::{
AuditVerdict, AuditedBucket, Change, CountBucket, CountReason, CountsAudit, CountsRebuild,
KeptCount, ReceiptQuery, Version,
};
pub(super) struct View<'a> {
reader: &'a ReadOnlyDatabase,
at: &'a Captured,
ready: bool,
}
impl<'a> View<'a> {
pub(super) fn new(reader: &'a ReadOnlyDatabase, at: &'a Captured) -> Result<Self, Failure> {
Ok(Self {
reader,
at,
ready: is_ready(reader, at)?,
})
}
pub(super) const fn unavailable(&self, allowed: bool) -> Option<CountReason> {
if !self.ready {
Some(CountReason::CountsNotReady)
} else if !allowed {
Some(CountReason::Withheld)
} else {
None
}
}
pub(super) fn documents(
&self,
access: &Access,
collection: &str,
) -> Result<KeptCount, Failure> {
if let Some(reason) = self.unavailable(access.can_read_collection(collection)) {
return Ok(KeptCount {
value: None,
reason: Some(reason),
});
}
Ok(KeptCount {
value: Some(required_documents(self.reader, self.at, collection)?),
reason: None,
})
}
pub(super) fn relationships(
&self,
access: &Access,
relation: &str,
) -> Result<KeptCount, Failure> {
if let Some(reason) = self.unavailable(access.can_read_store()) {
return Ok(KeptCount {
value: None,
reason: Some(reason),
});
}
let key = format!("k\0e\0{relation}");
let value = read(self.reader, self.at, key.as_bytes())?
.ok_or_else(|| Failure::new("counts_corrupt", "declared relation has no count"))?;
Ok(KeptCount {
value: Some(decode(&value)?),
reason: None,
})
}
}
const READY_KEY: &[u8] = b"k\0ready";
const READY_VALUE: &[u8] = b"1";
#[derive(Default)]
struct BucketCount {
counted: u64,
kept: Option<u64>,
present: bool,
}
impl CountBucket {
fn key(&self) -> Vec<u8> {
match self {
Self::Documents { collection } => document_count_key(collection),
Self::Relationships { relation } => format!("k\0e\0{relation}").into_bytes(),
Self::Vectors { collection, field } => {
format!("k\0v\0{collection}\0{field}").into_bytes()
}
}
}
}
fn counted_buckets(
reader: &ReadOnlyDatabase,
at: &Captured,
) -> Result<BTreeMap<CountBucket, BucketCount>, Failure> {
let mut counts = BTreeMap::new();
for (prefix, parts, kind) in [
(b"c\0".as_slice(), 1, 0),
(b"n\0".as_slice(), 1, 1),
(b"w\0".as_slice(), 2, 2),
(b"d\0".as_slice(), 2, 3),
(b"e\0".as_slice(), 5, 4),
(b"v\0".as_slice(), 3, 5),
] {
reads::scan_rows(reader, at, prefix, parts, |parts, bytes| {
let bucket = match kind {
0 | 3 => CountBucket::Documents {
collection: parts[0].to_owned(),
},
1 => {
super::relations::decode(parts[0], bytes)?;
CountBucket::Relationships {
relation: parts[0].to_owned(),
}
}
4 => CountBucket::Relationships {
relation: parts[2].to_owned(),
},
_ => {
if kind == 2
&& matches!(
super::embedded::read_field(bytes)?,
super::embedded::Field::Embedded(_)
)
{
return Ok(());
}
CountBucket::Vectors {
collection: parts[0].to_owned(),
field: parts[1].to_owned(),
}
}
};
let counter: &mut BucketCount = counts.entry(bucket).or_default();
counter.present = true;
if kind >= 3 {
counter.counted = added_count(counter.counted, 1)?;
}
Ok(())
})?;
}
Ok(counts)
}
fn count_rows(
reader: &ReadOnlyDatabase,
at: &Captured,
counts: &mut BTreeMap<CountBucket, BucketCount>,
audit: bool,
) -> Result<(), Failure> {
for (prefix, parts, kind) in [
(b"k\0d\0".as_slice(), 1, 0),
(b"k\0e\0".as_slice(), 1, 1),
(b"k\0v\0".as_slice(), 2, 2),
] {
reads::scan_rows(reader, at, prefix, parts, |parts, bytes| {
let bucket = match kind {
0 => CountBucket::Documents {
collection: parts[0].to_owned(),
},
1 => CountBucket::Relationships {
relation: parts[0].to_owned(),
},
_ => CountBucket::Vectors {
collection: parts[0].to_owned(),
field: parts[1].to_owned(),
},
};
let counter = counts.entry(bucket).or_default();
if audit && counter.kept.replace(decode(bytes)?).is_some() {
return Err(Failure::new("counts_corrupt", "duplicate count bucket"));
}
Ok(())
})?;
}
Ok(())
}
impl Storage {
pub fn rebuild_counts(&mut self, request: &CountsRebuild) -> Result<Receipt, Failure> {
identifier(&request.branch)?;
identifier(&request.request_id)?;
let name = branch_name(&request.branch);
if self.pending_discards.contains(&name) {
return Err(Failure::new(
"counts_discard_pending",
"branch discard is incomplete",
));
}
let record = self
.refs
.get(&name)
.ok_or_else(|| Failure::new("counts_branch_missing", &request.branch))?
.clone();
let fingerprint =
Hash::of(&serde_json::to_vec(&("counts.rebuild", request)).map_err(Failure::storage)?)
.to_string();
if let Some(receipt) = self.receipt(&ReceiptQuery {
branch: request.branch.clone(),
request_id: request.request_id.clone(),
})? {
return if receipt.fingerprint == fingerprint {
Ok(receipt)
} else {
Err(Failure::new(
"request_conflict",
"request identity has different content",
))
};
}
let at = Captured {
branch: request.branch.clone(),
version: version(&record),
shards: record.shards.clone(),
};
self.refuse_count_merge(&at).map_err(|error| {
if error.code == "counts_merge_pending" {
error
} else {
count_scan_failure(&error)
}
})?;
let next = Version {
generation: record.created,
sequence: record
.seq
.checked_add(1)
.ok_or_else(|| Failure::new("sequence_exhausted", "branch sequence exhausted"))?,
};
let counts = rebuild_buckets(&self.reader, &at)?;
#[cfg(test)]
rebuild_probe::pause()?;
let receipt = Receipt {
retention: None,
branch: request.branch.clone(),
request_id: request.request_id.clone(),
fingerprint,
version: next,
};
self.commit_rebuilt_counts(&name, &record, &at, receipt, counts)
}
fn commit_rebuilt_counts(
&mut self,
name: &str,
record: &BranchRefRecord,
at: &Captured,
receipt: Receipt,
counts: BTreeMap<CountBucket, BucketCount>,
) -> Result<Receipt, Failure> {
self.rows_before_advance(&at.branch, record)?;
let branch =
open_branch(name, &self.refs, &self.registry).map_err(Failure::branch_commit)?;
for (bucket, counter) in counts {
let key = bucket.key();
let shard = self.writer.shard_for(&key);
if counter.present {
branch.put(shard, key, counter.counted.to_string().into_bytes())
} else {
branch.delete(shard, key)
}
.map_err(Failure::branch)?;
}
branch
.put(self.writer.shard_for(READY_KEY), READY_KEY, READY_VALUE)
.map_err(Failure::branch)?;
let key = receipt_key(
&receipt.branch,
receipt.version.generation,
&receipt.request_id,
);
branch
.put(
self.writer.shard_for(&key),
key,
serde_json::to_vec(&receipt).map_err(Failure::storage)?,
)
.map_err(Failure::branch)?;
#[cfg(test)]
rebuild_probe::before_commit();
self.writer
.commit(
&branch,
&self.registry,
CommitRequest {
durability: CommitDurability::Durable {
refs: &mut self.refs,
},
extra_parents: &[],
timestamp: current_timestamp(),
},
)
.map_err(|error| {
Failure::new(
if self
.refs
.get(name)
.is_some_and(|now| version(now) != at.version)
{
"durability_unknown"
} else {
"write_refused"
},
error.to_string(),
)
})?;
self.publish_record(name)?;
Ok(receipt)
}
fn refuse_count_merge(&self, at: &Captured) -> Result<(), Failure> {
if reads::get_bytes(&self.reader, at, b"b\0merge-candidate-progress")?.is_some() {
return Err(Failure::new(
"counts_merge_pending",
"branch holds a merge candidate",
));
}
let service = self.merge_service_selection()?;
reads::scan_rows(&self.reader, &service, b"j\0merge\0", 1, |parts, bytes| {
let plan: MergePlan = serde_json::from_slice(bytes).map_err(Failure::storage)?;
identifier(&plan.id)?;
for pin in [&plan.source, &plan.target, &plan.origin] {
identifier(&pin.branch)?;
}
if Hash::of(plan.id.as_bytes()).to_string() != parts[0]
|| plan.origin.branch != plan.target.branch
{
return Err(Failure::new(
"merge_corrupt",
"merge plan key differs from identity",
));
}
if plan.state == MergeState::Open
&& [&plan.source, &plan.target, &plan.origin]
.iter()
.any(|pin| pin.branch == at.branch)
{
return Err(Failure::new(
"counts_merge_pending",
"branch has an open merge",
));
}
Ok(())
})
}
pub fn audit_counts(
access: &Access,
reader: &ReadOnlyDatabase,
at: &Captured,
) -> Result<CountsAudit, Failure> {
if !access.can_read_store() {
return Ok(CountsAudit::CannotTell {
reason: CountReason::Withheld,
});
}
if !is_ready(reader, at)? {
return Ok(CountsAudit::CannotTell {
reason: CountReason::CountsNotReady,
});
}
let mut counts = counted_buckets(reader, at)?;
count_rows(reader, at, &mut counts, true)?;
Ok(CountsAudit::Compared {
buckets: counts
.into_iter()
.map(|(bucket, count)| {
let verdict = if count.kept == Some(count.counted) {
AuditVerdict::Agree
} else {
AuditVerdict::Differ {
kept: count.kept,
counted: count.counted,
}
};
AuditedBucket { bucket, verdict }
})
.collect(),
})
}
#[cfg(test)]
pub(in crate::serve) fn arm_count_rebuild_pause(
paused: std::sync::mpsc::Sender<()>,
release: std::sync::mpsc::Receiver<()>,
) {
rebuild_probe::arm(paused, release);
}
}
fn rebuild_buckets(
reader: &ReadOnlyDatabase,
at: &Captured,
) -> Result<BTreeMap<CountBucket, BucketCount>, Failure> {
let mut counts = counted_buckets(reader, at).map_err(|error| count_scan_failure(&error))?;
count_rows(reader, at, &mut counts, false).map_err(|error| count_scan_failure(&error))?;
Ok(counts)
}
fn count_scan_failure(error: &Failure) -> Failure {
Failure::new(
"counts_scan_refused",
format!("{}: {}", error.code, error.message),
)
}
#[cfg(test)]
pub(super) mod rebuild_probe {
use super::Failure;
use std::cell::{Cell, RefCell};
use std::sync::mpsc::{Receiver, Sender};
thread_local! {
static CRASH: Cell<bool> = const { Cell::new(false) };
static PAUSE: RefCell<Option<(Sender<()>, Receiver<()>)>> = const { RefCell::new(None) };
}
pub fn crash_next() {
CRASH.with(|slot| slot.set(true));
}
pub fn arm(paused: Sender<()>, release: Receiver<()>) {
PAUSE.with(|slot| *slot.borrow_mut() = Some((paused, release)));
}
pub fn pause() -> Result<(), Failure> {
if let Some((paused, release)) = PAUSE.with(|slot| slot.borrow_mut().take()) {
paused.send(()).map_err(Failure::storage)?;
release.recv().map_err(Failure::storage)?;
}
Ok(())
}
pub fn before_commit() {
assert!(
!CRASH.with(|slot| slot.replace(false)),
"injected rebuild interruption"
);
}
}
pub(in crate::serve::storage) fn initialize_new_store(
writer: &mut BranchDatabase,
refs: &mut BranchRefStore,
registry: &BranchRegistry,
) -> Result<(), Failure> {
let service = open_branch(SERVICE, refs, registry).map_err(Failure::branch_commit)?;
service
.put(writer.shard_for(READY_KEY), READY_KEY, READY_VALUE)
.map_err(Failure::branch)?;
writer
.commit(
&service,
registry,
CommitRequest {
durability: CommitDurability::Durable { refs },
extra_parents: &[],
timestamp: current_timestamp(),
},
)
.map_err(Failure::branch_commit)?;
Ok(())
}
fn document_count_key(collection: &str) -> Vec<u8> {
format!("k\0d\0{collection}").into_bytes()
}
fn read(reader: &ReadOnlyDatabase, at: &Captured, key: &[u8]) -> Result<Option<Vec<u8>>, Failure> {
#[cfg(test)]
probe::read(key);
reads::get_bytes(reader, at, key)
}
fn is_ready(reader: &ReadOnlyDatabase, at: &Captured) -> Result<bool, Failure> {
match read(reader, at, READY_KEY)? {
None => Ok(false),
Some(value) if value.as_slice() == READY_VALUE => Ok(true),
Some(_) => Err(Failure::new(
"counts_corrupt",
"count readiness format is invalid",
)),
}
}
fn decode(value: &[u8]) -> Result<u64, Failure> {
if value.is_empty()
|| value.len() > 20
|| (value.len() > 1 && value.first() == Some(&b'0'))
|| !value.iter().all(u8::is_ascii_digit)
{
return Err(Failure::new(
"counts_corrupt",
"count is not a canonical decimal u64",
));
}
std::str::from_utf8(value)
.map_err(|error| Failure::new("counts_corrupt", error.to_string()))?
.parse()
.map_err(|error: std::num::ParseIntError| Failure::new("counts_corrupt", error.to_string()))
}
fn required_documents(
reader: &ReadOnlyDatabase,
at: &Captured,
collection: &str,
) -> Result<u64, Failure> {
let value = read(reader, at, &document_count_key(collection))?
.ok_or_else(|| Failure::new("counts_corrupt", "ready collection has no document count"))?;
decode(&value)
}
fn document_exists(
reader: &ReadOnlyDatabase,
at: &Captured,
collection: &str,
id: &str,
) -> Result<bool, Failure> {
let key = revision_key(collection, id);
Ok(read(reader, at, &key)?.is_some())
}
fn increased(count: u64, added: u64) -> Result<Option<u64>, Failure> {
if added == 0 {
return Ok(None);
}
added_count(count, added).map(Some)
}
fn added_count(count: u64, added: u64) -> Result<u64, Failure> {
count
.checked_add(added)
.ok_or_else(|| Failure::new("count_overflow", "object count exceeds u64"))
}
fn removed_count(count: u64, removed: u64) -> Result<u64, Failure> {
count
.checked_sub(removed)
.ok_or_else(|| Failure::new("count_underflow", "removed objects exceed count"))
}
pub(super) fn append(
reader: &ReadOnlyDatabase,
at: &Captured,
change: &Change,
writes: &mut Vec<PreparedWrite>,
) -> Result<(), Failure> {
let collection = match change {
Change::Collection(write) => Some(&write.collection),
Change::Put(write) => Some(&write.collection),
Change::AppendPrune(write) => Some(&write.collection),
Change::Delete(write) => Some(&write.collection),
Change::DeleteMany(delete) => Some(&delete.collection),
Change::PutMany(import) => Some(&import.collection),
Change::EdgePut(_)
| Change::EdgeDelete(_)
| Change::EdgePutMany(_)
| Change::VectorPut(_)
| Change::VectorPutMany(_)
| Change::VectorDelete(_)
| Change::RelationPut(_) => None,
_ => return Ok(()),
};
if !is_ready(reader, at)? {
return Ok(());
}
if let Some(collection) = collection {
append_documents(reader, at, change, collection, writes)?;
}
let mut counters = declared_counters(reader, at, writes)?;
append_objects(reader, at, writes, &mut counters)?;
for (key, counter) in counters {
if counter.write {
writes.push((key, Some(counter.value.to_string().into_bytes())));
}
}
#[cfg(test)]
probe::write_count(writes);
Ok(())
}
fn append_documents(
reader: &ReadOnlyDatabase,
at: &Captured,
change: &Change,
collection: &str,
writes: &mut Vec<PreparedWrite>,
) -> Result<(), Failure> {
let count = match change {
Change::Collection(_) => {
if read(reader, at, &collection_key(collection))?.is_some() {
required_documents(reader, at, collection)?;
None
} else {
if read(reader, at, &document_count_key(collection))?.is_some() {
return Err(Failure::new(
"counts_corrupt",
"document count has no collection",
));
}
Some(0)
}
}
Change::Put(write) => {
let present = document_exists(reader, at, collection, &write.document_id)?;
increased(
required_documents(reader, at, collection)?,
u64::from(!present),
)?
}
Change::Delete(write) => {
if document_exists(reader, at, collection, &write.document_id)? {
Some(removed_count(
required_documents(reader, at, collection)?,
1,
)?)
} else {
None
}
}
Change::DeleteMany(delete) => Some(removed_count(
required_documents(reader, at, collection)?,
u64::try_from(delete.document_ids.len()).map_err(Failure::storage)?,
)?),
Change::AppendPrune(_) => {
let prefix = revision_key(collection, "");
let mut added = 0_u64;
let mut removed = 0_u64;
for (key, value) in writes.iter() {
let Some(id) = key.strip_prefix(prefix.as_slice()) else {
continue;
};
let id = std::str::from_utf8(id).map_err(Failure::storage)?;
let present = document_exists(reader, at, collection, id)?;
if value.is_some() && !present {
added = added_count(added, 1)?;
}
if value.is_none() && present {
removed = added_count(removed, 1)?;
}
}
Some(added_count(
removed_count(required_documents(reader, at, collection)?, removed)?,
added,
)?)
}
Change::PutMany(import) => {
let mut added = 0_u64;
for document in &import.documents {
if !document_exists(reader, at, collection, &document.document_id)? {
added = added.checked_add(1).ok_or_else(|| {
Failure::new("count_overflow", "import delta exceeds u64")
})?;
}
}
increased(required_documents(reader, at, collection)?, added)?
}
_ => return Ok(()),
};
if let Some(count) = count {
writes.push((
document_count_key(collection),
Some(count.to_string().into_bytes()),
));
}
Ok(())
}
#[derive(Default)]
struct Delta {
added: u64,
removed: u64,
existing: bool,
}
struct Counter {
value: u64,
write: bool,
initialized: bool,
}
pub(super) struct CandidateCounts {
ready: bool,
changes: BTreeMap<Vec<u8>, (bool, bool)>,
}
impl CandidateCounts {
pub(super) fn new(reader: &ReadOnlyDatabase, at: &Captured) -> Result<Self, Failure> {
Ok(Self {
ready: is_ready(reader, at)?,
changes: BTreeMap::new(),
})
}
pub(super) fn record(
&mut self,
reader: &ReadOnlyDatabase,
at: &Captured,
key: &[u8],
presence_key: &[u8],
after: bool,
) -> Result<(), Failure> {
if !self.ready {
return Ok(());
}
if let Some(change) = self.changes.get_mut(key) {
change.1 = after;
} else {
let before = read(reader, at, presence_key)?.is_some();
self.changes.insert(key.to_vec(), (before, after));
}
Ok(())
}
pub(super) fn append(
self,
reader: &ReadOnlyDatabase,
at: &Captured,
rows: &mut BTreeMap<Vec<u8>, Option<Vec<u8>>>,
) -> Result<(), Failure> {
if !self.ready {
return Ok(());
}
let mut counters = BTreeMap::new();
for (key, value) in &*rows {
let Some(value) = value else { continue };
if key.starts_with(b"w\0")
&& matches!(
super::embedded::read_field(value)?,
super::embedded::Field::Embedded(_)
)
{
continue;
}
let count_key = if let Some(name) = key.strip_prefix(b"c\0") {
let collection = std::str::from_utf8(name).map_err(Failure::storage)?;
identifier(collection)?;
Some(document_count_key(collection))
} else {
declaration_count_key(key)?
};
let Some(count_key) = count_key else { continue };
let stored = read(reader, at, &count_key)?;
let count = match &stored {
Some(bytes) => decode(bytes)?,
None if read(reader, at, key)?.is_none() => 0,
None => return Err(Failure::new("counts_corrupt", "declared kind has no count")),
};
counters.insert(
count_key,
Counter {
value: count,
write: stored.is_none(),
initialized: stored.is_none(),
},
);
}
let deltas = candidate_deltas(self.changes)?;
for (key, delta) in deltas {
let counter = match counters.entry(key) {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => {
let value = match read(reader, at, entry.key())? {
Some(bytes) => decode(&bytes)?,
None if !delta.existing => {
candidate_missing_count(reader, at, entry.key())?
}
None => {
return Err(Failure::new(
"counts_corrupt",
"present object has no count",
));
}
};
entry.insert(Counter {
value,
write: false,
initialized: false,
})
}
};
if counter.initialized && delta.existing {
return Err(Failure::new(
"counts_corrupt",
"present object has no previous count",
));
}
counter.value = added_count(removed_count(counter.value, delta.removed)?, delta.added)?;
counter.write |= delta.added != 0 || delta.removed != 0;
}
for (key, counter) in counters {
if counter.write {
rows.insert(key, Some(counter.value.to_string().into_bytes()));
}
}
Ok(())
}
}
fn candidate_deltas(
changes: BTreeMap<Vec<u8>, (bool, bool)>,
) -> Result<BTreeMap<Vec<u8>, Delta>, Failure> {
let mut deltas: BTreeMap<Vec<u8>, Delta> = BTreeMap::new();
for (key, (before, after)) in changes {
let count_key = if let Some(suffix) = key.strip_prefix(b"d\0") {
let text = std::str::from_utf8(suffix).map_err(Failure::storage)?;
let (collection, document) = text
.split_once('\0')
.ok_or_else(|| Failure::new("corrupt_key", "document owner has no identifier"))?;
identifier(collection)?;
identifier(document)?;
document_count_key(collection)
} else {
object_count_key(&key)?
.ok_or_else(|| Failure::new("counts_corrupt", "candidate object has no bucket"))?
};
let delta = deltas.entry(count_key).or_default();
delta.existing |= before;
if after && !before {
delta.added = added_count(delta.added, 1)?;
}
if before && !after {
delta.removed = added_count(delta.removed, 1)?;
}
}
Ok(deltas)
}
fn candidate_missing_count(
reader: &ReadOnlyDatabase,
at: &Captured,
key: &[u8],
) -> Result<u64, Failure> {
if let Some(collection) = key.strip_prefix(b"k\0d\0") {
let collection = std::str::from_utf8(collection).map_err(Failure::storage)?;
identifier(collection)?;
if read(reader, at, &collection_key(collection))?.is_some() {
return Err(Failure::new(
"counts_corrupt",
"declared collection has no count",
));
}
Ok(0)
} else {
unregistered_count(reader, at, key)
}
}
fn declaration_count_key(key: &[u8]) -> Result<Option<Vec<u8>>, Failure> {
if let Some(name) = key.strip_prefix(b"n\0") {
let name = std::str::from_utf8(name).map_err(Failure::storage)?;
identifier(name)?;
return Ok(Some(format!("k\0e\0{name}").into_bytes()));
}
if !key.starts_with(b"w\0") {
return Ok(None);
}
let text = std::str::from_utf8(key).map_err(Failure::storage)?;
let parts: Vec<_> = text.split('\0').collect();
let ["w", collection, field] = parts.as_slice() else {
return Err(invalid_vector_key());
};
identifier(collection)?;
identifier(field)?;
Ok(Some(format!("k\0v\0{collection}\0{field}").into_bytes()))
}
fn declared_counters(
reader: &ReadOnlyDatabase,
at: &Captured,
writes: &[PreparedWrite],
) -> Result<BTreeMap<Vec<u8>, Counter>, Failure> {
let mut counters = BTreeMap::new();
for (key, value) in writes {
if value.is_none() {
continue;
}
let Some(count_key) = declaration_count_key(key)? else {
continue;
};
if counters.contains_key(&count_key) {
continue;
}
let declared = read(reader, at, key)?.is_some();
let counter = match read(reader, at, &count_key)? {
Some(value) => Counter {
value: decode(&value)?,
write: false,
initialized: false,
},
None if !declared => Counter {
value: 0,
write: true,
initialized: true,
},
None => return Err(Failure::new("counts_corrupt", "declared kind has no count")),
};
counters.insert(count_key, counter);
}
Ok(counters)
}
fn object_count_key(key: &[u8]) -> Result<Option<Vec<u8>>, Failure> {
if key.starts_with(b"e\0") {
let edge = edges::decode_key(key)?;
return Ok(Some(format!("k\0e\0{}", edge.relation).into_bytes()));
}
if !key.starts_with(b"v\0") {
return Ok(None);
}
let text = std::str::from_utf8(key).map_err(Failure::storage)?;
let mut parts = text.split('\0');
if parts.next() != Some("v") {
return Err(invalid_vector_key());
}
let collection = parts.next().ok_or_else(invalid_vector_key)?;
let field = parts.next().ok_or_else(invalid_vector_key)?;
let document = parts.next().ok_or_else(invalid_vector_key)?;
if parts.next().is_some() {
return Err(invalid_vector_key());
}
identifier(collection)?;
identifier(field)?;
identifier(document)?;
Ok(Some(format!("k\0v\0{collection}\0{field}").into_bytes()))
}
fn invalid_vector_key() -> Failure {
Failure::new("vector_key", "vector key must have three identifier parts")
}
fn append_objects(
reader: &ReadOnlyDatabase,
at: &Captured,
writes: &[PreparedWrite],
counters: &mut BTreeMap<Vec<u8>, Counter>,
) -> Result<(), Failure> {
let mut final_presence = BTreeMap::new();
for (key, value) in writes {
if key.starts_with(b"e\0") || key.starts_with(b"v\0") {
final_presence.insert(key.as_slice(), value.is_some());
}
}
let mut deltas: BTreeMap<Vec<u8>, Delta> = BTreeMap::new();
for (key, present) in final_presence {
let count_key = object_count_key(key)?
.ok_or_else(|| Failure::new("counts_corrupt", "counted object has no bucket"))?;
let existed = read(reader, at, key)?.is_some();
let delta = deltas.entry(count_key).or_default();
delta.existing |= existed;
if present && !existed {
delta.added = delta
.added
.checked_add(1)
.ok_or_else(|| Failure::new("count_overflow", "object delta exceeds u64"))?;
} else if existed && !present {
delta.removed = delta
.removed
.checked_add(1)
.ok_or_else(|| Failure::new("count_overflow", "object delta exceeds u64"))?;
}
}
for (key, delta) in deltas {
if let Some(counter) = counters.get_mut(&key) {
if counter.initialized && delta.existing {
return Err(Failure::new(
"counts_corrupt",
"unregistered kind has stored objects",
));
}
counter.value = added_count(removed_count(counter.value, delta.removed)?, delta.added)?;
counter.write |= delta.added != 0 || delta.removed != 0;
continue;
}
let count = match read(reader, at, &key)? {
Some(value) => decode(&value)?,
None if !delta.existing => unregistered_count(reader, at, &key)?,
None => {
return Err(Failure::new(
"counts_corrupt",
"present object has no count",
));
}
};
let remaining = removed_count(count, delta.removed)?;
let next = added_count(remaining, delta.added)?;
if delta.added != 0 || delta.removed != 0 {
counters.insert(
key,
Counter {
value: next,
write: true,
initialized: false,
},
);
}
}
Ok(())
}
fn unregistered_count(
reader: &ReadOnlyDatabase,
at: &Captured,
key: &[u8],
) -> Result<u64, Failure> {
let declaration = if let Some(name) = key.strip_prefix(b"k\0e\0") {
let mut declaration = b"n\0".to_vec();
declaration.extend_from_slice(name);
declaration
} else if let Some(field) = key.strip_prefix(b"k\0v\0") {
let mut declaration = b"w\0".to_vec();
declaration.extend_from_slice(field);
declaration
} else {
return Err(Failure::new(
"counts_corrupt",
"count key has no declaration family",
));
};
if read(reader, at, &declaration)?.is_some() {
return Err(Failure::new("counts_corrupt", "declared kind has no count"));
}
Ok(0)
}
pub(super) fn sum_rows(
reader: &ReadOnlyDatabase,
at: &Captured,
prefix: &[u8],
parts: usize,
) -> Result<(u64, std::collections::BTreeSet<String>), Failure> {
let mut total = 0;
let mut names = std::collections::BTreeSet::new();
reads::scan_rows(reader, at, prefix, parts, |parts, value| {
let name = parts.join("\0");
if !names.insert(name) {
return Err(Failure::new(
"counts_corrupt",
"count exists in multiple shards",
));
}
total = added_count(total, decode(value)?)?;
Ok(())
})?;
Ok((total, names))
}
#[cfg(test)]
mod probe {
use std::cell::Cell;
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(super) struct Reads {
pub(super) ready: usize,
pub(super) revisions: usize,
pub(super) bodies: usize,
pub(super) objects: usize,
pub(super) counts: usize,
pub(super) count_writes: usize,
}
std::thread_local! {
static READS: Cell<Reads> = const { Cell::new(Reads {
ready: 0, revisions: 0, bodies: 0, objects: 0,
counts: 0, count_writes: 0,
}) };
}
pub(super) fn read(key: &[u8]) {
READS.with(|slot| {
let mut value = slot.get();
if key == super::READY_KEY {
value.ready += 1;
} else if key.starts_with(b"m\0") {
value.revisions += 1;
} else if key.starts_with(b"d\0") {
value.bodies += 1;
} else if key.starts_with(b"e\0") || key.starts_with(b"v\0") {
value.objects += 1;
} else if key.starts_with(b"k\0") {
value.counts += 1;
}
slot.set(value);
});
}
pub(super) fn write_count(writes: &[super::PreparedWrite]) {
READS.with(|slot| {
let mut value = slot.get();
value.count_writes += writes
.iter()
.filter(|(key, _)| key.starts_with(b"k\0"))
.count();
slot.set(value);
});
}
pub(super) fn take() -> Reads {
READS.with(Cell::take)
}
}
#[cfg(test)]
#[path = "counts_tests.rs"]
mod tests;