1use haematite::branch::history::Key as HistoryKey;
2use haematite::branch::{
3 CommitDurability, CommitRequest, ForkWithChanges, ForkWrite, RefRecord, current_timestamp,
4 open as open_branch,
5};
6use haematite::db::branches::ForkKeptError;
7use haematite::tree::Hash;
8
9use super::{
10 Captured, Collection, DocumentRevision, Failure, Receipt, Storage, branch_name, collection_key,
11 counts, document_key, identifier, reads, receipt_key, revision_key, version,
12};
13use crate::serve::contracts::attribution::Row;
14use crate::serve::contracts::{Caller, local_owner};
15use crate::serve::protocol::{Change, DocumentImport, Mutation, ReceiptQuery, Selection, Version};
16use crate::serve::schema::Collection as CollectionSchema;
17
18pub(super) type PreparedWrite = (Vec<u8>, Option<Vec<u8>>);
19
20impl Storage {
21 pub fn receipt(&mut self, query: &ReceiptQuery) -> Result<Option<Receipt>, Failure> {
27 identifier(&query.request_id)?;
28 identifier(&query.branch)?;
29 if let Some(receipt) = self.discard_receipt(query)? {
30 return Ok(Some(receipt));
31 }
32 if self.pending_discards.contains(&branch_name(&query.branch)) {
33 return Err(Failure::new(
34 "discard_pending",
35 "branch removal requires receipt reconciliation",
36 ));
37 }
38 let record = self
41 .refs
42 .get(&branch_name(&query.branch))
43 .ok_or_else(|| Failure::new("branch_missing", &query.branch))?;
44 let at = Captured {
45 branch: query.branch.clone(),
46 version: version(record),
47 shards: record.shards.clone(),
48 };
49 let Some(bytes) = reads::get_bytes(
50 &self.reader,
51 &at,
52 &receipt_key(&query.branch, at.version.generation, &query.request_id),
53 )?
54 else {
55 return Ok(None);
56 };
57 let receipt: Receipt = serde_json::from_slice(&bytes).map_err(Failure::storage)?;
58 if receipt.request_id != query.request_id
59 || receipt.branch != query.branch
60 || receipt.version.generation != at.version.generation
61 {
62 return Err(Failure::new(
63 "receipt_corrupt",
64 "receipt identity does not match its stored key",
65 ));
66 }
67 self.refs
68 .confirm_version(
69 &branch_name(&query.branch),
70 at.version.generation,
71 at.version.sequence,
72 )
73 .map_err(|error| Failure::new("durability_unknown", error.to_string()))?;
74 self.publish_record(&branch_name(&query.branch))?;
75 Ok(Some(receipt))
76 }
77
78 pub fn mutate(&mut self, request: &Mutation) -> Result<Receipt, Failure> {
84 let author = local_owner()?;
85 self.mutate_with_import_bound(request, &author, None, super::AppendPruneLimits::default())
86 }
87
88 pub fn mutate_bounded(
93 &mut self,
94 request: &Mutation,
95 max_page: std::num::NonZeroU32,
96 ) -> Result<Receipt, Failure> {
97 let author = local_owner()?;
98 self.mutate_with_import_bound(
99 request,
100 &author,
101 Some(max_page),
102 super::AppendPruneLimits::default(),
103 )
104 }
105
106 pub fn mutate_with_append_prune_limits(
111 &mut self,
112 request: &Mutation,
113 max_page: std::num::NonZeroU32,
114 limits: super::AppendPruneLimits,
115 ) -> Result<Receipt, Failure> {
116 let author = local_owner()?;
117 self.mutate_with_import_bound(request, &author, Some(max_page), limits)
118 }
119
120 pub fn mutate_by(
129 &mut self,
130 request: &Mutation,
131 author: &Caller,
132 max_page: std::num::NonZeroU32,
133 limits: super::AppendPruneLimits,
134 ) -> Result<Receipt, Failure> {
135 self.mutate_with_import_bound(request, author, Some(max_page), limits)
136 }
137
138 fn expected_record(&self, request: &Mutation) -> Result<RefRecord, Failure> {
140 let record = self
141 .refs
142 .get(&branch_name(&request.branch))
143 .ok_or_else(|| Failure::new("branch_missing", &request.branch))?;
144 if version(record) != request.expected {
145 return Err(Failure::new(
146 "version_conflict",
147 "expected branch generation or sequence differs",
148 ));
149 }
150 if self.pending_discards.contains(&record.name) {
151 return Err(Failure::new(
152 "discard_pending",
153 "branch removal requires receipt reconciliation",
154 ));
155 }
156 Ok(record.clone())
157 }
158
159 fn mutate_with_import_bound(
160 &mut self,
161 request: &Mutation,
162 author: &Caller,
163 import_max_page: Option<std::num::NonZeroU32>,
164 retention_limits: super::AppendPruneLimits,
165 ) -> Result<Receipt, Failure> {
166 identifier(&request.request_id)?;
167 identifier(&request.branch)?;
168 let destination = match &request.change {
169 Change::Fork(fork) => &fork.name,
170 _ => &request.branch,
171 };
172 identifier(destination)?;
173 if let Some(stage) = &request.stage {
174 super::lineage::validate(stage, &request.change)?;
175 }
176 let fingerprint =
177 Hash::of(&serde_json::to_vec(request).map_err(Failure::storage)?).to_string();
178 if let Some(receipt) = self.replay(request, destination, &fingerprint)? {
179 return self.replayed(request, receipt);
180 }
181 validate_import_bound(&request.change, import_max_page)?;
182 let record = self.expected_record(request)?;
183 let at = Captured {
184 branch: request.branch.clone(),
185 version: request.expected,
186 shards: record.shards.clone(),
187 };
188 if let Change::Fork(fork) = &request.change {
189 return self.fork_request(request, author, &fingerprint, &fork.name, fork.from);
190 }
191 if let Change::Discard(discard) = &request.change {
192 return self.discard_request(request, &fingerprint, &discard.name, &record);
193 }
194 let next = Version {
195 generation: record.created,
196 sequence: record.seq.checked_add(1).ok_or_else(|| {
197 Failure::new("sequence_exhausted", "branch sequence cannot advance")
198 })?,
199 };
200 let (mut writes, retention) = match &request.change {
201 Change::AppendPrune(write) => {
202 let (writes, result) =
203 self.prepare_append_prune(&at, write, next, retention_limits)?;
204 (writes, Some(result))
205 }
206 change => (self.prepare_change(&at, change, next)?, None),
207 };
208 counts::append(&self.reader, &at, &request.change, &mut writes)?;
209 let authored = author_write(request, author, next)?;
212 self.rows_before_advance(&request.branch, &record)?;
215 #[cfg(test)]
216 super::history::after_rows::reached()?;
217 let receipt = Receipt {
218 request_id: request.request_id.clone(),
219 fingerprint,
220 branch: request.branch.clone(),
221 version: next,
222 retention,
223 versions_released: None,
224 };
225 push_own_rows(&mut writes, request, &receipt, authored)?;
226 let branch = open_branch(&record.name, &self.refs, &self.registry)
227 .map_err(Failure::branch_commit)?;
228 for (key, value) in writes {
229 let shard_id = self.writer.shard_for(&key);
230 match value {
231 Some(value) => branch.put(shard_id, key, value),
232 None => branch.delete(shard_id, key),
233 }
234 .map_err(Failure::branch)?;
235 }
236 #[cfg(test)]
237 super::lineage::before_commit::reached()?;
238 let result = self.writer.commit(
239 &branch,
240 &self.registry,
241 CommitRequest {
242 durability: CommitDurability::Durable {
243 refs: &mut self.refs,
244 },
245 extra_parents: &[],
246 timestamp: current_timestamp(),
247 note: None,
248 },
249 );
250 if let Err(error) = result {
251 let installed = self
252 .refs
253 .get(&record.name)
254 .is_some_and(|now| version(now) != request.expected);
255 let mut failure = Failure::branch_commit(error);
256 let code = if installed {
257 "durability_unknown"
258 } else {
259 "write_refused"
260 };
261 code.clone_into(&mut failure.code);
262 return Err(failure);
263 }
264 #[cfg(test)]
265 lose_mutation_answer()?;
266 #[cfg(test)]
267 publish_pause_point();
268 self.publish_record(&record.name)?;
269 Ok(receipt)
270 }
271
272 fn replayed(&mut self, request: &Mutation, receipt: Receipt) -> Result<Receipt, Failure> {
277 if matches!(&request.change, Change::Discard(_)) {
278 self.release_after_discard(&receipt.branch, receipt.version.generation)?;
279 }
280 Ok(receipt)
281 }
282
283 fn replay(
284 &mut self,
285 request: &Mutation,
286 destination: &str,
287 fingerprint: &str,
288 ) -> Result<Option<Receipt>, Failure> {
289 if matches!(&request.change, Change::Discard(_))
291 || self.refs.get(&branch_name(destination)).is_some()
292 {
293 if let Some(receipt) = self.receipt(&ReceiptQuery {
294 branch: destination.to_owned(),
295 request_id: request.request_id.clone(),
296 })? {
297 return if receipt.fingerprint == fingerprint {
298 Ok(Some(receipt))
299 } else {
300 Err(Failure::new(
301 "request_conflict",
302 "request identity was used with different content",
303 ))
304 };
305 }
306 if matches!(&request.change, Change::Fork(_)) {
307 return Err(Failure::new(
308 "branch_exists",
309 "destination exists without this fork receipt",
310 ));
311 }
312 }
313 Ok(None)
314 }
315
316 pub(super) fn prepare_change(
317 &self,
318 at: &Captured,
319 change: &Change,
320 next: Version,
321 ) -> Result<Vec<PreparedWrite>, Failure> {
322 match change {
323 Change::Collection(write) => {
324 identifier(&write.collection)?;
325 if let Some(schema) = &write.schema {
326 CollectionSchema::compile(schema).map_err(schema_failure)?;
327 }
328 let collection = Collection {
329 name: write.collection.clone(),
330 schema: write.schema.clone(),
331 changed_in: at.branch.clone(),
332 changed_at: next,
333 };
334 Ok(vec![(
335 collection_key(&write.collection),
336 Some(serde_json::to_vec(&collection).map_err(Failure::storage)?),
337 )])
338 }
339 Change::Put(write) => {
340 identifier(&write.document_id)?;
341 let collection = reads::collection(&self.reader, at, &write.collection)?
342 .ok_or_else(|| Failure::new("collection_missing", &write.collection))?;
343 if let Some(schema) = &collection.schema {
344 CollectionSchema::compile(schema)
345 .map_err(schema_failure)?
346 .validate(
347 &serde_json::from_str(write.document.get())
348 .map_err(Failure::storage)?,
349 )
350 .map_err(schema_failure)?;
351 }
352 let revision = DocumentRevision {
353 collection_branch: collection.changed_in,
354 collection_version: collection.changed_at,
355 written_in: at.branch.clone(),
356 written_at: next,
357 };
358 Ok(vec![
359 (
360 document_key(&write.collection, &write.document_id),
361 Some(write.document.get().as_bytes().to_vec()),
362 ),
363 (
364 revision_key(&write.collection, &write.document_id),
365 Some(serde_json::to_vec(&revision).map_err(Failure::storage)?),
366 ),
367 ])
368 }
369 Change::AppendPrune(_) => Err(Failure::new(
370 "append_prune_preparation_required",
371 "retention uses bounded mutation preparation",
372 )),
373 Change::Delete(write) => {
374 identifier(&write.collection)?;
375 identifier(&write.document_id)?;
376 let mut writes = vec![
379 (document_key(&write.collection, &write.document_id), None),
380 (revision_key(&write.collection, &write.document_id), None),
381 ];
382 writes.extend(super::edges::removals_for_document(
383 &self.reader,
384 at,
385 &write.collection,
386 &write.document_id,
387 )?);
388 writes.extend(super::vectors::removals_for_document(
389 &self.reader,
390 at,
391 &write.collection,
392 &write.document_id,
393 )?);
394 Ok(writes)
395 }
396 Change::DeleteMany(delete) => self.prepare_delete_many(at, delete, next),
397 Change::Merge(merge) => self.merge_writes(at, merge, next),
398 Change::Fork(_) | Change::Discard(_) => Err(Failure::new(
399 "invalid_operation",
400 "branch control uses the publication path",
401 )),
402 Change::EdgePut(edge) => super::edges::put(&self.reader, at, edge),
403 Change::EdgeDelete(edge) => super::edges::delete(edge),
404 Change::EdgePutMany(import) => super::edges::put_many(&self.reader, at, import),
405 Change::IndexDeclare(field) => {
406 super::field_index::declare(&self.reader, at, field, next)
407 }
408 Change::VectorDeclare(field) => super::vectors::declare(&self.reader, at, field),
409 Change::VectorPut(vector) => super::vectors::put(&self.reader, at, vector),
410 Change::VectorPutMany(import) => {
411 super::vector_import::prepare(&self.reader, at, import)
412 }
413 Change::VectorDelete(vector) => super::vectors::delete(vector),
414 Change::PutMany(import) => self.prepare_import(at, import, next),
415 Change::RelationPut(relation) => super::relations::put(&self.reader, at, relation),
416 Change::RelationDelete(relation) => super::relations::delete(relation),
417 }
418 }
419
420 fn prepare_import(
421 &self,
422 at: &Captured,
423 import: &DocumentImport,
424 next: Version,
425 ) -> Result<Vec<PreparedWrite>, Failure> {
426 if import.documents.is_empty() {
427 return Err(Failure {
428 code: "import_empty".to_owned(),
429 message: "an import must contain at least one document".to_owned(),
430 path: Some("/documents".to_owned()),
431 cause: None,
432 });
433 }
434 let collection = reads::collection(&self.reader, at, &import.collection)
435 .map_err(|mut failure| {
436 failure.path = Some("/collection".to_owned());
437 failure
438 })?
439 .ok_or_else(|| Failure {
440 code: "collection_missing".to_owned(),
441 message: import.collection.clone(),
442 path: Some("/collection".to_owned()),
443 cause: None,
444 })?;
445 let schema = collection
446 .schema
447 .as_ref()
448 .map(CollectionSchema::compile)
449 .transpose()
450 .map_err(schema_failure)?;
451 validate_import_rows(import, schema.as_ref())?;
452 let revision = serde_json::to_vec(&DocumentRevision {
453 collection_branch: collection.changed_in,
454 collection_version: collection.changed_at,
455 written_in: at.branch.clone(),
456 written_at: next,
457 })
458 .map_err(Failure::storage)?;
459 let mut writes = Vec::new();
460 for document in &import.documents {
461 writes.push((
462 document_key(&import.collection, &document.document_id),
463 Some(document.document.get().as_bytes().to_vec()),
464 ));
465 writes.push((
466 revision_key(&import.collection, &document.document_id),
467 Some(revision.clone()),
468 ));
469 }
470 Ok(writes)
471 }
472
473 fn fork_request(
474 &mut self,
475 request: &Mutation,
476 author: &Caller,
477 fingerprint: &str,
478 destination: &str,
479 from: Version,
480 ) -> Result<Receipt, Failure> {
481 if self.pending_discards.contains(&branch_name(destination)) {
482 return Err(Failure::new(
483 "discard_pending",
484 "destination removal requires receipt reconciliation",
485 ));
486 }
487 self.capture(&Selection {
488 branch: request.branch.clone(),
489 version: Some(from),
490 })?;
491 let current = self
492 .refs
493 .get(&branch_name(&request.branch))
494 .ok_or_else(|| Failure::new("branch_missing", &request.branch))?
495 .clone();
496 if from == version(¤t) {
499 self.retain_current(&request.branch, ¤t)?;
500 }
501 self.row_for_kept(&request.branch, ¤t, from)?;
504 #[cfg(test)]
505 super::history::after_rows::reached()?;
506 let service =
507 open_branch(super::SERVICE, &self.refs, &self.registry).map_err(|source| {
508 fork_kept_failure(&ForkKeptError::ServiceBindingStale { source }, false)
509 })?;
510 let line = matches!(&request.change, Change::Fork(fork) if fork.line);
511 let history_key = HistoryKey {
512 branch: Hash::of(request.branch.as_bytes()),
513 generation: from.generation,
514 sequence: from.sequence,
515 };
516 let created = current_timestamp();
517 let receipt = Receipt {
518 request_id: request.request_id.clone(),
519 fingerprint: fingerprint.to_owned(),
520 branch: destination.to_owned(),
521 version: Version {
522 generation: created,
523 sequence: 0,
524 },
525 retention: None,
526 versions_released: None,
527 };
528 let key = receipt_key(destination, created, &request.request_id);
529 let (author_key, row) = author_write(request, author, receipt.version)?;
530 let writes = [
531 ForkWrite {
532 shard_id: self.writer.shard_for(&key),
533 key,
534 value: Some(serde_json::to_vec(&receipt).map_err(Failure::storage)?),
535 },
536 super::branches::provenance_write(&self.writer, &request.branch, from, line)?,
537 ForkWrite {
538 shard_id: self.writer.shard_for(&author_key),
539 key: author_key,
540 value: Some(row),
541 },
542 ];
543 let name = branch_name(destination);
544 self.writer
545 .fork_kept(
546 &service,
547 history_key,
548 &mut self.refs,
549 &self.registry,
550 ForkWithChanges {
551 name: &name,
552 timestamp: created,
553 writes: &writes,
554 note: None,
555 },
556 )
557 .map_err(|error| fork_kept_failure(&error, self.refs.get(&name).is_some()))?;
558 #[cfg(test)]
559 lose_mutation_answer()?;
560 self.publish_record(&name)?;
561 Ok(receipt)
562 }
563
564 #[cfg(test)]
569 pub fn arm_publish_pause(
570 paused: std::sync::mpsc::Sender<()>,
571 release: std::sync::mpsc::Receiver<()>,
572 ) {
573 PUBLISH_PAUSE.with(|slot| {
574 *slot.borrow_mut() = Some(PublishPause { paused, release });
575 });
576 }
577
578 #[cfg(test)]
579 pub(crate) fn lose_next_mutation_answer() {
580 Self::lose_next_mutation_answer_with_message(
581 "injected lost answer after reference commit".to_owned(),
582 );
583 }
584
585 #[cfg(test)]
586 pub(crate) fn lose_next_mutation_answer_with_message(message: String) {
587 LOST_MUTATION_ANSWER.with(|slot| {
588 *slot.borrow_mut() = Some(message);
589 });
590 }
591}
592
593#[cfg(test)]
594thread_local! {
595 static LOST_MUTATION_ANSWER: std::cell::RefCell<Option<String>> = const { std::cell::RefCell::new(None) };
596}
597
598#[cfg(test)]
599fn lose_mutation_answer() -> Result<(), Failure> {
600 if let Some(message) = LOST_MUTATION_ANSWER.with(|slot| slot.borrow_mut().take()) {
601 return Err(Failure::new("durability_unknown", message));
602 }
603 Ok(())
604}
605
606pub(in crate::serve) fn author_write(
614 request: &Mutation,
615 author: &Caller,
616 written_at: Version,
617) -> Result<(Vec<u8>, Vec<u8>), Failure> {
618 let row = Row::by(&request.request_id, author, request.expected, written_at)?;
619 Ok((row.key()?, row.encode()?))
620}
621
622fn push_own_rows(
625 writes: &mut Vec<PreparedWrite>,
626 request: &Mutation,
627 receipt: &Receipt,
628 (author_key, author_row): (Vec<u8>, Vec<u8>),
629) -> Result<(), Failure> {
630 let key = receipt_key(
631 &request.branch,
632 receipt.version.generation,
633 &request.request_id,
634 );
635 writes.push((
636 key,
637 Some(serde_json::to_vec(receipt).map_err(Failure::storage)?),
638 ));
639 writes.push((author_key, Some(author_row)));
640 if let Some(stage) = &request.stage {
641 writes.push(super::lineage::stage_write(receipt.version, stage)?);
642 }
643 Ok(())
644}
645
646fn fork_kept_failure(error: &ForkKeptError, installed: bool) -> Failure {
647 let code = match error {
648 ForkKeptError::ServiceBindingStale { .. } => "kept_service_binding_stale",
649 ForkKeptError::ServiceInvalid { .. } => "kept_service_invalid",
650 ForkKeptError::ShardSetIncomplete { .. } => "kept_shard_set_incomplete",
651 ForkKeptError::RowMissing { .. } => "kept_row_missing",
652 ForkKeptError::RowInvalid { .. } => "kept_row_invalid",
653 ForkKeptError::RowRead { .. } => "kept_row_read",
654 ForkKeptError::RootUnavailable { .. } => "kept_root_unavailable",
655 ForkKeptError::Fork(_) if installed => "durability_unknown",
656 ForkKeptError::Fork(_) => "write_refused",
657 };
658 Failure::new(code, error.to_string())
659}
660
661#[cfg(test)]
662#[path = "mutations_kept_tests.rs"]
663mod kept_tests;
664
665fn validate_import_rows(
666 import: &DocumentImport,
667 schema: Option<&CollectionSchema>,
668) -> Result<(), Failure> {
669 let mut ids = std::collections::BTreeSet::new();
670 for (index, document) in import.documents.iter().enumerate() {
671 identifier(&document.document_id).map_err(|mut failure| {
672 failure.path = Some(format!("/documents/{index}/document_id"));
673 failure
674 })?;
675 if !ids.insert(document.document_id.as_str()) {
676 return Err(Failure {
677 code: "import_duplicate".to_owned(),
678 message: format!("repeated document_id {}", document.document_id),
679 path: Some(format!("/documents/{index}/document_id")),
680 cause: None,
681 });
682 }
683 if let Some(schema) = schema {
684 let value = serde_json::from_str(document.document.get()).map_err(|error| Failure {
685 code: "invalid_document".to_owned(),
686 message: error.to_string(),
687 path: Some(format!("/documents/{index}/document")),
688 cause: None,
689 })?;
690 schema.validate(&value).map_err(|error| Failure {
691 code: error.code().to_owned(),
692 message: error.reason,
693 path: Some(format!("/documents/{index}/document{}", error.path)),
694 cause: None,
695 })?;
696 }
697 }
698 Ok(())
699}
700
701fn validate_import_bound(
702 change: &Change,
703 max_page: Option<std::num::NonZeroU32>,
704) -> Result<(), Failure> {
705 let (rows, kind) = match change {
706 Change::PutMany(import) => (import.documents.len(), "documents"),
707 Change::DeleteMany(delete) => (delete.document_ids.len(), "document_ids"),
708 Change::EdgePutMany(import) => (import.edges.len(), "edges"),
709 Change::VectorPutMany(import) => (import.vectors.len(), "vectors"),
710 _ => return Ok(()),
711 };
712 let max_page = max_page.ok_or_else(|| {
713 Failure::new(
714 "import_limit_required",
715 "an import requires the service page bound",
716 )
717 })?;
718 let bound = u64::from(max_page.get())
719 .checked_mul(64)
720 .ok_or_else(|| Failure::new("import_limit", "import row bound exceeds u64"))?;
721 if u64::try_from(rows).map_err(Failure::storage)? > bound {
722 return Err(Failure {
723 code: "import_limit".to_owned(),
724 message: format!("maximum import size is {bound} {kind}"),
725 path: Some(format!("/{kind}/{bound}")),
726 cause: None,
727 });
728 }
729 Ok(())
730}
731
732fn schema_failure(error: crate::serve::schema::Violation) -> Failure {
733 Failure {
734 code: error.code().to_owned(),
735 message: error.reason,
736 path: Some(error.path),
737 cause: None,
738 }
739}
740
741#[cfg(test)]
742struct PublishPause {
743 paused: std::sync::mpsc::Sender<()>,
744 release: std::sync::mpsc::Receiver<()>,
745}
746
747#[cfg(test)]
748thread_local! {
749 static PUBLISH_PAUSE: std::cell::RefCell<Option<PublishPause>> =
750 const { std::cell::RefCell::new(None) };
751}
752
753#[cfg(test)]
756fn publish_pause_point() {
757 let Some(pause) = PUBLISH_PAUSE.with(|slot| slot.borrow_mut().take()) else {
758 return;
759 };
760 assert!(
761 pause.paused.send(()).is_ok(),
762 "paused signal receiver was dropped"
763 );
764 let wait = std::time::Duration::from_secs(10);
765 let released = pause.release.recv_timeout(wait);
766 assert!(released.is_ok(), "release signal not received within 10 s");
767}