1use super::*;
4use crate::storage::collection_read::{
5 CollectionCursor, CollectionRead, FrozenOverlay, MergedCollection,
6};
7use crate::storage::error::StoreError;
8use crate::storage::metadata_codec::serialize_entity_record;
9use std::collections::HashSet;
10
11pub(super) struct BackingRead {
12 pub source: Arc<dyn CollectionRead + Send + Sync>,
13 pub root_page: u32,
14 pub total_entities: u64,
15 pub format_version: u32,
16 pub read_bytes_max: usize,
17 pub removed: HashSet<EntityId>,
18}
19
20pub struct PagedEntityCursor {
23 cursor: Box<dyn CollectionCursor>,
24 format_version: u32,
25 read_bytes_max: usize,
26 snapshot: Option<SnapshotContext>,
27 visible_only: bool,
28 schema: Option<Arc<Vec<String>>>,
29 key: Option<String>,
30 finished: bool,
31}
32
33impl Iterator for PagedEntityCursor {
34 type Item = Result<(UnifiedEntity, Option<Metadata>), StoreError>;
35 fn next(&mut self) -> Option<Self::Item> {
36 if self.finished {
37 return None;
38 }
39 loop {
40 let record = match self.cursor.next_record(self.read_bytes_max) {
41 Ok(Some(record)) => record,
42 Ok(None) => {
43 self.finished = true;
44 return None;
45 }
46 Err(error) => {
47 self.finished = true;
48 return Some(Err(error));
49 }
50 };
51 let (mut entity, metadata) =
52 match record.decode_with_budget(self.format_version, self.read_bytes_max) {
53 Ok(record) => record,
54 Err(error) => {
55 self.finished = true;
56 return Some(Err(error));
57 }
58 };
59 if self.visible_only
60 && !SegmentManager::scan_entity_visible(self.snapshot.as_ref(), &entity)
61 {
62 continue;
63 }
64 if let super::super::entity::EntityData::Row(row) = &mut entity.data {
65 if row.named.is_none() && row.schema.is_none() {
66 row.schema = self.schema.clone();
67 }
68 }
69 if self.key.as_deref().is_some_and(|key| {
70 crate::storage::segment_key_lookup::SegmentKeyLookup::key(&entity) != Some(key)
71 }) {
72 continue;
73 }
74 return Some(Ok((entity, metadata)));
75 }
76 }
77}
78
79impl SegmentManager {
80 pub fn try_for_each_entity_unordered(
86 &self,
87 limit_bytes: usize,
88 mut callback: impl FnMut(&UnifiedEntity) -> bool,
89 ) -> Result<(), StoreError> {
90 let backing_guard = self.backing.read();
91 let Some(backing) = backing_guard.as_ref() else {
92 self.for_each_entity(callback);
93 return Ok(());
94 };
95 let limit_bytes = limit_bytes.min(backing.read_bytes_max);
96 let mut cursor = backing.source.cursor()?;
97 while let Some(id) = cursor.peek_id()? {
98 if backing.removed.contains(&id) || self.get_with(id, |_| ()).is_some() {
99 cursor.skip_record()?;
100 continue;
101 }
102 let record = cursor
103 .next_record(limit_bytes)?
104 .expect("invariant: peeked physical row exists");
105 let (mut entity, _) = record.decode_with_budget(backing.format_version, limit_bytes)?;
106 self.restore_owned_row_schemas(std::iter::once(&mut entity));
107 if !callback(&entity) {
108 return Ok(());
109 }
110 }
111 let mut result = Ok(());
112 self.for_each_entity(|entity| {
113 let required_bytes = crate::storage::memory_size::entity_bytes(entity);
114 if required_bytes > limit_bytes {
115 result = Err(StoreError::ReadLimitExceeded {
116 required_bytes,
117 limit_bytes,
118 });
119 return false;
120 }
121 callback(entity)
122 });
123 result
124 }
125
126 pub fn try_paged_entity_cursor(
129 &self,
130 snapshot: Option<SnapshotContext>,
131 ) -> Result<Option<PagedEntityCursor>, StoreError> {
132 self.try_paged_cursor(snapshot, true, None)
133 }
134
135 pub fn try_paged_physical_cursor(&self) -> Result<Option<PagedEntityCursor>, StoreError> {
138 self.try_paged_cursor(None, false, None)
139 }
140
141 pub fn try_paged_key_entity_cursor(
144 &self,
145 key: &str,
146 snapshot: Option<SnapshotContext>,
147 ) -> Result<Option<PagedEntityCursor>, StoreError> {
148 self.try_paged_cursor(snapshot, true, Some(key))
149 }
150
151 pub fn try_paged_key_physical_cursor(
153 &self,
154 key: &str,
155 ) -> Result<Option<PagedEntityCursor>, StoreError> {
156 self.try_paged_cursor(None, false, Some(key))
157 }
158
159 fn try_paged_cursor(
160 &self,
161 snapshot: Option<SnapshotContext>,
162 visible_only: bool,
163 key: Option<&str>,
164 ) -> Result<Option<PagedEntityCursor>, StoreError> {
165 let Some((format_version, read_bytes_max)) = self
166 .backing
167 .read()
168 .as_ref()
169 .map(|base| (base.format_version, base.read_bytes_max))
170 else {
171 return Ok(None);
172 };
173 let view = match key {
174 Some(key) => self.current_backing_key_overlay(key)?,
175 None => self.current_backing_overlay(),
176 };
177 let Some((base, overlay)) = view else {
178 return Ok(None);
179 };
180 let merged = MergedCollection::new(base.as_ref(), overlay);
181 let cursor = match key {
182 Some(key) => match merged.owned_key_cursor(key)? {
183 Some(cursor) => cursor,
184 None => merged.owned_cursor()?,
185 },
186 None => merged.owned_cursor()?,
187 };
188 Ok(Some(PagedEntityCursor {
189 cursor,
190 format_version,
191 read_bytes_max,
192 snapshot,
193 visible_only,
194 schema: self.column_schema(),
195 key: key.map(str::to_owned),
196 finished: false,
197 }))
198 }
199
200 pub fn has_backing(&self) -> bool {
201 self.backing.read().is_some()
202 }
203
204 pub fn backing_root_page(&self) -> Option<u32> {
206 self.backing.read().as_ref().map(|base| base.root_page)
207 }
208
209 pub fn backing_read_bytes_max(&self) -> Option<usize> {
212 self.backing
213 .read()
214 .as_ref()
215 .map(|backing| backing.read_bytes_max)
216 }
217
218 pub fn install_backing(
221 &self,
222 source: Arc<dyn CollectionRead + Send + Sync>,
223 root_page: u32,
224 format_version: u32,
225 read_bytes_max: usize,
226 total_entities: u64,
227 ) {
228 let mut backing = self.backing.write();
229 let mut consolidation = self.consolidation.write();
230 let mut growing = self.growing.write();
231 let mut sealed = self.sealed.write();
232 let mut retired = self.retired.write();
233 if let Some(source) = growing.take() {
234 self.restore_sequence_id_next(source.read().sequence_id_next());
235 retired.push(Arc::downgrade(&source));
236 }
237 for source in sealed.drain(..) {
238 self.restore_sequence_id_next(source.read().sequence_id_next());
239 retired.push(Arc::downgrade(&source));
240 }
241 *consolidation = None;
242 self.entity_segment.write().clear();
243 self.total_entities_atomic
244 .store(total_entities, Ordering::Relaxed);
245 *backing = Some(BackingRead {
246 source,
247 root_page,
248 total_entities,
249 format_version,
250 read_bytes_max,
251 removed: HashSet::new(),
252 });
253 }
254
255 #[doc(hidden)]
261 pub fn try_with_unshared_backing<T>(
262 &self,
263 root_page: u32,
264 update: impl FnOnce() -> (T, Option<(u32, u64)>),
265 ) -> Option<T> {
266 let mut backing = self.backing.write();
267 let backing = backing.as_mut()?;
268 if backing.root_page != root_page
269 || !backing.removed.is_empty()
270 || Arc::get_mut(&mut backing.source).is_none()
271 {
272 return None;
273 }
274 let (result, published) = update();
275 if let Some((root_page, total_entities)) = published {
276 backing.root_page = root_page;
277 backing.total_entities = total_entities;
278 }
279 Some(result)
280 }
281
282 pub fn backing_count_at_root(&self, root_page: u32) -> Option<u64> {
283 self.backing
284 .read()
285 .as_ref()
286 .filter(|base| base.root_page == root_page)
287 .map(|base| base.total_entities)
288 }
289
290 pub fn backing_row_is_shadowed(&self, id: EntityId) -> bool {
294 let backing = self.backing.read();
295 backing
296 .as_ref()
297 .is_none_or(|base| base.removed.contains(&id))
298 || self.get_with(id, |_| ()).is_some()
299 }
300
301 pub fn try_get(
303 &self,
304 id: EntityId,
305 limit_bytes: usize,
306 ) -> Result<Option<UnifiedEntity>, StoreError> {
307 let backing = self.backing.read();
308 let limit_bytes = backing
309 .as_ref()
310 .map_or(limit_bytes, |base| limit_bytes.min(base.read_bytes_max));
311 if let Some(result) = self.get_with(id, |entity| {
312 let required_bytes = if limit_bytes == usize::MAX {
313 0
314 } else {
315 crate::storage::memory_size::entity_bytes(entity)
316 };
317 if required_bytes > limit_bytes {
318 Err(StoreError::ReadLimitExceeded {
319 required_bytes,
320 limit_bytes,
321 })
322 } else {
323 Ok(entity.clone())
324 }
325 }) {
326 let mut entity = result?;
327 self.restore_owned_row_schemas(std::iter::once(&mut entity));
328 return Ok(Some(entity));
329 }
330 let Some(backing) = backing.as_ref() else {
331 return Ok(None);
332 };
333 if backing.removed.contains(&id) {
334 return Ok(None);
335 }
336 let Some(record) = backing
337 .source
338 .get_record(id, limit_bytes.min(backing.read_bytes_max))?
339 else {
340 return Ok(None);
341 };
342 let (mut entity, _) = record.decode_with_budget(
343 backing.format_version,
344 limit_bytes.min(backing.read_bytes_max),
345 )?;
346 self.restore_owned_row_schemas(std::iter::once(&mut entity));
347 Ok(Some(entity))
348 }
349
350 pub fn try_project_key_entity<T>(
355 &self,
356 key: &str,
357 snapshot: Option<&SnapshotContext>,
358 versioned: bool,
359 limit_bytes: usize,
360 mut project: impl FnMut(&UnifiedEntity) -> T,
361 ) -> Result<Option<T>, StoreError> {
362 let backing = self.backing.read();
363 let limit_bytes = backing
364 .as_ref()
365 .map_or(limit_bytes, |base| limit_bytes.min(base.read_bytes_max));
366 let mut best = self
367 .project_key_entity(key, snapshot, versioned, |entity| {
368 let required_bytes = if limit_bytes == usize::MAX {
369 0
370 } else {
371 crate::storage::memory_size::entity_bytes(entity)
372 };
373 if required_bytes > limit_bytes {
374 Err(StoreError::ReadLimitExceeded {
375 required_bytes,
376 limit_bytes,
377 })
378 } else {
379 Ok((entity.xmin, project(entity)))
380 }
381 })
382 .transpose()?;
383 let Some(backing) = backing.as_ref() else {
384 return Ok(best.map(|(_, projected)| projected));
385 };
386 if !versioned && best.is_some() {
387 return Ok(best.map(|(_, projected)| projected));
388 }
389 let mut cursor = match backing.source.key_cursor(key)? {
390 Some(cursor) => cursor,
391 None => backing.source.cursor()?,
392 };
393 while let Some(id) = cursor.peek_id()? {
394 if backing.removed.contains(&id) || self.get_with(id, |_| ()).is_some() {
397 cursor.skip_record()?;
398 continue;
399 }
400 let record = cursor
401 .next_record(limit_bytes)?
402 .expect("invariant: peeked physical row exists");
403 let (mut entity, _) = record.decode_with_budget(backing.format_version, limit_bytes)?;
404 self.restore_owned_row_schemas(std::iter::once(&mut entity));
405 if crate::storage::segment_key_lookup::SegmentKeyLookup::key(&entity) != Some(key)
406 || (versioned && !Self::scan_entity_visible(snapshot, &entity))
407 {
408 continue;
409 }
410 if best.as_ref().is_none_or(|(xmin, _)| entity.xmin >= *xmin) {
411 best = Some((entity.xmin, project(&entity)));
412 }
413 if !versioned {
414 break;
415 }
416 }
417 Ok(best.map(|(_, projected)| projected))
418 }
419
420 pub fn try_for_each_entity(
424 &self,
425 limit_bytes: usize,
426 mut callback: impl FnMut(&UnifiedEntity) -> bool,
427 ) -> Result<(), StoreError> {
428 let backing_guard = self.backing.read();
429 let Some(backing) = backing_guard.as_ref() else {
430 self.for_each_entity(callback);
431 return Ok(());
432 };
433 let limit_bytes = limit_bytes.min(backing.read_bytes_max);
434 let mut identities = Vec::new();
435 self.for_each_entity(|entity| {
436 identities.push(entity.id);
437 true
438 });
439 identities.sort_unstable();
440 identities.dedup();
441 let mut identities = identities.into_iter().peekable();
442 let mut cursor = backing.source.cursor()?;
443 loop {
444 let base = cursor.peek_id()?;
445 let resident = identities.peek().copied();
446 if let Some(id) = resident.filter(|id| base.is_none_or(|base| *id <= base)) {
447 if base == Some(id) {
448 cursor.skip_record()?;
449 }
450 identities.next();
451 let result = self.get_with(id, |entity| {
452 let required_bytes = crate::storage::memory_size::entity_bytes(entity);
453 if required_bytes > limit_bytes {
454 return Err(StoreError::ReadLimitExceeded {
455 required_bytes,
456 limit_bytes,
457 });
458 }
459 Ok(callback(entity))
460 });
461 if result.transpose()? == Some(false) {
462 return Ok(());
463 }
464 continue;
465 }
466 let Some(id) = base else {
467 return Ok(());
468 };
469 if backing.removed.contains(&id) {
470 cursor.skip_record()?;
471 continue;
472 }
473 let record = cursor
474 .next_record(limit_bytes)?
475 .expect("invariant: peeked physical row exists");
476 let (mut entity, _) = record.decode_with_budget(
477 backing.format_version,
478 limit_bytes.min(backing.read_bytes_max),
479 )?;
480 self.restore_owned_row_schemas(std::iter::once(&mut entity));
481 if !callback(&entity) {
482 return Ok(());
483 }
484 }
485 }
486
487 pub fn try_for_each_key_entity(
490 &self,
491 key: &str,
492 mut callback: impl FnMut(&UnifiedEntity) -> bool,
493 ) -> Result<(), StoreError> {
494 if let Some(cursor) = self.try_paged_key_physical_cursor(key)? {
495 for record in cursor {
496 let (entity, _) = record?;
497 if !callback(&entity) {
498 break;
499 }
500 }
501 return Ok(());
502 }
503 self.try_for_each_entity(usize::MAX, |entity| {
504 crate::storage::segment_key_lookup::SegmentKeyLookup::key(entity) != Some(key)
505 || callback(entity)
506 })
507 }
508
509 pub fn try_scan_for_each(
510 &self,
511 snapshot: Option<&SnapshotContext>,
512 limit_bytes: usize,
513 mut callback: impl FnMut(&UnifiedEntity) -> bool,
514 ) -> Result<(), StoreError> {
515 self.try_for_each_entity(limit_bytes, |entity| {
516 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
517 })
518 }
519
520 pub fn try_scan_for_each_zoned_with_stats(
523 &self,
524 snapshot: Option<&SnapshotContext>,
525 zone_preds: &[(&str, ZoneColPred<'_>)],
526 limit_bytes: usize,
527 callback: impl FnMut(&UnifiedEntity) -> bool,
528 ) -> Result<SegmentScanStats, StoreError> {
529 if !self.has_backing() {
530 return Ok(self.scan_for_each_zoned_with_stats(snapshot, zone_preds, callback));
531 }
532 self.try_scan_for_each(snapshot, limit_bytes, callback)?;
533 Ok(SegmentScanStats {
534 segments_total: 1,
535 segments_scanned: 1,
536 segments_pruned: 0,
537 })
538 }
539
540 pub fn try_scan_project_zoned_with_stats<T, F>(
543 &self,
544 snapshot: Option<&SnapshotContext>,
545 zone_preds: &[(&str, ZoneColPred<'_>)],
546 read_bytes_max: usize,
547 project: F,
548 ) -> Result<(Vec<T>, SegmentScanStats), StoreError>
549 where
550 T: Send,
551 F: Fn(&UnifiedEntity) -> Option<T> + Sync,
552 {
553 if !self.has_backing() {
554 return Ok(self.scan_project_zoned_with_stats(snapshot, zone_preds, project));
555 }
556 let mut projected = Vec::new();
557 let stats = self.try_scan_for_each_zoned_with_stats(
558 snapshot,
559 zone_preds,
560 read_bytes_max,
561 |entity| {
562 if let Some(value) = project(entity) {
563 projected.push(value);
564 }
565 true
566 },
567 )?;
568 Ok((projected, stats))
569 }
570
571 pub fn try_get_many(
574 &self,
575 ids: &[EntityId],
576 read_bytes_max: usize,
577 ) -> Result<Vec<Option<UnifiedEntity>>, StoreError> {
578 if !self.has_backing() && read_bytes_max == usize::MAX {
579 return Ok(self.get_many(ids));
580 }
581 ids.iter()
582 .map(|id| self.try_get(*id, read_bytes_max))
583 .collect()
584 }
585
586 pub fn try_for_each_id(
587 &self,
588 ids: &[EntityId],
589 read_bytes_max: usize,
590 mut visit: impl FnMut(usize, &UnifiedEntity) -> bool,
591 ) -> Result<(), StoreError> {
592 if self.has_backing() {
593 for (position, id) in ids.iter().enumerate() {
594 if let Some(entity) = self.try_get(*id, read_bytes_max)? {
595 if !visit(position, &entity) {
596 break;
597 }
598 }
599 }
600 return Ok(());
601 }
602 let mut stopped = false;
603 let mut failure = None;
604 self.for_each_id(ids, |position, entity| {
605 if stopped {
606 return;
607 }
608 let required_bytes = if read_bytes_max == usize::MAX {
609 0
610 } else {
611 crate::storage::memory_size::entity_bytes(entity)
612 };
613 if required_bytes > read_bytes_max {
614 failure = Some(StoreError::ReadLimitExceeded {
615 required_bytes,
616 limit_bytes: read_bytes_max,
617 });
618 stopped = true;
619 } else {
620 stopped = !visit(position, entity);
621 }
622 });
623 failure.map_or(Ok(()), Err)
624 }
625
626 pub(super) fn reject_backing_duplicate(&self, id: EntityId) -> Result<(), SegmentError> {
628 let backing = self.backing.read();
629 let Some(backing) = backing.as_ref() else {
630 return Ok(());
631 };
632 if !backing.removed.contains(&id)
633 && backing
634 .source
635 .get_record(id, backing.read_bytes_max)
636 .map_err(|error| SegmentError::BackingRead(error.to_string()))?
637 .is_some()
638 {
639 return Err(SegmentError::AlreadyExists(id));
640 }
641 Ok(())
642 }
643
644 pub fn try_get_metadata(
646 &self,
647 id: EntityId,
648 limit_bytes: usize,
649 ) -> Result<Option<Metadata>, StoreError> {
650 let backing = self.backing.read();
651 if self.get_with(id, |_| ()).is_some() {
652 return Ok(self.get_metadata(id));
653 }
654 let Some(backing) = backing.as_ref() else {
655 return Ok(None);
656 };
657 if backing.removed.contains(&id) {
658 return Ok(None);
659 }
660 backing
661 .source
662 .get_record(id, limit_bytes.min(backing.read_bytes_max))?
663 .map(|record| {
664 record
665 .decode_with_budget(
666 backing.format_version,
667 limit_bytes.min(backing.read_bytes_max),
668 )
669 .map(|(_, metadata)| metadata)
670 })
671 .transpose()
672 .map(Option::flatten)
673 }
674
675 pub(super) fn materialize_backing_entity(&self, id: EntityId) -> Result<(), SegmentError> {
676 if self.get_with(id, |_| ()).is_some() {
677 return Ok(());
678 }
679 let record = {
680 let backing = self.backing.read();
681 let Some(backing) = backing.as_ref() else {
682 return Ok(());
683 };
684 if backing.removed.contains(&id) {
685 return Ok(());
686 }
687 backing
688 .source
689 .get_record(id, backing.read_bytes_max)
690 .and_then(|record| {
691 record
692 .map(|record| {
693 record
694 .decode_with_budget(backing.format_version, backing.read_bytes_max)
695 })
696 .transpose()
697 })
698 .map_err(|error| SegmentError::BackingRead(error.to_string()))?
699 };
700 let Some((mut entity, metadata)) = record else {
701 return Ok(());
702 };
703 self.restore_owned_row_schemas(std::iter::once(&mut entity));
704 self.restore_snapshot_entity(entity)?;
705 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
707 if let Some(metadata) = metadata {
708 self.set_metadata(id, metadata)?;
709 }
710 Ok(())
711 }
712
713 pub(super) fn remember_backing_delete(&self, id: EntityId) {
714 if let Some(backing) = self.backing.write().as_mut() {
715 backing.removed.insert(id);
716 }
717 }
718
719 pub fn backing_overlay(
720 &self,
721 mut encode: impl FnMut(&UnifiedEntity, Option<Metadata>) -> Option<Vec<u8>>,
722 ) -> Option<(Arc<dyn CollectionRead + Send + Sync>, Arc<FrozenOverlay>)> {
723 let backing = self.backing.read();
724 let backing = backing.as_ref()?;
725 let mut overlay: FrozenOverlay = backing.removed.iter().map(|id| (*id, None)).collect();
726 self.for_each_entity_with_metadata(|entity, metadata| {
727 if let Some(bytes) = encode(entity, metadata) {
728 overlay.insert(entity.id, Some(Arc::from(bytes)));
729 }
730 true
731 });
732 Some((Arc::clone(&backing.source), Arc::new(overlay)))
733 }
734
735 pub fn current_backing_overlay(
736 &self,
737 ) -> Option<(Arc<dyn CollectionRead + Send + Sync>, Arc<FrozenOverlay>)> {
738 let version = self.backing.read().as_ref()?.format_version;
739 self.backing_overlay(|entity, metadata| {
740 Some(serialize_entity_record(entity, metadata.as_ref(), version))
741 })
742 }
743
744 fn current_backing_key_overlay(
745 &self,
746 key: &str,
747 ) -> Result<Option<(Arc<dyn CollectionRead + Send + Sync>, Arc<FrozenOverlay>)>, StoreError>
748 {
749 let backing = self.backing.read();
750 let Some(backing) = backing.as_ref() else {
751 return Ok(None);
752 };
753 let mut overlay: FrozenOverlay = backing.removed.iter().map(|id| (*id, None)).collect();
754 let schema = self.column_schema();
755 let key_column = schema
756 .as_ref()
757 .and_then(|schema| schema.iter().position(|name| name == "key"));
758 let mut failure = None;
759 self.for_each_entity_source(|segment, entity| {
760 let key_matches = match (&entity.data, key_column) {
761 (crate::storage::entity::EntityData::Row(row), Some(index)) if row.named.is_none() && row.schema.is_none() => {
762 matches!(row.columns.get(index), Some(reddb_types::Value::Text(stored)) if stored.as_ref() == key)
763 }
764 _ => crate::storage::segment_key_lookup::SegmentKeyLookup::key(entity) == Some(key),
765 };
766 let bytes = if key_matches {
767 let required_bytes = crate::storage::memory_size::entity_bytes(entity);
768 if required_bytes > backing.read_bytes_max {
769 failure = Some(StoreError::ReadLimitExceeded {
770 required_bytes,
771 limit_bytes: backing.read_bytes_max,
772 });
773 return false;
774 }
775 let metadata = segment.get_metadata(entity.id);
776 Some(Arc::from(serialize_entity_record(entity, metadata.as_ref(), backing.format_version)))
777 } else {
778 None
781 };
782 overlay.insert(entity.id, bytes);
783 true
784 });
785 if let Some(error) = failure {
786 return Err(error);
787 }
788 Ok(Some((Arc::clone(&backing.source), Arc::new(overlay))))
789 }
790}