use std::collections::{HashMap, HashSet, VecDeque};
use std::fmt;
use std::io;
use serde_json::Value;
use crate::errors::{
HaysonCipherError, HaysonPartitionError, HaysonPendingConflictError, HaysonUniqueViolationError,
HaysonValidationError,
};
use crate::registry::composite_index::unique_key_of;
use crate::registry::table_registry::TableRegistry;
use crate::schema::{
partition_of, primary_key_of, validate_record, HaysonRecord, TableSchema, ValidateOptions, UNPARTITIONED,
};
use crate::storage::{serialize_entry, LogFile, LogMeta, LogOp, LogSet, Slot};
use super::pending_map::{PendingMap, ALL_FIELDS};
pub trait LogAppender {
fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>>;
fn append_lines_in(&mut self, part: &str, lines: &[String]) -> io::Result<Vec<Slot>> {
let _ = part;
self.append_lines(lines)
}
fn sync(&mut self) -> io::Result<()>;
}
impl LogAppender for LogFile {
fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
self.append(lines)
}
fn sync(&mut self) -> io::Result<()> {
LogFile::sync(self)
}
}
impl LogAppender for LogSet {
fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
self.append(UNPARTITIONED, lines)
}
fn append_lines_in(&mut self, part: &str, lines: &[String]) -> io::Result<Vec<Slot>> {
self.append(part, lines)
}
fn sync(&mut self) -> io::Result<()> {
LogSet::sync(self)
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum StageError {
Validation(HaysonValidationError),
Unique(HaysonUniqueViolationError),
Conflict(HaysonPendingConflictError),
Partition(HaysonPartitionError),
}
impl fmt::Display for StageError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
StageError::Validation(e) => e.fmt(f),
StageError::Unique(e) => e.fmt(f),
StageError::Conflict(e) => e.fmt(f),
StageError::Partition(e) => e.fmt(f),
}
}
}
impl std::error::Error for StageError {}
#[derive(Debug)]
pub enum CommitError {
Io(io::Error),
Cipher(HaysonCipherError),
}
impl fmt::Display for CommitError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
CommitError::Io(e) => e.fmt(f),
CommitError::Cipher(e) => e.fmt(f),
}
}
}
impl std::error::Error for CommitError {}
impl From<io::Error> for CommitError {
fn from(value: io::Error) -> Self {
CommitError::Io(value)
}
}
impl From<HaysonCipherError> for CommitError {
fn from(value: HaysonCipherError) -> Self {
CommitError::Cipher(value)
}
}
impl From<HaysonValidationError> for StageError {
fn from(value: HaysonValidationError) -> Self {
StageError::Validation(value)
}
}
impl From<HaysonUniqueViolationError> for StageError {
fn from(value: HaysonUniqueViolationError) -> Self {
StageError::Unique(value)
}
}
impl From<HaysonPartitionError> for StageError {
fn from(value: HaysonPartitionError) -> Self {
StageError::Partition(value)
}
}
impl From<HaysonPendingConflictError> for StageError {
fn from(value: HaysonPendingConflictError) -> Self {
StageError::Conflict(value)
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum WriteOutcome {
Put(HaysonRecord),
Deleted(bool),
}
pub type Seal<'a> = dyn Fn(&str, &HaysonRecord, &[String]) -> Result<HaysonRecord, HaysonCipherError> + 'a;
enum EntryKind {
Put(HaysonRecord),
Delete,
}
struct QueuedEntry {
key: String,
fields: Vec<String>,
kind: EntryKind,
sealed: Vec<String>,
part: String,
}
pub struct WriteQueue {
table: String,
max_batch: usize,
pending: PendingMap,
queued: VecDeque<QueuedEntry>,
}
impl WriteQueue {
pub fn new(table: impl Into<String>) -> Self {
Self::with_max_batch(table, 256)
}
pub fn with_max_batch(table: impl Into<String>, max_batch: usize) -> Self {
WriteQueue {
table: table.into(),
max_batch: max_batch.max(1),
pending: PendingMap::new(),
queued: VecDeque::new(),
}
}
pub fn depth(&self) -> usize {
self.queued.len()
}
pub fn pending(&self, key: &str) -> HashSet<String> {
self.pending.pending(key)
}
pub fn queued_keys(&self) -> impl Iterator<Item = &str> {
self.queued.iter().map(|entry| entry.key.as_str())
}
pub fn stage_put(
&mut self,
schema: &TableSchema,
registry: &TableRegistry,
document: &Value,
force: bool,
) -> Result<String, StageError> {
self.stage_entry(schema, registry, document, None, Vec::new(), force)
}
pub fn stage_put_carrying(
&mut self,
schema: &TableSchema,
registry: &TableRegistry,
document: &Value,
sealed: Vec<String>,
force: bool,
) -> Result<String, StageError> {
self.stage_entry(schema, registry, document, None, sealed, force)
}
pub fn stage_put_claiming(
&mut self,
schema: &TableSchema,
registry: &TableRegistry,
document: &Value,
claim: &[String],
force: bool,
) -> Result<String, StageError> {
self.stage_entry(schema, registry, document, Some(claim), Vec::new(), force)
}
fn stage_entry(
&mut self,
schema: &TableSchema,
registry: &TableRegistry,
document: &Value,
claim: Option<&[String]>,
sealed: Vec<String>,
force: bool,
) -> Result<String, StageError> {
let record = validate_record(schema, document, ValidateOptions::default())?;
let key = primary_key_of(schema, &record)?;
assert_unique(schema, registry, &key, &record)?;
let part = partition_of(schema, &record)?;
assert_partition_stable(schema, registry, &key, &part)?;
let fields: Vec<String> = match claim {
Some(claim) => claim.to_vec(),
None => record.keys().cloned().collect(),
};
self.reserve(&key, &fields, force)?;
self.queued.push_back(QueuedEntry {
key: key.clone(),
fields,
kind: EntryKind::Put(record),
sealed,
part,
});
Ok(key)
}
pub fn last_put_record(&self) -> Option<&HaysonRecord> {
match self.queued.back() {
Some(QueuedEntry {
kind: EntryKind::Put(record),
..
}) => Some(record),
_ => None,
}
}
pub fn stage_delete(
&mut self,
key: impl Into<String>,
force: bool,
part: &str,
) -> Result<(), HaysonPendingConflictError> {
let key = key.into();
let fields = vec![ALL_FIELDS.to_string()];
self.reserve(&key, &fields, force)?;
self.queued.push_back(QueuedEntry {
key,
fields,
kind: EntryKind::Delete,
sealed: Vec::new(),
part: part.to_string(),
});
Ok(())
}
fn reserve(&mut self, key: &str, fields: &[String], force: bool) -> Result<(), HaysonPendingConflictError> {
let conflicts = self.pending.conflicts_of(key, fields);
if !conflicts.is_empty() && !force {
return Err(HaysonPendingConflictError {
table: self.table.clone(),
key: key.to_string(),
fields: conflicts,
});
}
self.pending.claim(key, fields);
Ok(())
}
#[allow(dead_code)]
pub fn commit<L: LogAppender>(
&mut self,
log: &mut L,
registry: &mut TableRegistry,
seq: &mut u64,
sync: bool,
) -> io::Result<Vec<WriteOutcome>> {
self.commit_sealed(log, registry, seq, sync, &|_key, record, _sealed| Ok(record.clone()))
.map_err(|error| match error {
CommitError::Io(error) => error,
CommitError::Cipher(error) => io::Error::other(error.to_string()),
})
}
pub fn commit_sealed<L: LogAppender>(
&mut self,
log: &mut L,
registry: &mut TableRegistry,
seq: &mut u64,
sync: bool,
seal: &Seal<'_>,
) -> Result<Vec<WriteOutcome>, CommitError> {
let mut results = Vec::with_capacity(self.queued.len());
while !self.queued.is_empty() {
let take = self.max_batch.min(self.queued.len());
let batch: Vec<QueuedEntry> = self.queued.drain(..take).collect();
let now = chrono::Utc::now().timestamp_millis();
let mut lines = Vec::with_capacity(batch.len());
let mut stamps = Vec::with_capacity(batch.len());
for entry in &batch {
*seq += 1;
stamps.push(*seq);
let line = match &entry.kind {
EntryKind::Put(record) => {
let sealed = seal(&entry.key, record, &entry.sealed)?;
serialize_entry(
&LogMeta {
key: entry.key.clone(),
op: LogOp::Put,
seq: *seq,
ts: now,
},
Some(&sealed),
)
}
EntryKind::Delete => serialize_entry(
&LogMeta {
key: entry.key.clone(),
op: LogOp::Del,
seq: *seq,
ts: now,
},
None,
),
};
lines.push(line);
}
let slots = append_by_part(log, &batch, &lines)?;
if sync {
log.sync()?;
}
for ((entry, slot), commit_seq) in batch.into_iter().zip(slots).zip(stamps) {
match entry.kind {
EntryKind::Put(record) => {
registry.apply_put_in(&entry.key, &record, slot, commit_seq, &entry.part);
results.push(WriteOutcome::Put(record));
}
EntryKind::Delete => {
let existed = registry.slot_of(&entry.key).is_some();
registry.apply_delete_in(&entry.key, slot.length, commit_seq, &entry.part);
results.push(WriteOutcome::Deleted(existed));
}
}
self.pending.release(&entry.key, &entry.fields);
}
}
Ok(results)
}
}
fn append_by_part<L: LogAppender + ?Sized>(
log: &mut L,
batch: &[QueuedEntry],
lines: &[String],
) -> io::Result<Vec<Slot>> {
let mut grouped: Vec<(&str, Vec<usize>)> = Vec::new();
for (position, entry) in batch.iter().enumerate() {
match grouped.iter_mut().find(|(part, _)| *part == entry.part) {
Some((_, positions)) => positions.push(position),
None => grouped.push((&entry.part, vec![position])),
}
}
let mut slots: Vec<Option<Slot>> = vec![None; batch.len()];
for (part, positions) in grouped {
let payload: Vec<String> = positions.iter().map(|position| lines[*position].clone()).collect();
let appended = log.append_lines_in(part, &payload)?;
for (position, slot) in positions.into_iter().zip(appended) {
slots[position] = Some(slot);
}
}
slots
.into_iter()
.map(|slot| slot.ok_or_else(|| io::Error::other("the log returned fewer slots than lines appended")))
.collect()
}
pub(crate) fn assert_partition_stable(
schema: &TableSchema,
registry: &TableRegistry,
key: &str,
part: &str,
) -> Result<(), HaysonPartitionError> {
if part == UNPARTITIONED {
return Ok(());
}
let Some(held) = registry.part_of(key) else {
return Ok(());
};
if held == part {
return Ok(());
}
Err(HaysonPartitionError::new(
schema.table.clone(),
schema.partition_by.clone().unwrap_or_default(),
format!("key \"{key}\" already lives in partition \"{held}\" — delete it before writing it into \"{part}\""),
))
}
pub(crate) type UnitShadow = HashMap<String, Option<HaysonRecord>>;
pub(crate) fn assert_unique(
schema: &TableSchema,
registry: &TableRegistry,
key: &str,
record: &HaysonRecord,
) -> Result<(), HaysonUniqueViolationError> {
assert_unique_in_unit(schema, registry, key, record, &UnitShadow::new())
}
fn virtual_path_of<'a>(record: &'a HaysonRecord, field: &str) -> Option<&'a str> {
record.get(field)?.as_object()?.get("virtualPath")?.as_str()
}
fn violation(schema: &TableSchema, field: String, holder: String) -> HaysonUniqueViolationError {
HaysonUniqueViolationError {
table: schema.table.clone(),
field,
holder,
}
}
fn shadow_released(
shadow: &UnitShadow,
holder: &str,
mine: &str,
of: impl Fn(&HaysonRecord) -> Option<String>,
) -> bool {
match shadow.get(holder) {
None => false,
Some(None) => true,
Some(Some(record)) => of(record).as_deref() != Some(mine),
}
}
pub(crate) fn assert_unique_in_unit(
schema: &TableSchema,
registry: &TableRegistry,
key: &str,
record: &HaysonRecord,
shadow: &UnitShadow,
) -> Result<(), HaysonUniqueViolationError> {
for index in registry.indexes() {
if !index.unique {
continue;
}
if index.fields.len() == 1 && index.fields[0] == schema.primary_key {
continue;
}
let unique_key = |candidate: &HaysonRecord| unique_key_of(candidate, &index.fields);
let Some(mine) = unique_key(record) else { continue };
for (other_key, other) in shadow {
if other_key == key {
continue;
}
let Some(other) = other else { continue };
if unique_key(other).as_deref() == Some(mine.as_str()) {
return Err(violation(schema, index.fields.join(", "), other_key.clone()));
}
}
if let Some(holder) = registry.unique_holder_of(index, record) {
if holder != key && !shadow_released(shadow, &holder, &mine, unique_key) {
return Err(violation(schema, index.fields.join(", "), holder));
}
}
}
for field in registry.virtual_path_fields() {
let path_of = |candidate: &HaysonRecord| virtual_path_of(candidate, field).map(str::to_string);
let Some(mine) = virtual_path_of(record, field) else {
continue;
};
for (other_key, other) in shadow {
if other_key == key {
continue;
}
let Some(other) = other else { continue };
if virtual_path_of(other, field) == Some(mine) {
return Err(violation(schema, field.clone(), other_key.clone()));
}
}
if let Some(holder) = registry.virtual_path_holder(field, mine) {
if holder != key && !shadow_released(shadow, &holder, mine, path_of) {
return Err(violation(schema, field.clone(), holder));
}
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::schema::parse_schema;
use serde_json::json;
fn users_schema() -> TableSchema {
parse_schema(&json!({
"schemaVersion": 1,
"table": "users",
"primaryKey": "id",
"fields": {
"id": { "type": "string" },
"name": { "type": "string", "required": true },
"email": { "type": "string", "unique": true },
},
}))
.unwrap()
}
struct CountingLog {
calls: usize,
synced: usize,
appended: Vec<String>,
}
impl CountingLog {
fn new() -> Self {
CountingLog {
calls: 0,
synced: 0,
appended: Vec::new(),
}
}
}
impl LogAppender for CountingLog {
fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
self.calls += 1;
let start = self.appended.len() as u64;
let slots = lines
.iter()
.enumerate()
.map(|(i, line)| Slot {
offset: start + i as u64,
length: line.len() as u64,
})
.collect();
self.appended.extend(lines.iter().cloned());
Ok(slots)
}
fn sync(&mut self) -> io::Result<()> {
self.synced += 1;
Ok(())
}
}
#[test]
fn commits_every_entry_queued_since_the_last_call_in_one_append() {
let schema = users_schema();
let mut registry = TableRegistry::new(schema.clone());
let mut queue = WriteQueue::new("users");
let mut log = CountingLog::new();
let mut seq = 0u64;
queue
.stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
.unwrap();
queue
.stage_put(&schema, ®istry, &json!({ "id": "u2", "name": "Piotr" }), false)
.unwrap();
assert_eq!(queue.depth(), 2);
let results = queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
assert_eq!(log.calls, 1);
assert_eq!(log.synced, 1);
assert_eq!(results.len(), 2);
assert_eq!(queue.depth(), 0);
assert_eq!(registry.records(), 2);
}
#[test]
fn a_second_write_to_the_same_field_conflicts_unless_forced() {
let schema = users_schema();
let registry = TableRegistry::new(schema.clone());
let mut queue = WriteQueue::new("users");
queue
.stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
.unwrap();
let err = queue
.stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Zosia" }), false)
.unwrap_err();
assert!(matches!(err, StageError::Conflict(_)));
queue
.stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Zosia" }), true)
.unwrap();
assert_eq!(queue.depth(), 2);
}
#[test]
fn a_delete_claims_every_field_and_conflicts_with_a_write_on_any_of_them() {
let schema = users_schema();
let registry = TableRegistry::new(schema.clone());
let mut queue = WriteQueue::new("users");
queue.stage_delete("u1", false, "").unwrap();
let err = queue
.stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
.unwrap_err();
assert!(matches!(err, StageError::Conflict(_)));
}
#[test]
fn releases_the_claim_only_after_the_batch_commits() {
let schema = users_schema();
let mut registry = TableRegistry::new(schema.clone());
let mut queue = WriteQueue::new("users");
let mut log = CountingLog::new();
let mut seq = 0u64;
queue
.stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
.unwrap();
assert_eq!(queue.pending("u1"), HashSet::from(["id".into(), "name".into()]));
queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
assert!(queue.pending("u1").is_empty());
}
#[test]
fn rejects_a_second_record_claiming_a_unique_field_already_held() {
let schema = users_schema();
let mut registry = TableRegistry::new(schema.clone());
let mut queue = WriteQueue::new("users");
let mut log = CountingLog::new();
let mut seq = 0u64;
queue
.stage_put(
&schema,
®istry,
&json!({ "id": "u1", "name": "Ania", "email": "a@x.io" }),
false,
)
.unwrap();
queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
let err = queue
.stage_put(
&schema,
®istry,
&json!({ "id": "u2", "name": "Zosia", "email": "a@x.io" }),
false,
)
.unwrap_err();
assert!(matches!(err, StageError::Unique(_)));
}
fn docs_schema() -> TableSchema {
parse_schema(&json!({
"schemaVersion": 1,
"table": "docs",
"primaryKey": "id",
"fields": {
"id": { "type": "string" },
"page": { "type": "markdown" },
},
}))
.unwrap()
}
#[test]
fn rejects_a_second_record_claiming_an_already_taken_markdown_virtual_path() {
let schema = docs_schema();
let mut registry = TableRegistry::new(schema.clone());
let mut queue = WriteQueue::new("docs");
let mut log = CountingLog::new();
let mut seq = 0u64;
queue
.stage_put(
&schema,
®istry,
&json!({ "id": "d1", "page": { "content": "a", "virtualPath": "/docs/a.md" } }),
false,
)
.unwrap();
queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
let err = queue
.stage_put(
&schema,
®istry,
&json!({ "id": "d2", "page": { "content": "b", "virtualPath": "/docs/a.md" } }),
false,
)
.unwrap_err();
assert!(matches!(err, StageError::Unique(_)));
}
#[test]
fn dead_bytes_grow_when_an_existing_key_is_overwritten_or_removed() {
let schema = users_schema();
let mut registry = TableRegistry::new(schema.clone());
let mut queue = WriteQueue::new("users");
let mut log = CountingLog::new();
let mut seq = 0u64;
queue
.stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
.unwrap();
queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
assert_eq!(registry.stats().dead_bytes, 0);
queue
.stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Zosia" }), false)
.unwrap();
queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
let after_overwrite = registry.stats().dead_bytes;
assert!(after_overwrite > 0);
queue.stage_delete("u1", false, "").unwrap();
queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
assert!(registry.stats().dead_bytes > after_overwrite);
assert_eq!(registry.records(), 0);
}
#[test]
fn splits_a_batch_larger_than_max_batch_into_several_appends() {
let schema = users_schema();
let mut registry = TableRegistry::new(schema.clone());
let mut queue = WriteQueue::with_max_batch("users", 2);
let mut log = CountingLog::new();
let mut seq = 0u64;
for i in 0..5 {
queue
.stage_put(
&schema,
®istry,
&json!({ "id": format!("u{i}"), "name": "x" }),
false,
)
.unwrap();
}
let results = queue.commit(&mut log, &mut registry, &mut seq, false).unwrap();
assert_eq!(results.len(), 5);
assert_eq!(log.calls, 3);
assert_eq!(log.synced, 0);
}
}