1use std::collections::{HashMap, HashSet, VecDeque};
2use std::fmt;
3use std::io;
4
5use serde_json::Value;
6
7use crate::errors::{
8 HaysonCipherError, HaysonPartitionError, HaysonPendingConflictError, HaysonUniqueViolationError,
9 HaysonValidationError,
10};
11use crate::registry::composite_index::unique_key_of;
12use crate::registry::table_registry::TableRegistry;
13use crate::schema::{
14 partition_of, primary_key_of, validate_record, HaysonRecord, TableSchema, ValidateOptions, UNPARTITIONED,
15};
16use crate::storage::{serialize_entry, LogFile, LogMeta, LogOp, LogSet, Slot};
17
18use super::pending_map::{PendingMap, ALL_FIELDS};
19
20pub trait LogAppender {
21 fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>>;
22
23 fn append_lines_in(&mut self, part: &str, lines: &[String]) -> io::Result<Vec<Slot>> {
27 let _ = part;
28 self.append_lines(lines)
29 }
30
31 fn sync(&mut self) -> io::Result<()>;
32}
33
34impl LogAppender for LogFile {
35 fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
36 self.append(lines)
37 }
38
39 fn sync(&mut self) -> io::Result<()> {
40 LogFile::sync(self)
41 }
42}
43
44impl LogAppender for LogSet {
45 fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
46 self.append(UNPARTITIONED, lines)
47 }
48
49 fn append_lines_in(&mut self, part: &str, lines: &[String]) -> io::Result<Vec<Slot>> {
50 self.append(part, lines)
51 }
52
53 fn sync(&mut self) -> io::Result<()> {
54 LogSet::sync(self)
55 }
56}
57
58#[derive(Debug, Clone, PartialEq)]
59pub enum StageError {
60 Validation(HaysonValidationError),
61 Unique(HaysonUniqueViolationError),
62 Conflict(HaysonPendingConflictError),
63 Partition(HaysonPartitionError),
64}
65
66impl fmt::Display for StageError {
67 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
68 match self {
69 StageError::Validation(e) => e.fmt(f),
70 StageError::Unique(e) => e.fmt(f),
71 StageError::Conflict(e) => e.fmt(f),
72 StageError::Partition(e) => e.fmt(f),
73 }
74 }
75}
76
77impl std::error::Error for StageError {}
78
79#[derive(Debug)]
82pub enum CommitError {
83 Io(io::Error),
84 Cipher(HaysonCipherError),
85}
86
87impl fmt::Display for CommitError {
88 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
89 match self {
90 CommitError::Io(e) => e.fmt(f),
91 CommitError::Cipher(e) => e.fmt(f),
92 }
93 }
94}
95
96impl std::error::Error for CommitError {}
97
98impl From<io::Error> for CommitError {
99 fn from(value: io::Error) -> Self {
100 CommitError::Io(value)
101 }
102}
103
104impl From<HaysonCipherError> for CommitError {
105 fn from(value: HaysonCipherError) -> Self {
106 CommitError::Cipher(value)
107 }
108}
109
110impl From<HaysonValidationError> for StageError {
111 fn from(value: HaysonValidationError) -> Self {
112 StageError::Validation(value)
113 }
114}
115
116impl From<HaysonUniqueViolationError> for StageError {
117 fn from(value: HaysonUniqueViolationError) -> Self {
118 StageError::Unique(value)
119 }
120}
121
122impl From<HaysonPartitionError> for StageError {
123 fn from(value: HaysonPartitionError) -> Self {
124 StageError::Partition(value)
125 }
126}
127
128impl From<HaysonPendingConflictError> for StageError {
129 fn from(value: HaysonPendingConflictError) -> Self {
130 StageError::Conflict(value)
131 }
132}
133
134#[derive(Debug, Clone, PartialEq)]
135pub enum WriteOutcome {
136 Put(HaysonRecord),
137 Deleted(bool),
138}
139
140pub type Seal<'a> = dyn Fn(&str, &HaysonRecord, &[String]) -> Result<HaysonRecord, HaysonCipherError> + 'a;
144
145enum EntryKind {
146 Put(HaysonRecord),
147 Delete,
148}
149
150struct QueuedEntry {
151 key: String,
152 fields: Vec<String>,
153 kind: EntryKind,
154 sealed: Vec<String>,
159 part: String,
161}
162
163pub struct WriteQueue {
167 table: String,
168 max_batch: usize,
169 pending: PendingMap,
170 queued: VecDeque<QueuedEntry>,
171}
172
173impl WriteQueue {
174 pub fn new(table: impl Into<String>) -> Self {
175 Self::with_max_batch(table, 256)
176 }
177
178 pub fn with_max_batch(table: impl Into<String>, max_batch: usize) -> Self {
179 WriteQueue {
180 table: table.into(),
181 max_batch: max_batch.max(1),
182 pending: PendingMap::new(),
183 queued: VecDeque::new(),
184 }
185 }
186
187 pub fn depth(&self) -> usize {
188 self.queued.len()
189 }
190
191 pub fn pending(&self, key: &str) -> HashSet<String> {
192 self.pending.pending(key)
193 }
194
195 pub fn queued_keys(&self) -> impl Iterator<Item = &str> {
199 self.queued.iter().map(|entry| entry.key.as_str())
200 }
201
202 pub fn stage_put(
206 &mut self,
207 schema: &TableSchema,
208 registry: &TableRegistry,
209 document: &Value,
210 force: bool,
211 ) -> Result<String, StageError> {
212 self.stage_entry(schema, registry, document, None, Vec::new(), force)
213 }
214
215 pub fn stage_put_carrying(
220 &mut self,
221 schema: &TableSchema,
222 registry: &TableRegistry,
223 document: &Value,
224 sealed: Vec<String>,
225 force: bool,
226 ) -> Result<String, StageError> {
227 self.stage_entry(schema, registry, document, None, sealed, force)
228 }
229
230 pub fn stage_put_claiming(
236 &mut self,
237 schema: &TableSchema,
238 registry: &TableRegistry,
239 document: &Value,
240 claim: &[String],
241 force: bool,
242 ) -> Result<String, StageError> {
243 self.stage_entry(schema, registry, document, Some(claim), Vec::new(), force)
244 }
245
246 fn stage_entry(
247 &mut self,
248 schema: &TableSchema,
249 registry: &TableRegistry,
250 document: &Value,
251 claim: Option<&[String]>,
252 sealed: Vec<String>,
253 force: bool,
254 ) -> Result<String, StageError> {
255 let record = validate_record(schema, document, ValidateOptions::default())?;
256 let key = primary_key_of(schema, &record)?;
257 assert_unique(schema, registry, &key, &record)?;
258 let part = partition_of(schema, &record)?;
259 assert_partition_stable(schema, registry, &key, &part)?;
260 let fields: Vec<String> = match claim {
261 Some(claim) => claim.to_vec(),
262 None => record.keys().cloned().collect(),
263 };
264 self.reserve(&key, &fields, force)?;
265 self.queued.push_back(QueuedEntry {
266 key: key.clone(),
267 fields,
268 kind: EntryKind::Put(record),
269 sealed,
270 part,
271 });
272 Ok(key)
273 }
274
275 pub fn last_put_record(&self) -> Option<&HaysonRecord> {
279 match self.queued.back() {
280 Some(QueuedEntry {
281 kind: EntryKind::Put(record),
282 ..
283 }) => Some(record),
284 _ => None,
285 }
286 }
287
288 pub fn stage_delete(
292 &mut self,
293 key: impl Into<String>,
294 force: bool,
295 part: &str,
296 ) -> Result<(), HaysonPendingConflictError> {
297 let key = key.into();
298 let fields = vec![ALL_FIELDS.to_string()];
299 self.reserve(&key, &fields, force)?;
300 self.queued.push_back(QueuedEntry {
301 key,
302 fields,
303 kind: EntryKind::Delete,
304 sealed: Vec::new(),
305 part: part.to_string(),
306 });
307 Ok(())
308 }
309
310 fn reserve(&mut self, key: &str, fields: &[String], force: bool) -> Result<(), HaysonPendingConflictError> {
311 let conflicts = self.pending.conflicts_of(key, fields);
312 if !conflicts.is_empty() && !force {
313 return Err(HaysonPendingConflictError {
314 table: self.table.clone(),
315 key: key.to_string(),
316 fields: conflicts,
317 });
318 }
319 self.pending.claim(key, fields);
320 Ok(())
321 }
322
323 #[allow(dead_code)]
329 pub fn commit<L: LogAppender>(
330 &mut self,
331 log: &mut L,
332 registry: &mut TableRegistry,
333 seq: &mut u64,
334 sync: bool,
335 ) -> io::Result<Vec<WriteOutcome>> {
336 self.commit_sealed(log, registry, seq, sync, &|_key, record, _sealed| Ok(record.clone()))
337 .map_err(|error| match error {
338 CommitError::Io(error) => error,
339 CommitError::Cipher(error) => io::Error::other(error.to_string()),
340 })
341 }
342
343 pub fn commit_sealed<L: LogAppender>(
350 &mut self,
351 log: &mut L,
352 registry: &mut TableRegistry,
353 seq: &mut u64,
354 sync: bool,
355 seal: &Seal<'_>,
356 ) -> Result<Vec<WriteOutcome>, CommitError> {
357 let mut results = Vec::with_capacity(self.queued.len());
358 while !self.queued.is_empty() {
359 let take = self.max_batch.min(self.queued.len());
360 let batch: Vec<QueuedEntry> = self.queued.drain(..take).collect();
361 let now = chrono::Utc::now().timestamp_millis();
362
363 let mut lines = Vec::with_capacity(batch.len());
364 let mut stamps = Vec::with_capacity(batch.len());
365 for entry in &batch {
366 *seq += 1;
367 stamps.push(*seq);
368 let line = match &entry.kind {
369 EntryKind::Put(record) => {
370 let sealed = seal(&entry.key, record, &entry.sealed)?;
371 serialize_entry(
372 &LogMeta {
373 key: entry.key.clone(),
374 op: LogOp::Put,
375 seq: *seq,
376 ts: now,
377 },
378 Some(&sealed),
379 )
380 }
381 EntryKind::Delete => serialize_entry(
382 &LogMeta {
383 key: entry.key.clone(),
384 op: LogOp::Del,
385 seq: *seq,
386 ts: now,
387 },
388 None,
389 ),
390 };
391 lines.push(line);
392 }
393
394 let slots = append_by_part(log, &batch, &lines)?;
395 if sync {
396 log.sync()?;
397 }
398
399 for ((entry, slot), commit_seq) in batch.into_iter().zip(slots).zip(stamps) {
400 match entry.kind {
401 EntryKind::Put(record) => {
402 registry.apply_put_in(&entry.key, &record, slot, commit_seq, &entry.part);
403 results.push(WriteOutcome::Put(record));
404 }
405 EntryKind::Delete => {
406 let existed = registry.slot_of(&entry.key).is_some();
407 registry.apply_delete_in(&entry.key, slot.length, commit_seq, &entry.part);
408 results.push(WriteOutcome::Deleted(existed));
409 }
410 }
411 self.pending.release(&entry.key, &entry.fields);
412 }
413 }
414 Ok(results)
415 }
416}
417
418fn append_by_part<L: LogAppender + ?Sized>(
423 log: &mut L,
424 batch: &[QueuedEntry],
425 lines: &[String],
426) -> io::Result<Vec<Slot>> {
427 let mut grouped: Vec<(&str, Vec<usize>)> = Vec::new();
428 for (position, entry) in batch.iter().enumerate() {
429 match grouped.iter_mut().find(|(part, _)| *part == entry.part) {
430 Some((_, positions)) => positions.push(position),
431 None => grouped.push((&entry.part, vec![position])),
432 }
433 }
434
435 let mut slots: Vec<Option<Slot>> = vec![None; batch.len()];
436 for (part, positions) in grouped {
437 let payload: Vec<String> = positions.iter().map(|position| lines[*position].clone()).collect();
438 let appended = log.append_lines_in(part, &payload)?;
439 for (position, slot) in positions.into_iter().zip(appended) {
440 slots[position] = Some(slot);
441 }
442 }
443 slots
444 .into_iter()
445 .map(|slot| slot.ok_or_else(|| io::Error::other("the log returned fewer slots than lines appended")))
446 .collect()
447}
448
449pub(crate) fn assert_partition_stable(
453 schema: &TableSchema,
454 registry: &TableRegistry,
455 key: &str,
456 part: &str,
457) -> Result<(), HaysonPartitionError> {
458 if part == UNPARTITIONED {
459 return Ok(());
460 }
461 let Some(held) = registry.part_of(key) else {
462 return Ok(());
463 };
464 if held == part {
465 return Ok(());
466 }
467 Err(HaysonPartitionError::new(
468 schema.table.clone(),
469 schema.partition_by.clone().unwrap_or_default(),
470 format!("key \"{key}\" already lives in partition \"{held}\" — delete it before writing it into \"{part}\""),
471 ))
472}
473
474pub(crate) type UnitShadow = HashMap<String, Option<HaysonRecord>>;
478
479pub(crate) fn assert_unique(
480 schema: &TableSchema,
481 registry: &TableRegistry,
482 key: &str,
483 record: &HaysonRecord,
484) -> Result<(), HaysonUniqueViolationError> {
485 assert_unique_in_unit(schema, registry, key, record, &UnitShadow::new())
486}
487
488fn virtual_path_of<'a>(record: &'a HaysonRecord, field: &str) -> Option<&'a str> {
489 record.get(field)?.as_object()?.get("virtualPath")?.as_str()
490}
491
492fn violation(schema: &TableSchema, field: String, holder: String) -> HaysonUniqueViolationError {
493 HaysonUniqueViolationError {
494 table: schema.table.clone(),
495 field,
496 holder,
497 }
498}
499
500fn shadow_released(
504 shadow: &UnitShadow,
505 holder: &str,
506 mine: &str,
507 of: impl Fn(&HaysonRecord) -> Option<String>,
508) -> bool {
509 match shadow.get(holder) {
510 None => false,
511 Some(None) => true,
512 Some(Some(record)) => of(record).as_deref() != Some(mine),
513 }
514}
515
516pub(crate) fn assert_unique_in_unit(
521 schema: &TableSchema,
522 registry: &TableRegistry,
523 key: &str,
524 record: &HaysonRecord,
525 shadow: &UnitShadow,
526) -> Result<(), HaysonUniqueViolationError> {
527 for index in registry.indexes() {
528 if !index.unique {
529 continue;
530 }
531 if index.fields.len() == 1 && index.fields[0] == schema.primary_key {
532 continue;
533 }
534 let unique_key = |candidate: &HaysonRecord| unique_key_of(candidate, &index.fields);
535 let Some(mine) = unique_key(record) else { continue };
536 for (other_key, other) in shadow {
537 if other_key == key {
538 continue;
539 }
540 let Some(other) = other else { continue };
541 if unique_key(other).as_deref() == Some(mine.as_str()) {
542 return Err(violation(schema, index.fields.join(", "), other_key.clone()));
543 }
544 }
545 if let Some(holder) = registry.unique_holder_of(index, record) {
546 if holder != key && !shadow_released(shadow, &holder, &mine, unique_key) {
547 return Err(violation(schema, index.fields.join(", "), holder));
548 }
549 }
550 }
551 for field in registry.virtual_path_fields() {
552 let path_of = |candidate: &HaysonRecord| virtual_path_of(candidate, field).map(str::to_string);
553 let Some(mine) = virtual_path_of(record, field) else {
554 continue;
555 };
556 for (other_key, other) in shadow {
557 if other_key == key {
558 continue;
559 }
560 let Some(other) = other else { continue };
561 if virtual_path_of(other, field) == Some(mine) {
562 return Err(violation(schema, field.clone(), other_key.clone()));
563 }
564 }
565 if let Some(holder) = registry.virtual_path_holder(field, mine) {
566 if holder != key && !shadow_released(shadow, &holder, mine, path_of) {
567 return Err(violation(schema, field.clone(), holder));
568 }
569 }
570 }
571 Ok(())
572}
573
574#[cfg(test)]
575mod tests {
576 use super::*;
577 use crate::schema::parse_schema;
578 use serde_json::json;
579
580 fn users_schema() -> TableSchema {
581 parse_schema(&json!({
582 "schemaVersion": 1,
583 "table": "users",
584 "primaryKey": "id",
585 "fields": {
586 "id": { "type": "string" },
587 "name": { "type": "string", "required": true },
588 "email": { "type": "string", "unique": true },
589 },
590 }))
591 .unwrap()
592 }
593
594 struct CountingLog {
595 calls: usize,
596 synced: usize,
597 appended: Vec<String>,
598 }
599
600 impl CountingLog {
601 fn new() -> Self {
602 CountingLog {
603 calls: 0,
604 synced: 0,
605 appended: Vec::new(),
606 }
607 }
608 }
609
610 impl LogAppender for CountingLog {
611 fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
612 self.calls += 1;
613 let start = self.appended.len() as u64;
614 let slots = lines
615 .iter()
616 .enumerate()
617 .map(|(i, line)| Slot {
618 offset: start + i as u64,
619 length: line.len() as u64,
620 })
621 .collect();
622 self.appended.extend(lines.iter().cloned());
623 Ok(slots)
624 }
625
626 fn sync(&mut self) -> io::Result<()> {
627 self.synced += 1;
628 Ok(())
629 }
630 }
631
632 #[test]
633 fn commits_every_entry_queued_since_the_last_call_in_one_append() {
634 let schema = users_schema();
635 let mut registry = TableRegistry::new(schema.clone());
636 let mut queue = WriteQueue::new("users");
637 let mut log = CountingLog::new();
638 let mut seq = 0u64;
639
640 queue
641 .stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
642 .unwrap();
643 queue
644 .stage_put(&schema, ®istry, &json!({ "id": "u2", "name": "Piotr" }), false)
645 .unwrap();
646 assert_eq!(queue.depth(), 2);
647
648 let results = queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
649
650 assert_eq!(log.calls, 1);
651 assert_eq!(log.synced, 1);
652 assert_eq!(results.len(), 2);
653 assert_eq!(queue.depth(), 0);
654 assert_eq!(registry.records(), 2);
655 }
656
657 #[test]
658 fn a_second_write_to_the_same_field_conflicts_unless_forced() {
659 let schema = users_schema();
660 let registry = TableRegistry::new(schema.clone());
661 let mut queue = WriteQueue::new("users");
662
663 queue
664 .stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
665 .unwrap();
666 let err = queue
667 .stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Zosia" }), false)
668 .unwrap_err();
669 assert!(matches!(err, StageError::Conflict(_)));
670
671 queue
672 .stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Zosia" }), true)
673 .unwrap();
674 assert_eq!(queue.depth(), 2);
675 }
676
677 #[test]
678 fn a_delete_claims_every_field_and_conflicts_with_a_write_on_any_of_them() {
679 let schema = users_schema();
680 let registry = TableRegistry::new(schema.clone());
681 let mut queue = WriteQueue::new("users");
682
683 queue.stage_delete("u1", false, "").unwrap();
684 let err = queue
685 .stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
686 .unwrap_err();
687 assert!(matches!(err, StageError::Conflict(_)));
688 }
689
690 #[test]
691 fn releases_the_claim_only_after_the_batch_commits() {
692 let schema = users_schema();
693 let mut registry = TableRegistry::new(schema.clone());
694 let mut queue = WriteQueue::new("users");
695 let mut log = CountingLog::new();
696 let mut seq = 0u64;
697
698 queue
699 .stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
700 .unwrap();
701 assert_eq!(queue.pending("u1"), HashSet::from(["id".into(), "name".into()]));
702
703 queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
704 assert!(queue.pending("u1").is_empty());
705 }
706
707 #[test]
708 fn rejects_a_second_record_claiming_a_unique_field_already_held() {
709 let schema = users_schema();
710 let mut registry = TableRegistry::new(schema.clone());
711 let mut queue = WriteQueue::new("users");
712 let mut log = CountingLog::new();
713 let mut seq = 0u64;
714
715 queue
716 .stage_put(
717 &schema,
718 ®istry,
719 &json!({ "id": "u1", "name": "Ania", "email": "a@x.io" }),
720 false,
721 )
722 .unwrap();
723 queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
724
725 let err = queue
726 .stage_put(
727 &schema,
728 ®istry,
729 &json!({ "id": "u2", "name": "Zosia", "email": "a@x.io" }),
730 false,
731 )
732 .unwrap_err();
733 assert!(matches!(err, StageError::Unique(_)));
734 }
735
736 fn docs_schema() -> TableSchema {
737 parse_schema(&json!({
738 "schemaVersion": 1,
739 "table": "docs",
740 "primaryKey": "id",
741 "fields": {
742 "id": { "type": "string" },
743 "page": { "type": "markdown" },
744 },
745 }))
746 .unwrap()
747 }
748
749 #[test]
750 fn rejects_a_second_record_claiming_an_already_taken_markdown_virtual_path() {
751 let schema = docs_schema();
752 let mut registry = TableRegistry::new(schema.clone());
753 let mut queue = WriteQueue::new("docs");
754 let mut log = CountingLog::new();
755 let mut seq = 0u64;
756
757 queue
758 .stage_put(
759 &schema,
760 ®istry,
761 &json!({ "id": "d1", "page": { "content": "a", "virtualPath": "/docs/a.md" } }),
762 false,
763 )
764 .unwrap();
765 queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
766
767 let err = queue
768 .stage_put(
769 &schema,
770 ®istry,
771 &json!({ "id": "d2", "page": { "content": "b", "virtualPath": "/docs/a.md" } }),
772 false,
773 )
774 .unwrap_err();
775 assert!(matches!(err, StageError::Unique(_)));
776 }
777
778 #[test]
779 fn dead_bytes_grow_when_an_existing_key_is_overwritten_or_removed() {
780 let schema = users_schema();
781 let mut registry = TableRegistry::new(schema.clone());
782 let mut queue = WriteQueue::new("users");
783 let mut log = CountingLog::new();
784 let mut seq = 0u64;
785
786 queue
787 .stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Ania" }), false)
788 .unwrap();
789 queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
790 assert_eq!(registry.stats().dead_bytes, 0);
791
792 queue
793 .stage_put(&schema, ®istry, &json!({ "id": "u1", "name": "Zosia" }), false)
794 .unwrap();
795 queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
796 let after_overwrite = registry.stats().dead_bytes;
797 assert!(after_overwrite > 0);
798
799 queue.stage_delete("u1", false, "").unwrap();
800 queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
801 assert!(registry.stats().dead_bytes > after_overwrite);
802 assert_eq!(registry.records(), 0);
803 }
804
805 #[test]
806 fn splits_a_batch_larger_than_max_batch_into_several_appends() {
807 let schema = users_schema();
808 let mut registry = TableRegistry::new(schema.clone());
809 let mut queue = WriteQueue::with_max_batch("users", 2);
810 let mut log = CountingLog::new();
811 let mut seq = 0u64;
812
813 for i in 0..5 {
814 queue
815 .stage_put(
816 &schema,
817 ®istry,
818 &json!({ "id": format!("u{i}"), "name": "x" }),
819 false,
820 )
821 .unwrap();
822 }
823 let results = queue.commit(&mut log, &mut registry, &mut seq, false).unwrap();
824 assert_eq!(results.len(), 5);
825 assert_eq!(log.calls, 3);
826 assert_eq!(log.synced, 0);
827 }
828}