1use std::num::NonZeroUsize;
4
5use haematite::TreeError;
6use haematite::db::{ObserverError, ReadOnlyDatabase};
7use serde::Serialize;
8use serde_json::value::RawValue;
9
10use super::{
11 Captured, Collection, DocumentRevision, Failure, collection_key, document_key, identifier,
12 revision_key,
13};
14use crate::serve::access::Access;
15use crate::serve::protocol::{CountReason, DocumentPage, KeptCount, Version};
16
17#[derive(Debug, Serialize)]
18pub struct Document {
19 pub id: String,
20 pub raw_json: String,
21 pub revision: DocumentRevision,
22}
23
24#[derive(Debug, Serialize)]
25#[serde(untagged)]
26pub(in crate::serve) enum DocumentListEntry {
27 Full(Document),
28 Metadata {
29 id: String,
30 revision: DocumentRevision,
31 },
32}
33
34#[derive(Clone, Copy)]
35pub(in crate::serve) enum DocumentContent {
36 Full,
37 None,
45}
46
47#[derive(Debug, Serialize)]
48pub struct CollectionEntry {
49 pub name: String,
50 pub changed_in: String,
51 pub changed_at: Version,
52 pub schema_json: Option<String>,
53 pub documents: Option<u64>,
54 #[serde(skip_serializing_if = "Option::is_none")]
55 pub documents_reason: Option<CountReason>,
56}
57
58fn render_collection(row: Collection, total: &KeptCount) -> Result<CollectionEntry, Failure> {
59 let schema_json = row
60 .schema
61 .as_ref()
62 .map(serde_json::to_string)
63 .transpose()
64 .map_err(Failure::storage)?;
65 Ok(CollectionEntry {
66 name: row.name,
67 changed_in: row.changed_in,
68 changed_at: row.changed_at,
69 schema_json,
70 documents: total.value,
71 documents_reason: total.reason,
72 })
73}
74
75pub(in crate::serve) fn render_collections(
76 access: &Access,
77 reader: &ReadOnlyDatabase,
78 at: &Captured,
79 collections: Vec<Collection>,
80) -> Result<Vec<CollectionEntry>, Failure> {
81 let totals = collection_totals(access, reader, at, &collections)?;
82 collections
83 .into_iter()
84 .zip(totals)
85 .map(|(row, total)| render_collection(row, &total))
86 .collect()
87}
88
89#[derive(Debug, Serialize)]
90pub struct Page<T> {
91 pub branch: String,
92 pub version: Version,
93 pub entries: Vec<T>,
94 pub next_after: Option<String>,
95}
96
97#[derive(Serialize)]
98pub(in crate::serve) struct CollectionCount {
99 name: String,
100 documents: Option<u64>,
101 #[serde(skip_serializing_if = "Option::is_none")]
102 documents_reason: Option<CountReason>,
103}
104
105#[derive(Serialize)]
106pub(in crate::serve) struct EmbeddedFieldCount {
107 collection: String,
108 field: String,
109 kind: super::embedded::VectorKind,
110 storage: &'static str,
111 count: Option<u64>,
112 count_reason: &'static str,
113}
114
115pub(in crate::serve) struct Stats {
116 pub(in crate::serve) collections: Vec<CollectionCount>,
117 pub(in crate::serve) collection_count: KeptCount,
118 pub(in crate::serve) relation_count: KeptCount,
119 pub(in crate::serve) documents: KeptCount,
120 pub(in crate::serve) edges: KeptCount,
121 pub(in crate::serve) vectors: KeptCount,
122 pub(in crate::serve) embedded_fields: Option<Vec<EmbeddedFieldCount>>,
123 pub(in crate::serve) embedded_fields_reason: Option<CountReason>,
124}
125
126pub(in crate::serve) fn collection_totals(
127 access: &Access,
128 reader: &ReadOnlyDatabase,
129 at: &Captured,
130 collections: &[Collection],
131) -> Result<Vec<KeptCount>, Failure> {
132 let view = super::counts::View::new(reader, at)?;
133 collections
134 .iter()
135 .map(|row| view.documents(access, &row.name))
136 .collect()
137}
138
139pub(in crate::serve) fn stats(
140 access: &Access,
141 reader: &ReadOnlyDatabase,
142 at: &Captured,
143) -> Result<Stats, Failure> {
144 let view = super::counts::View::new(reader, at)?;
145 let mut collections = std::collections::BTreeSet::new();
146 scan_rows(reader, at, b"c\0", 1, |parts, _| {
147 if !collections.insert(parts[0].to_owned()) {
148 return Err(Failure::new(
149 "corrupt_key",
150 "collection key exists in multiple shards",
151 ));
152 }
153 Ok(())
154 })?;
155 let allowed = access.can_read_store();
156 let collection_count = declaration_total(allowed, collections.len())?;
157 let mut declared_relations = std::collections::BTreeSet::new();
158 if allowed {
159 scan_rows(reader, at, b"n\0", 1, |parts, bytes| {
160 super::relations::decode(parts[0], bytes)?;
161 if !declared_relations.insert(parts[0].to_owned()) {
162 return Err(Failure::new(
163 "counts_corrupt",
164 "duplicate relation declaration",
165 ));
166 }
167 Ok(())
168 })?;
169 }
170 let relation_count = declaration_total(allowed, declared_relations.len())?;
171 let mut documents = 0_u64;
172 let mut entries = Vec::new();
173 for name in collections {
174 let total = view.documents(access, &name)?;
175 if let Some(count) = total.value {
176 documents = documents
177 .checked_add(count)
178 .ok_or_else(|| Failure::new("count_overflow", "document total exceeds u64"))?;
179 }
180 entries.push(CollectionCount {
181 name,
182 documents: total.value,
183 documents_reason: total.reason,
184 });
185 }
186 if let Some(reason) = view.unavailable(access.can_read_store()) {
187 return Ok(Stats {
188 collections: entries,
189 collection_count,
190 relation_count,
191 documents: KeptCount {
192 value: None,
193 reason: Some(reason),
194 },
195 edges: KeptCount {
196 value: None,
197 reason: Some(reason),
198 },
199 vectors: KeptCount {
200 value: None,
201 reason: Some(reason),
202 },
203 embedded_fields: if allowed {
204 Some(embedded_field_counts(reader, at, None)?)
205 } else {
206 None
207 },
208 embedded_fields_reason: (!allowed).then_some(CountReason::Withheld),
209 });
210 }
211 let (edges, relations) = super::counts::sum_rows(reader, at, b"k\0e\0", 1)?;
212 let (vectors, fields) = super::counts::sum_rows(reader, at, b"k\0v\0", 2)?;
213 for name in declared_relations {
214 if !relations.contains(&name) {
215 return Err(Failure::new(
216 "counts_corrupt",
217 "declared relation has no count",
218 ));
219 }
220 }
221 let embedded_fields = embedded_field_counts(reader, at, Some(&fields))?;
222 Ok(Stats {
223 collections: entries,
224 collection_count,
225 relation_count,
226 documents: KeptCount {
227 value: Some(documents),
228 reason: None,
229 },
230 edges: KeptCount {
231 value: Some(edges),
232 reason: None,
233 },
234 vectors: KeptCount {
235 value: Some(vectors),
236 reason: None,
237 },
238 embedded_fields: Some(embedded_fields),
239 embedded_fields_reason: None,
240 })
241}
242
243fn embedded_field_counts(
244 reader: &ReadOnlyDatabase,
245 at: &Captured,
246 required_supplied: Option<&std::collections::BTreeSet<String>>,
247) -> Result<Vec<EmbeddedFieldCount>, Failure> {
248 let mut embedded_fields = Vec::new();
249 scan_rows(reader, at, b"w\0", 2, |parts, bytes| {
250 match super::embedded::read_field(bytes)? {
251 super::embedded::Field::Supplied { .. } => {
252 if required_supplied.is_some_and(|fields| !fields.contains(&parts.join("\0"))) {
253 return Err(Failure::new(
254 "counts_corrupt",
255 "declared field has no count",
256 ));
257 }
258 }
259 super::embedded::Field::Embedded(declaration) => {
260 embedded_fields.push(EmbeddedFieldCount {
262 collection: parts[0].to_owned(),
263 field: parts[1].to_owned(),
264 kind: declaration.kind,
265 storage: "companion",
266 count: None,
267 count_reason: "companion_count_not_kept",
268 });
269 }
270 }
271 Ok(())
272 })?;
273 Ok(embedded_fields)
274}
275
276fn declaration_total(allowed: bool, count: usize) -> Result<KeptCount, Failure> {
277 if !allowed {
278 return Ok(KeptCount {
279 value: None,
280 reason: Some(CountReason::Withheld),
281 });
282 }
283 let value =
284 u64::try_from(count).map_err(|error| Failure::new("count_overflow", error.to_string()))?;
285 Ok(KeptCount {
286 value: Some(value),
287 reason: None,
288 })
289}
290
291#[cfg(test)]
292pub(super) mod test_support {
293 use super::{Captured, Failure, ReadOnlyDatabase, scan_rows};
294
295 pub(in crate::serve::storage) fn count_keys(
296 reader: &ReadOnlyDatabase,
297 at: &Captured,
298 prefix: &[u8],
299 parts: usize,
300 ) -> Result<u64, Failure> {
301 let mut count = 0_u64;
302 scan_rows(reader, at, prefix, parts, |_, _| {
303 count = increment(count)?;
304 Ok(())
305 })?;
306 Ok(count)
307 }
308
309 fn increment(value: u64) -> Result<u64, Failure> {
310 value
311 .checked_add(1)
312 .ok_or_else(|| Failure::new("count_overflow", "stored key count exceeds u64"))
313 }
314}
315
316pub(super) fn scan_rows(
318 reader: &ReadOnlyDatabase,
319 at: &Captured,
320 prefix: &[u8],
321 expected_parts: usize,
322 mut visit: impl FnMut(&[&str], &[u8]) -> Result<(), Failure>,
323) -> Result<(), Failure> {
324 let rows = NonZeroUsize::new(64)
325 .ok_or_else(|| Failure::new("page_limit", "key scan page size must be nonzero"))?;
326 let mut from = prefix.to_vec();
327 loop {
328 let (_, next_after) = bounded_key_page(
329 reader,
330 at,
331 prefix,
332 &from,
333 PageLimits {
334 rows,
335 bytes: NonZeroUsize::MAX,
336 },
337 |suffix, bytes| {
338 let parts: Vec<_> = suffix.split('\0').collect();
339 if parts.len() != expected_parts {
340 return Err(Failure::new(
341 "corrupt_key",
342 "stored key has the wrong number of parts",
343 ));
344 }
345 for part in &parts {
346 identifier(part)?;
347 }
348 visit(&parts, &bytes)
349 },
350 |&()| Ok(0),
351 )?;
352 let Some(next_after) = next_after else {
353 return Ok(());
354 };
355 from = prefix.to_vec();
356 from.extend_from_slice(next_after.as_bytes());
357 from.push(0);
358 }
359}
360
361pub fn get_bytes(
366 reader: &ReadOnlyDatabase,
367 at: &Captured,
368 key: &[u8],
369) -> Result<Option<Vec<u8>>, Failure> {
370 let shard_id = reader.shard_for(key);
371 let shard = at
372 .shards
373 .iter()
374 .find(|shard| shard.shard_id == shard_id)
375 .ok_or_else(|| {
376 Failure::new(
377 "incomplete_roots",
378 "captured version does not name the owning shard",
379 )
380 })?;
381 #[cfg(test)]
382 count_point_read();
383 #[cfg(test)]
384 count_content_read(key);
385 reader
386 .get_at(shard_id, shard.head, key)
387 .map_err(|error| point_read_failure(at, error))
388}
389
390fn point_read_failure(at: &Captured, error: ObserverError) -> Failure {
395 let gone = match error {
396 ObserverError::Tree(TreeError::MissingNode { hash }) => format!("node {hash} is gone"),
397 ObserverError::RootUnavailable { shard_id, root } => {
398 format!("shard {shard_id} cannot read root {root}")
399 }
400 other => return Failure::observer(other),
401 };
402 Failure::new(
403 "version_released",
404 format!(
405 "version {}.{} of branch {} was released by the store's tidying and cannot be \
406 read: {gone}",
407 at.version.generation, at.version.sequence, at.branch
408 ),
409 )
410}
411
412pub fn collection(
416 reader: &ReadOnlyDatabase,
417 at: &Captured,
418 name: &str,
419) -> Result<Option<Collection>, Failure> {
420 identifier(name)?;
421 get_bytes(reader, at, &collection_key(name))?
422 .map(|bytes| serde_json::from_slice(&bytes).map_err(Failure::storage))
423 .transpose()
424}
425
426pub fn document(
453 access: &crate::serve::access::Access,
454 reader: &ReadOnlyDatabase,
455 at: &Captured,
456 collection_name: &str,
457 id: &str,
458) -> Result<Option<Document>, Failure> {
459 identifier(collection_name)?;
460 identifier(id)?;
461 if !access.can_read_document(collection_name, id) {
462 return Ok(None);
463 }
464 get_bytes(reader, at, &document_key(collection_name, id))?
465 .map(|bytes| render_document(reader, at, collection_name, id.to_owned(), &bytes))
466 .transpose()
467}
468
469pub(super) fn document_exists_for_constraint(
470 reader: &ReadOnlyDatabase,
471 at: &Captured,
472 collection_name: &str,
473 id: &str,
474) -> Result<bool, Failure> {
475 identifier(collection_name)?;
476 identifier(id)?;
477 let Some(bytes) = get_bytes(reader, at, &document_key(collection_name, id))? else {
478 return Ok(false);
479 };
480 render_document(reader, at, collection_name, id.to_owned(), &bytes)?;
481 Ok(true)
482}
483
484pub fn documents(
489 access: &Access,
490 reader: &ReadOnlyDatabase,
491 at: &Captured,
492 collection_name: &str,
493 after: Option<&str>,
494 limit: NonZeroUsize,
495) -> Result<Page<Document>, Failure> {
496 document_page(
497 access,
498 reader,
499 at,
500 DocumentRange {
501 collection: collection_name,
502 prefix: None,
503 after,
504 metadata: false,
505 },
506 DocumentLimits {
507 rows: limit,
508 bytes: None,
509 },
510 render_document,
511 )
512}
513
514pub fn documents_bounded(
528 access: &Access,
529 reader: &ReadOnlyDatabase,
530 at: &Captured,
531 collection_name: &str,
532 after: Option<&str>,
533 limits: PageLimits,
534) -> Result<Page<Document>, Failure> {
535 document_page(
536 access,
537 reader,
538 at,
539 DocumentRange {
540 collection: collection_name,
541 prefix: None,
542 after,
543 metadata: false,
544 },
545 DocumentLimits {
546 rows: limits.rows,
547 bytes: Some(limits.bytes),
548 },
549 render_document,
550 )
551}
552
553pub fn documents_prefixed(
558 access: &Access,
559 reader: &ReadOnlyDatabase,
560 at: &Captured,
561 collection: &str,
562 prefix: Option<&str>,
563 after: Option<&str>,
564 limit: NonZeroUsize,
565) -> Result<Page<Document>, Failure> {
566 document_page(
567 access,
568 reader,
569 at,
570 DocumentRange {
571 collection,
572 prefix,
573 after,
574 metadata: false,
575 },
576 DocumentLimits {
577 rows: limit,
578 bytes: None,
579 },
580 render_document,
581 )
582}
583
584pub fn documents_prefixed_bounded(
589 access: &Access,
590 reader: &ReadOnlyDatabase,
591 at: &Captured,
592 collection: &str,
593 prefix: Option<&str>,
594 after: Option<&str>,
595 limits: PageLimits,
596) -> Result<Page<Document>, Failure> {
597 document_page(
598 access,
599 reader,
600 at,
601 DocumentRange {
602 collection,
603 prefix,
604 after,
605 metadata: false,
606 },
607 DocumentLimits {
608 rows: limits.rows,
609 bytes: Some(limits.bytes),
610 },
611 render_document,
612 )
613}
614
615pub(in crate::serve) fn validate_document_content(
620 content: Option<&str>,
621) -> Result<DocumentContent, Failure> {
622 match content {
623 None | Some("full") => Ok(DocumentContent::Full),
624 Some("none") => Ok(DocumentContent::None),
625 Some(_) => Err(Failure::new(
626 "invalid_document_content",
627 "document content must be full or none",
628 )),
629 }
630}
631
632pub(in crate::serve) fn list_documents(
637 access: &Access,
638 reader: &ReadOnlyDatabase,
639 at: &Captured,
640 page: &DocumentPage,
641 limit: NonZeroUsize,
642 byte_limit: Option<NonZeroUsize>,
643) -> Result<Page<DocumentListEntry>, Failure> {
644 let content = validate_document_content(page.content.as_deref())?;
645 document_page(
646 access,
647 reader,
648 at,
649 DocumentRange {
650 collection: &page.collection,
651 prefix: page.prefix.as_deref(),
652 after: page.after.as_deref(),
653 metadata: matches!(content, DocumentContent::None),
654 },
655 DocumentLimits {
656 rows: limit,
657 bytes: byte_limit,
658 },
659 |reader, at, collection, id, bytes| match content {
660 DocumentContent::Full => {
661 render_document(reader, at, collection, id, bytes).map(DocumentListEntry::Full)
662 }
663 DocumentContent::None => Ok(DocumentListEntry::Metadata {
664 id,
665 revision: serde_json::from_slice(bytes).map_err(Failure::storage)?,
666 }),
667 },
668 )
669}
670
671pub(in crate::serve) fn validate_document_prefix(
676 prefix: Option<&str>,
677 after: Option<&str>,
678) -> Result<(), Failure> {
679 if let Some(prefix) = prefix {
680 identifier(prefix)?;
681 if let Some(after) = after {
682 identifier(after)?;
683 if !after.as_bytes().starts_with(prefix.as_bytes()) {
684 return Err(Failure::new(
685 "document_after_outside_prefix",
686 "document continuation does not start with the requested prefix",
687 ));
688 }
689 }
690 }
691 Ok(())
692}
693
694#[derive(Clone, Copy)]
695struct DocumentRange<'a> {
696 collection: &'a str,
697 prefix: Option<&'a str>,
698 after: Option<&'a str>,
699 metadata: bool,
700}
701
702#[derive(Clone, Copy)]
703struct DocumentLimits {
704 rows: NonZeroUsize,
705 bytes: Option<NonZeroUsize>,
706}
707
708const DOCUMENT_SCAN_ROWS_PER_LIMIT: usize = 2;
711
712fn document_page<T: Serialize>(
713 access: &Access,
714 reader: &ReadOnlyDatabase,
715 at: &Captured,
716 range: DocumentRange<'_>,
717 limits: DocumentLimits,
718 make_entry: impl Fn(&ReadOnlyDatabase, &Captured, &str, String, &[u8]) -> Result<T, Failure>,
719) -> Result<Page<T>, Failure> {
720 let DocumentRange {
721 collection,
722 prefix: id_prefix,
723 after,
724 metadata,
725 } = range;
726 identifier(collection)?;
727 validate_document_prefix(id_prefix, after)?;
728 let family = if metadata { "m" } else { "d" };
729 let prefix = format!("{family}\0{collection}\0").into_bytes();
730 let mut range_prefix = prefix.clone();
731 if let Some(id_prefix) = id_prefix {
732 range_prefix.extend_from_slice(id_prefix.as_bytes());
733 }
734 let to = prefix_end(&range_prefix)?;
735 let from = match after {
736 Some(after) => exclusive_from(&prefix, Some(after))?,
737 None => range_prefix,
738 };
739 let examined_limit = limits
740 .rows
741 .get()
742 .checked_mul(DOCUMENT_SCAN_ROWS_PER_LIMIT)
743 .and_then(|rows| rows.checked_add(1))
744 .and_then(|rows| rows.checked_mul(at.shards.len()))
745 .ok_or_else(|| Failure::new("limit_overflow", "document scan limit is too large"))?;
746 let mut examined = 0;
747 let mut candidates = Vec::with_capacity(at.shards.len());
748 for shard in &at.shards {
749 candidates.push((
750 shard,
751 document_candidate(reader, shard, &from, &to, &mut examined, examined_limit)?,
752 ));
753 }
754 let mut entries = Vec::new();
755 let mut last_id = None;
756 let mut budget = limits.bytes.map(WireBudget::new);
757 let next_after = loop {
758 let Some(index) = candidates
759 .iter()
760 .enumerate()
761 .filter_map(|(index, (_, row))| row.as_ref().map(|row| (index, &row.row.0)))
762 .min_by(|(_, left), (_, right)| left.cmp(right))
763 .map(|(index, _)| index)
764 else {
765 break None;
766 };
767 let DocumentCandidate {
768 row: (key, value),
769 next_from,
770 } = candidates[index]
771 .1
772 .take()
773 .ok_or_else(|| Failure::new("storage_unavailable", "page candidate disappeared"))?;
774 let suffix = key
775 .strip_prefix(prefix.as_slice())
776 .ok_or_else(|| Failure::new("corrupt_key", "row is outside requested prefix"))?;
777 let id = std::str::from_utf8(suffix).map_err(Failure::storage)?;
778 if access.can_read_document(collection, id) {
779 if entries.len() == limits.rows.get() {
780 break last_id;
781 }
782 let entry = make_entry(reader, at, collection, id.to_owned(), &value)?;
783 if let Some(budget) = budget.as_mut() {
784 let bytes = serde_json::to_vec(&entry).map_err(Failure::storage)?.len();
785 if !budget.admit_len(bytes)? {
786 break last_id;
787 }
788 }
789 last_id = Some(id.to_owned());
790 entries.push(entry);
791 }
792 candidates[index].1 = match next_from {
793 Some(from) => document_candidate(
794 reader,
795 candidates[index].0,
796 &from,
797 &to,
798 &mut examined,
799 examined_limit,
800 )?,
801 None => None,
802 };
803 };
804 Ok(Page {
805 branch: at.branch.clone(),
806 version: at.version,
807 entries,
808 next_after,
809 })
810}
811
812struct DocumentCandidate {
813 row: RawRow,
814 next_from: Option<Vec<u8>>,
815}
816
817fn document_candidate(
818 reader: &ReadOnlyDatabase,
819 shard: &haematite::branch::ShardRef,
820 from: &[u8],
821 to: &[u8],
822 examined: &mut usize,
823 limit: usize,
824) -> Result<Option<DocumentCandidate>, Failure> {
825 if *examined == limit {
828 return Err(Failure::new(
829 "access_scan_limit",
830 "document visibility scan exceeds the page examination limit",
831 ));
832 }
833 #[cfg(test)]
834 record_document_range(from, to);
835 #[cfg(test)]
836 count_range_read();
837 #[cfg(test)]
838 count_content_read(from);
839 let page = reader
840 .range_page_at(shard.shard_id, shard.head, from, to, NonZeroUsize::MIN)
841 .map_err(Failure::observer)?;
842 let Some(row) = page.entries.into_iter().next() else {
843 return Ok(None);
844 };
845 *examined += 1;
846 Ok(Some(DocumentCandidate {
847 row,
848 next_from: page.next_from,
849 }))
850}
851
852pub fn collections(
856 reader: &ReadOnlyDatabase,
857 at: &Captured,
858 after: Option<&str>,
859 limit: NonZeroUsize,
860) -> Result<Page<Collection>, Failure> {
861 let (items, next_after) = page(reader, at, b"c\0", after, limit)?;
862 let entries = items
863 .into_iter()
864 .map(|(_, bytes)| serde_json::from_slice(&bytes).map_err(Failure::storage))
865 .collect::<Result<Vec<_>, _>>()?;
866 Ok(Page {
867 branch: at.branch.clone(),
868 version: at.version,
869 entries,
870 next_after,
871 })
872}
873
874pub fn collections_bounded(
882 access: &Access,
883 reader: &ReadOnlyDatabase,
884 at: &Captured,
885 after: Option<&str>,
886 limit: NonZeroUsize,
887 entry_byte_limit: NonZeroUsize,
888) -> Result<Page<CollectionEntry>, Failure> {
889 let view = super::counts::View::new(reader, at)?;
890 let (entries, next_after) = bounded_page(
891 reader,
892 at,
893 b"c\0",
894 after,
895 PageLimits {
896 rows: limit,
897 bytes: entry_byte_limit,
898 },
899 |_, bytes| {
900 let row: Collection = serde_json::from_slice(&bytes).map_err(Failure::storage)?;
901 let total = view.documents(access, &row.name)?;
902 render_collection(row, &total)
903 },
904 |entry| {
905 serde_json::to_vec(entry)
906 .map(|bytes| bytes.len())
907 .map_err(Failure::storage)
908 },
909 )?;
910 Ok(Page {
911 branch: at.branch.clone(),
912 version: at.version,
913 entries,
914 next_after,
915 })
916}
917
918type NamedBytes = Vec<(String, Vec<u8>)>;
919
920fn page(
921 reader: &ReadOnlyDatabase,
922 at: &Captured,
923 prefix: &[u8],
924 after: Option<&str>,
925 limit: NonZeroUsize,
926) -> Result<(NamedBytes, Option<String>), Failure> {
927 let mut from = prefix.to_vec();
928 if let Some(after) = after {
929 identifier(after)?;
930 from.extend_from_slice(after.as_bytes());
931 from.push(0); }
933 let mut to = prefix.to_vec();
934 let last = to
935 .last_mut()
936 .ok_or_else(|| Failure::new("invalid_prefix", "empty key prefix"))?;
937 *last = last
938 .checked_add(1)
939 .ok_or_else(|| Failure::new("invalid_prefix", "unbounded key prefix"))?;
940 let fetch = limit
941 .checked_add(1)
942 .ok_or_else(|| Failure::new("limit_overflow", "page limit is too large"))?;
943 let mut rows = Vec::new();
944 for shard in &at.shards {
945 let page = reader
946 .range_page_at(shard.shard_id, shard.head, &from, &to, fetch)
947 .map_err(Failure::observer)?;
948 rows.extend(page.entries);
949 }
950 rows.sort_unstable_by(|left, right| left.0.cmp(&right.0));
951 let more = rows.len() > limit.get();
952 rows.truncate(limit.get());
953 let items = rows
954 .into_iter()
955 .map(|(key, value)| {
956 let suffix = key
957 .strip_prefix(prefix)
958 .ok_or_else(|| Failure::new("corrupt_key", "row is outside requested prefix"))?;
959 let name = std::str::from_utf8(suffix)
960 .map_err(Failure::storage)?
961 .to_owned();
962 Ok((name, value))
963 })
964 .collect::<Result<NamedBytes, Failure>>()?;
965 let next_after = if more {
966 items.last().map(|(id, _)| id.clone())
967 } else {
968 None
969 };
970 Ok((items, next_after))
971}
972
973#[derive(Clone, Copy)]
974pub struct PageLimits {
975 pub rows: NonZeroUsize,
976 pub bytes: NonZeroUsize,
977}
978
979fn bounded_page<T>(
980 reader: &ReadOnlyDatabase,
981 at: &Captured,
982 prefix: &[u8],
983 after: Option<&str>,
984 limits: PageLimits,
985 decode: impl FnMut(String, Vec<u8>) -> Result<T, Failure>,
986 serialized_len: impl FnMut(&T) -> Result<usize, Failure>,
987) -> Result<(Vec<T>, Option<String>), Failure> {
988 let from = exclusive_from(prefix, after)?;
989 bounded_key_page(reader, at, prefix, &from, limits, decode, serialized_len)
990}
991
992pub(super) fn bounded_key_page<T>(
995 reader: &ReadOnlyDatabase,
996 at: &Captured,
997 prefix: &[u8],
998 from: &[u8],
999 limits: PageLimits,
1000 mut decode: impl FnMut(String, Vec<u8>) -> Result<T, Failure>,
1001 mut serialized_len: impl FnMut(&T) -> Result<usize, Failure>,
1002) -> Result<(Vec<T>, Option<String>), Failure> {
1003 let to = prefix_end(prefix)?;
1004 if !from.starts_with(prefix) || from >= to.as_slice() {
1005 return Err(Failure::new(
1006 "invalid_continuation",
1007 "page boundary is outside the requested prefix",
1008 ));
1009 }
1010 let mut candidates = Vec::with_capacity(at.shards.len());
1011 for shard in &at.shards {
1012 candidates.push((
1013 shard.shard_id,
1014 first_row(reader, shard.shard_id, shard.head, from, &to)?,
1015 ));
1016 }
1017
1018 let mut entries = Vec::new();
1019 let mut last_id = None;
1020 let mut budget = WireBudget::new(limits.bytes);
1021 loop {
1022 let Some(index) = candidates
1023 .iter()
1024 .enumerate()
1025 .filter_map(|(index, (_, row))| row.as_ref().map(|row| (index, &row.0)))
1026 .min_by(|(_, left), (_, right)| left.cmp(right))
1027 .map(|(index, _)| index)
1028 else {
1029 return Ok((entries, None));
1030 };
1031 let (key, value) = candidates[index]
1032 .1
1033 .take()
1034 .ok_or_else(|| Failure::new("storage_unavailable", "page candidate disappeared"))?;
1035 let suffix = key
1036 .strip_prefix(prefix)
1037 .ok_or_else(|| Failure::new("corrupt_key", "row is outside requested prefix"))?;
1038 let id = std::str::from_utf8(suffix)
1039 .map_err(Failure::storage)?
1040 .to_owned();
1041 let entry = decode(id.clone(), value)?;
1042 if !budget.admit_len(serialized_len(&entry)?)? {
1043 return Ok((entries, last_id));
1044 }
1045 entries.push(entry);
1046 last_id = Some(id);
1047
1048 let shard_id = candidates[index].0;
1049 let shard = at
1050 .shards
1051 .iter()
1052 .find(|shard| shard.shard_id == shard_id)
1053 .ok_or_else(|| Failure::new("incomplete_roots", "page shard disappeared"))?;
1054 let mut next_from = key;
1055 next_from.push(0);
1056 candidates[index].1 = first_row(reader, shard_id, shard.head, &next_from, &to)?;
1057
1058 if entries.len() == limits.rows.get() {
1059 let next_after = if candidates.iter().any(|(_, row)| row.is_some()) {
1060 last_id
1061 } else {
1062 None
1063 };
1064 return Ok((entries, next_after));
1065 }
1066 }
1067}
1068
1069#[derive(Clone, Copy)]
1070pub(super) struct VisibleKeyPage<'a> {
1071 pub access: &'a Access,
1072 pub collection: &'a str,
1073 pub prefix: &'a [u8],
1074 pub from: &'a [u8],
1075 pub limits: PageLimits,
1076}
1077
1078pub(super) fn visible_key_page<T>(
1080 reader: &ReadOnlyDatabase,
1081 at: &Captured,
1082 request: VisibleKeyPage<'_>,
1083 decode: impl FnMut(String, Vec<u8>) -> Result<T, Failure>,
1084 serialized_len: impl FnMut(&T) -> Result<usize, Failure>,
1085) -> Result<(Vec<T>, Option<String>), Failure> {
1086 visible_row_page(
1087 reader,
1088 at,
1089 VisibleRowPage {
1090 prefix: request.prefix,
1091 from: request.from,
1092 limits: request.limits,
1093 },
1094 |id| {
1095 identifier(id)?;
1096 Ok(request.access.can_read_document(request.collection, id))
1097 },
1098 decode,
1099 serialized_len,
1100 )
1101}
1102
1103#[derive(Clone, Copy)]
1104pub(super) struct VisibleRowPage<'a> {
1105 pub prefix: &'a [u8],
1106 pub from: &'a [u8],
1107 pub limits: PageLimits,
1108}
1109
1110pub(super) fn visible_row_page<T>(
1112 reader: &ReadOnlyDatabase,
1113 at: &Captured,
1114 request: VisibleRowPage<'_>,
1115 mut visible: impl FnMut(&str) -> Result<bool, Failure>,
1116 mut decode: impl FnMut(String, Vec<u8>) -> Result<T, Failure>,
1117 mut serialized_len: impl FnMut(&T) -> Result<usize, Failure>,
1118) -> Result<(Vec<T>, Option<String>), Failure> {
1119 let to = prefix_end(request.prefix)?;
1120 if !request.from.starts_with(request.prefix) || request.from >= to.as_slice() {
1121 return Err(Failure::new(
1122 "invalid_continuation",
1123 "page boundary is outside the requested prefix",
1124 ));
1125 }
1126 let examined_limit = request
1127 .limits
1128 .rows
1129 .get()
1130 .checked_mul(DOCUMENT_SCAN_ROWS_PER_LIMIT)
1131 .and_then(|rows| rows.checked_add(1))
1132 .and_then(|rows| rows.checked_mul(at.shards.len()))
1133 .ok_or_else(|| Failure::new("limit_overflow", "visibility scan limit is too large"))?;
1134 let mut examined = 0;
1135 let mut candidates = Vec::with_capacity(at.shards.len());
1136 for shard in &at.shards {
1137 candidates.push((
1138 shard,
1139 document_candidate(
1140 reader,
1141 shard,
1142 request.from,
1143 &to,
1144 &mut examined,
1145 examined_limit,
1146 )?,
1147 ));
1148 }
1149 let mut entries = Vec::new();
1150 let mut last_id = None;
1151 let mut budget = WireBudget::new(request.limits.bytes);
1152 loop {
1153 let Some(index) = candidates
1154 .iter()
1155 .enumerate()
1156 .filter_map(|(index, (_, row))| row.as_ref().map(|row| (index, &row.row.0)))
1157 .min_by(|(_, left), (_, right)| left.cmp(right))
1158 .map(|(index, _)| index)
1159 else {
1160 return Ok((entries, None));
1161 };
1162 let DocumentCandidate {
1163 row: (key, value),
1164 next_from,
1165 } = candidates[index]
1166 .1
1167 .take()
1168 .ok_or_else(|| Failure::new("storage_unavailable", "page candidate disappeared"))?;
1169 let suffix = key
1170 .strip_prefix(request.prefix)
1171 .ok_or_else(|| Failure::new("corrupt_key", "row is outside requested prefix"))?;
1172 let id = std::str::from_utf8(suffix).map_err(Failure::storage)?;
1173 if visible(id)? {
1174 if entries.len() == request.limits.rows.get() {
1175 return Ok((entries, last_id));
1176 }
1177 let entry = decode(id.to_owned(), value)?;
1178 if !budget.admit_len(serialized_len(&entry)?)? {
1179 return Ok((entries, last_id));
1180 }
1181 last_id = Some(id.to_owned());
1182 entries.push(entry);
1183 }
1184 candidates[index].1 = match next_from {
1185 Some(from) => document_candidate(
1186 reader,
1187 candidates[index].0,
1188 &from,
1189 &to,
1190 &mut examined,
1191 examined_limit,
1192 )?,
1193 None => None,
1194 };
1195 }
1196}
1197
1198pub(super) fn prefix_keys(
1200 reader: &ReadOnlyDatabase,
1201 at: &Captured,
1202 prefix: &[u8],
1203 limit: NonZeroUsize,
1204) -> Result<Vec<Vec<u8>>, Failure> {
1205 let to = prefix_end(prefix)?;
1206 let mut candidates = Vec::with_capacity(at.shards.len());
1207 for shard in &at.shards {
1208 candidates.push((
1209 shard,
1210 first_row(reader, shard.shard_id, shard.head, prefix, &to)?,
1211 ));
1212 }
1213 let mut keys = Vec::new();
1214 while keys.len() < limit.get() {
1215 let Some(index) = candidates
1216 .iter()
1217 .enumerate()
1218 .filter_map(|(index, (_, row))| row.as_ref().map(|row| (index, &row.0)))
1219 .min_by(|(_, left), (_, right)| left.cmp(right))
1220 .map(|(index, _)| index)
1221 else {
1222 break;
1223 };
1224 let (key, _) = candidates[index].1.take().ok_or_else(|| {
1225 Failure::new("storage_unavailable", "retention key candidate disappeared")
1226 })?;
1227 if !key.starts_with(prefix) || keys.last() == Some(&key) {
1228 return Err(Failure::new(
1229 "corrupt_key",
1230 "retention key is outside its scope or duplicated",
1231 ));
1232 }
1233 let mut from = key.clone();
1234 from.push(0);
1235 keys.push(key);
1236 if keys.len() == limit.get() {
1237 break;
1238 }
1239 let shard = candidates[index].0;
1240 candidates[index].1 = first_row(reader, shard.shard_id, shard.head, &from, &to)?;
1241 }
1242 Ok(keys)
1243}
1244
1245fn exclusive_from(prefix: &[u8], after: Option<&str>) -> Result<Vec<u8>, Failure> {
1246 let mut from = prefix.to_vec();
1247 if let Some(after) = after {
1248 identifier(after)?;
1249 from.extend_from_slice(after.as_bytes());
1250 from.push(0);
1251 }
1252 Ok(from)
1253}
1254
1255fn prefix_end(prefix: &[u8]) -> Result<Vec<u8>, Failure> {
1256 let mut to = prefix.to_vec();
1257 let last = to
1258 .last_mut()
1259 .ok_or_else(|| Failure::new("invalid_prefix", "empty key prefix"))?;
1260 *last = last
1261 .checked_add(1)
1262 .ok_or_else(|| Failure::new("invalid_prefix", "unbounded key prefix"))?;
1263 Ok(to)
1264}
1265
1266type RawRow = (Vec<u8>, Vec<u8>);
1267
1268fn first_row(
1269 reader: &ReadOnlyDatabase,
1270 shard_id: usize,
1271 root: haematite::tree::Hash,
1272 from: &[u8],
1273 to: &[u8],
1274) -> Result<Option<RawRow>, Failure> {
1275 #[cfg(test)]
1276 count_range_read();
1277 #[cfg(test)]
1278 count_content_read(from);
1279 Ok(reader
1280 .range_page_at(shard_id, root, from, to, NonZeroUsize::MIN)
1281 .map_err(Failure::observer)?
1282 .entries
1283 .into_iter()
1284 .next())
1285}
1286
1287#[cfg(test)]
1288#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1289pub(in crate::serve) struct PageReadCounts {
1290 pub points: usize,
1291 pub ranges: usize,
1292}
1293
1294#[cfg(test)]
1295thread_local! {
1296 static DOCUMENT_RANGES: std::cell::RefCell<Vec<(Vec<u8>, Vec<u8>)>> = const {
1297 std::cell::RefCell::new(Vec::new())
1298 };
1299 static CONTENT_READS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
1300 static PAGE_READ_COUNTS: std::cell::Cell<PageReadCounts> = const {
1301 std::cell::Cell::new(PageReadCounts { points: 0, ranges: 0 })
1302 };
1303}
1304
1305#[cfg(test)]
1306pub(super) fn take_content_reads() -> usize {
1307 CONTENT_READS.with(std::cell::Cell::take)
1308}
1309
1310#[cfg(test)]
1311fn count_content_read(key: &[u8]) {
1312 if key.starts_with(b"d\0") || key.starts_with(b"e\0") || key.starts_with(b"v\0") {
1313 CONTENT_READS.with(|reads| reads.set(reads.get() + 1));
1314 }
1315}
1316
1317#[cfg(test)]
1318pub(in crate::serve) fn reset_page_read_counts() {
1319 PAGE_READ_COUNTS.with(|counts| counts.set(PageReadCounts::default()));
1320}
1321
1322#[cfg(test)]
1323pub(in crate::serve) fn page_read_counts() -> PageReadCounts {
1324 PAGE_READ_COUNTS.with(std::cell::Cell::get)
1325}
1326
1327#[cfg(test)]
1328fn count_point_read() {
1329 PAGE_READ_COUNTS.with(|counts| {
1330 let mut current = counts.get();
1331 current.points += 1;
1332 counts.set(current);
1333 });
1334}
1335
1336#[cfg(test)]
1337fn count_range_read() {
1338 PAGE_READ_COUNTS.with(|counts| {
1339 let mut current = counts.get();
1340 current.ranges += 1;
1341 counts.set(current);
1342 });
1343}
1344
1345struct WireBudget {
1346 limit: usize,
1347 used: usize,
1348}
1349
1350impl WireBudget {
1351 const fn new(limit: NonZeroUsize) -> Self {
1352 Self {
1353 limit: limit.get(),
1354 used: 0,
1355 }
1356 }
1357
1358 #[cfg(test)]
1359 fn admit(&mut self, entry: &impl Serialize) -> Result<bool, Failure> {
1360 let bytes = serde_json::to_vec(entry).map_err(Failure::storage)?.len();
1361 self.admit_len(bytes)
1362 }
1363
1364 fn admit_len(&mut self, bytes: usize) -> Result<bool, Failure> {
1365 let separator = usize::from(self.used != 0);
1366 let required = bytes
1367 .checked_add(separator)
1368 .ok_or_else(|| Failure::new("response_limit", "response byte count overflow"))?;
1369 if required > self.limit.saturating_sub(self.used) {
1370 if self.used == 0 {
1371 return Err(Failure::new(
1372 "response_limit",
1373 "one row exceeds the response byte budget",
1374 ));
1375 }
1376 return Ok(false);
1377 }
1378 self.used += required;
1379 Ok(true)
1380 }
1381}
1382
1383pub(super) fn render_document(
1384 reader: &ReadOnlyDatabase,
1385 at: &Captured,
1386 collection: &str,
1387 id: String,
1388 bytes: &[u8],
1389) -> Result<Document, Failure> {
1390 let raw: Box<RawValue> = serde_json::from_slice(bytes).map_err(Failure::storage)?;
1391 let revision = get_bytes(reader, at, &revision_key(collection, &id))?.ok_or_else(|| {
1392 Failure::new(
1393 "revision_missing",
1394 "document has no committed schema revision",
1395 )
1396 })?;
1397 Ok(Document {
1398 id,
1399 raw_json: raw.get().to_owned(),
1400 revision: serde_json::from_slice(&revision).map_err(Failure::storage)?,
1401 })
1402}
1403
1404#[cfg(test)]
1405#[path = "reads_tests.rs"]
1406mod tests;
1407
1408#[cfg(test)]
1409fn record_document_range(from: &[u8], to: &[u8]) {
1410 DOCUMENT_RANGES.with(|ranges| ranges.borrow_mut().push((from.to_vec(), to.to_vec())));
1411}
1412
1413#[cfg(test)]
1414fn take_document_ranges() -> Vec<(Vec<u8>, Vec<u8>)> {
1415 DOCUMENT_RANGES.with(|ranges| std::mem::take(&mut *ranges.borrow_mut()))
1416}
1417
1418#[cfg(test)]
1419#[path = "reads_prefix_tests.rs"]
1420mod prefix_tests;