1mod budget;
4mod cancel;
5mod clock;
6mod close;
7pub mod durability;
8mod expansion;
9pub(crate) mod graph_cache;
10mod hybrid;
11pub mod lock;
12mod materialize;
13mod pool;
14pub(crate) mod prepared;
15mod prepared_lexical;
16pub(crate) mod rescored_scan;
17#[cfg(test)]
18mod shared_bound_tests;
19mod snapshot;
20pub(crate) mod stats;
21
22use crate::diag::{timing_elapsed, timing_start};
23pub use cancel::{
24 CancelToken, Deadline, DeadlineError, QueryCancellation, QueryControl, QueryError,
25};
26#[cfg(any(test, feature = "test-support"))]
27pub use clock::ManualMonotonicClock;
28pub use clock::{MonotonicClock, SystemMonotonicClock};
29#[cfg(any(test, feature = "test-support"))]
30pub use hybrid::{
31 HybridScoreTestObservations, begin_hybrid_fresh_round_test_observations,
32 begin_hybrid_score_test_observations, take_hybrid_score_test_observations,
33};
34#[cfg(any(test, feature = "test-support"))]
35pub use materialize::MaterializationTestCounters;
36pub use materialize::{
37 HybridPreparationError, MaterializationError, MaterializedRow, QueryMaterializer,
38};
39
40#[cfg(test)]
41mod hybrid_overlap_tests;
42pub use rescored_scan::ScanRescoreOptions;
43pub use snapshot::{
44 InMemorySegment, InMemorySegmentFactors, PreparedSegment, PublishedSnapshot, SnapshotLease,
45};
46pub use stats::Stats;
47
48use std::collections::HashSet;
49#[cfg(any(test, feature = "test-support"))]
50use std::io::IoSlice;
51use std::path::{Path, PathBuf};
52use std::sync::atomic::{AtomicU64, Ordering};
53use std::sync::{Arc, Condvar, Mutex, RwLock};
54use std::time::Duration;
55
56use durability::{
57 CommitTier, DurabilityMode, DurabilityPolicy, DurabilityPolicyError, SyncRequirement,
58};
59use lock::{StoreLock, StoreLockError};
60
61use self::close::BackgroundThread;
62
63#[cfg(any(test, feature = "test-support"))]
65#[derive(Clone)]
66pub struct StoreTestDependencies {
67 vfs: Arc<dyn crate::vfs::Vfs>,
68 clock: Arc<dyn MonotonicClock>,
69 hybrid_leg_fault: Option<HybridLegTestFault>,
70 storage_fault_controller: Option<StorageFaultController>,
71 ingest_retention_fault_controller: Option<crate::ingest::IngestRetentionFaultController>,
72 metadata_test_controller: Option<Arc<crate::planner::MetadataTestController>>,
73 pub(crate) vector_fault_controller: Option<crate::scan::vector_fault::VectorFaultController>,
74 pub(crate) kernel_fault_controller: Option<crate::kernels::vector_fault::KernelFaultController>,
75 vector_seal_scheme: Option<crate::quant::QuantScheme>,
76}
77
78#[cfg(any(test, feature = "test-support"))]
79impl StoreTestDependencies {
80 #[must_use]
83 pub fn new(vfs: Arc<dyn crate::vfs::Vfs>, clock: Arc<dyn MonotonicClock>) -> Self {
84 Self {
85 vfs,
86 clock,
87 hybrid_leg_fault: None,
88 storage_fault_controller: None,
89 ingest_retention_fault_controller: None,
90 metadata_test_controller: None,
91 vector_fault_controller: None,
92 kernel_fault_controller: None,
93 vector_seal_scheme: None,
94 }
95 }
96
97 #[must_use]
99 pub const fn with_hybrid_leg_fault(mut self, fault: HybridLegTestFault) -> Self {
100 self.hybrid_leg_fault = Some(fault);
101 self
102 }
103
104 #[must_use]
106 pub fn with_storage_fault_controller(mut self, controller: StorageFaultController) -> Self {
107 self.storage_fault_controller = Some(controller);
108 self
109 }
110
111 #[must_use]
113 pub fn with_ingest_retention_fault_controller(
114 mut self,
115 controller: crate::ingest::IngestRetentionFaultController,
116 ) -> Self {
117 self.ingest_retention_fault_controller = Some(controller);
118 self
119 }
120
121 #[must_use]
124 pub fn with_metadata_test_controller(
125 mut self,
126 controller: Arc<crate::planner::MetadataTestController>,
127 ) -> Self {
128 self.metadata_test_controller = Some(controller);
129 self
130 }
131
132 #[must_use]
134 pub fn with_vector_fault_controller(
135 mut self,
136 controller: crate::scan::vector_fault::VectorFaultController,
137 ) -> Self {
138 self.vector_fault_controller = Some(controller);
139 self
140 }
141
142 #[must_use]
144 pub fn with_kernel_fault_controller(
145 mut self,
146 controller: crate::kernels::vector_fault::KernelFaultController,
147 ) -> Self {
148 self.kernel_fault_controller = Some(controller);
149 self
150 }
151
152 #[must_use]
157 pub const fn with_vector_seal_scheme(mut self, scheme: crate::quant::QuantScheme) -> Self {
158 self.vector_seal_scheme = Some(scheme);
159 self
160 }
161}
162
163#[cfg(any(test, feature = "test-support"))]
165#[derive(Clone, Debug, Eq, PartialEq)]
166#[doc(hidden)]
167pub enum StorageTestFault {
168 TornWalHeader,
169 TornWalBody,
170 TornWalChecksum,
171 PostCommitError,
172 ManifestPreRename,
173 ManifestPostRename,
174 CorruptSegmentRegion,
175 WrongManifestObject,
176 WrongSegmentObject,
177 ListOmission { file_name: String },
178 DeleteOmission { file_name: String },
179}
180
181#[cfg(any(test, feature = "test-support"))]
182impl StorageTestFault {
183 const fn operation(&self) -> &'static str {
184 match self {
185 Self::TornWalHeader | Self::TornWalBody | Self::TornWalChecksum => "wal-prefix",
186 Self::PostCommitError => "retry",
187 Self::ManifestPreRename | Self::ManifestPostRename => "publication",
188 Self::CorruptSegmentRegion | Self::WrongManifestObject | Self::WrongSegmentObject => {
189 "format-check"
190 }
191 Self::ListOmission { .. } | Self::DeleteOmission { .. } => "orphan-cleanup",
192 }
193 }
194
195 const fn key(&self) -> &'static str {
196 match self {
197 Self::TornWalHeader => "torn-wal-header",
198 Self::TornWalBody => "torn-wal-body",
199 Self::TornWalChecksum => "torn-wal-checksum",
200 Self::PostCommitError => "post-commit-error",
201 Self::ManifestPreRename => "manifest-pre-rename-crash",
202 Self::ManifestPostRename => "manifest-post-rename-crash",
203 Self::CorruptSegmentRegion => "corrupt-segment-region",
204 Self::WrongManifestObject => "wrong-manifest-object",
205 Self::WrongSegmentObject => "wrong-segment-object",
206 Self::ListOmission { .. } | Self::DeleteOmission { .. } => "list-delete-omission",
207 }
208 }
209}
210
211#[cfg(any(test, feature = "test-support"))]
213#[derive(Clone, Debug, Eq, PartialEq)]
214#[doc(hidden)]
215pub struct StorageFaultPlan {
216 op_index: u32,
217 artifact: String,
218 offset: Option<u64>,
219 segment: Option<crate::segment::SegmentId>,
220 region_kind: Option<u16>,
221 chunk: Option<u32>,
222}
223
224#[cfg(any(test, feature = "test-support"))]
225impl StorageFaultPlan {
226 #[must_use]
228 pub fn new(op_index: u32, artifact: impl Into<String>) -> Self {
229 Self {
230 op_index,
231 artifact: artifact.into(),
232 offset: None,
233 segment: None,
234 region_kind: None,
235 chunk: None,
236 }
237 }
238
239 #[must_use]
241 pub const fn with_offset(mut self, offset: u64) -> Self {
242 self.offset = Some(offset);
243 self
244 }
245
246 #[must_use]
248 pub const fn with_segment_region(
249 mut self,
250 segment: crate::segment::SegmentId,
251 region_kind: u16,
252 chunk: u32,
253 ) -> Self {
254 self.segment = Some(segment);
255 self.region_kind = Some(region_kind);
256 self.chunk = Some(chunk);
257 self
258 }
259
260 #[must_use]
262 pub const fn with_segment(mut self, segment: crate::segment::SegmentId) -> Self {
263 self.segment = Some(segment);
264 self
265 }
266
267 #[must_use]
269 pub const fn op_index(&self) -> u32 {
270 self.op_index
271 }
272
273 #[must_use]
275 pub fn artifact(&self) -> &str {
276 &self.artifact
277 }
278
279 #[must_use]
281 pub const fn offset(&self) -> Option<u64> {
282 self.offset
283 }
284
285 #[must_use]
287 pub const fn segment(&self) -> Option<crate::segment::SegmentId> {
288 self.segment
289 }
290
291 #[must_use]
293 pub const fn region_kind(&self) -> Option<u16> {
294 self.region_kind
295 }
296
297 #[must_use]
299 pub const fn chunk(&self) -> Option<u32> {
300 self.chunk
301 }
302}
303
304#[cfg(any(test, feature = "test-support"))]
306#[derive(Clone, Copy, Debug, Eq, PartialEq)]
307#[doc(hidden)]
308pub enum StorageReceiptSite {
309 WalOpenHeaderValidation,
310 WalOpenRecordValidation,
311 WalOpenRecordChecksum,
312 WalCommitAppendAfterInnerSuccess,
313 ManifestCommitBeforeRename,
314 ManifestCommitAfterRename,
315 SegmentReadRegionChecksum,
316 ManifestOpenFamilyValidation,
317 SegmentOpenFamilyValidation,
318 SegmentOpenObjectIdentity,
319 OrphanCleanupList,
320 OrphanCleanupDelete,
321}
322
323#[cfg(any(test, feature = "test-support"))]
324impl StorageReceiptSite {
325 const fn key(self) -> &'static str {
326 match self {
327 Self::WalOpenHeaderValidation => "WalOpen.HeaderValidation",
328 Self::WalOpenRecordValidation => "WalOpen.RecordValidation",
329 Self::WalOpenRecordChecksum => "WalOpen.RecordChecksum",
330 Self::WalCommitAppendAfterInnerSuccess => "WalCommit.AppendAfterInnerSuccess",
331 Self::ManifestCommitBeforeRename => "ManifestCommit.BeforeRename",
332 Self::ManifestCommitAfterRename => "ManifestCommit.AfterRename",
333 Self::SegmentReadRegionChecksum => "SegmentRead.RegionChecksum",
334 Self::ManifestOpenFamilyValidation => "ManifestOpen.FamilyValidation",
335 Self::SegmentOpenFamilyValidation => "SegmentOpen.FamilyValidation",
336 Self::SegmentOpenObjectIdentity => "SegmentOpen.ObjectIdentity",
337 Self::OrphanCleanupList => "OrphanCleanup.List",
338 Self::OrphanCleanupDelete => "OrphanCleanup.Delete",
339 }
340 }
341}
342
343#[cfg(any(test, feature = "test-support"))]
345#[derive(Clone, Debug, Eq, PartialEq)]
346#[doc(hidden)]
347pub enum StorageReceiptObserved {
348 WalHeader {
349 artifact: String,
350 reason: crate::wal::header::WalHeaderError,
351 },
352 WalRecord {
353 artifact: String,
354 offset: u64,
355 reason: crate::wal::replay::CorruptionReason,
356 },
357 WalAppend {
358 artifact: String,
359 encoded_len: u64,
360 first_seq: u64,
361 last_seq: u64,
362 inner_append_completed: bool,
363 caller_saw_error: bool,
364 },
365 ManifestRename {
366 temporary: String,
367 committed: String,
368 rename_performed: bool,
369 new_segment_final: bool,
370 directory_sync_returned: bool,
371 },
372 SegmentChecksum {
373 artifact: String,
374 segment: crate::segment::SegmentId,
375 region_kind: u16,
376 chunk: u32,
377 expected_checksum: u64,
378 actual_checksum: u64,
379 },
380 Format {
381 artifact: String,
382 check: crate::format::frame::FormatCheck,
383 expected_family: Option<u16>,
384 actual_family: Option<u16>,
385 expected_id: Option<crate::segment::SegmentId>,
386 actual_id: Option<crate::segment::SegmentId>,
387 },
388 Omission {
389 artifact: String,
390 deletion_observed: bool,
391 },
392}
393
394#[cfg(any(test, feature = "test-support"))]
396#[derive(Clone, Debug, Eq, PartialEq)]
397#[doc(hidden)]
398pub struct StorageFaultReceipt {
399 operation: &'static str,
400 fault: &'static str,
401 site: StorageReceiptSite,
402 plan: StorageFaultPlan,
403 observed: StorageReceiptObserved,
404 cardinality: u32,
405}
406
407#[cfg(any(test, feature = "test-support"))]
408impl StorageFaultReceipt {
409 #[must_use]
410 pub const fn campaign(&self) -> &'static str {
411 "storage-durability"
412 }
413
414 #[must_use]
415 pub const fn operation(&self) -> &'static str {
416 self.operation
417 }
418
419 #[must_use]
420 pub const fn fault(&self) -> &'static str {
421 self.fault
422 }
423
424 #[must_use]
425 pub const fn site(&self) -> &'static str {
426 self.site.key()
427 }
428
429 #[must_use]
430 pub const fn typed_site(&self) -> StorageReceiptSite {
431 self.site
432 }
433
434 #[must_use]
435 pub const fn cardinality(&self) -> u32 {
436 self.cardinality
437 }
438
439 #[must_use]
440 pub const fn plan(&self) -> &StorageFaultPlan {
441 &self.plan
442 }
443
444 #[must_use]
445 pub const fn observed(&self) -> &StorageReceiptObserved {
446 &self.observed
447 }
448
449 pub fn from_manifest_abort_acknowledgment(acknowledgment: &str) -> Result<Self, String> {
455 let body = acknowledgment
456 .strip_suffix('\n')
457 .filter(|body| !body.contains('\n'))
458 .ok_or_else(|| {
459 "manifest abort acknowledgment must end in exactly one newline".to_owned()
460 })?;
461 let fields = body
462 .split('|')
463 .map(|field| {
464 field.split_once('=').ok_or_else(|| {
465 format!("manifest abort acknowledgment field lacks '=': {field}")
466 })
467 })
468 .collect::<Result<Vec<_>, _>>()?;
469 let expected_keys = [
470 "campaign",
471 "operation",
472 "fault",
473 "site",
474 "op_index",
475 "cardinality",
476 "artifact",
477 "temporary",
478 "committed",
479 "rename_performed",
480 "new_segment_final",
481 "directory_sync_returned",
482 ];
483 if fields.len() != expected_keys.len()
484 || fields
485 .iter()
486 .map(|(key, _)| *key)
487 .ne(expected_keys.iter().copied())
488 {
489 return Err("manifest abort acknowledgment field inventory differs".to_owned());
490 }
491 let mut values = fields.into_iter().map(|(_, value)| value);
492 let campaign = values
493 .next()
494 .ok_or_else(|| "manifest abort acknowledgment lacks campaign".to_owned())?;
495 let operation = values
496 .next()
497 .ok_or_else(|| "manifest abort acknowledgment lacks operation".to_owned())?;
498 let fault = values
499 .next()
500 .ok_or_else(|| "manifest abort acknowledgment lacks fault".to_owned())?;
501 let site = values
502 .next()
503 .ok_or_else(|| "manifest abort acknowledgment lacks site".to_owned())?;
504 let op_index = values
505 .next()
506 .ok_or_else(|| "manifest abort acknowledgment lacks op index".to_owned())?
507 .parse::<u32>()
508 .map_err(|error| format!("manifest abort op index is invalid: {error}"))?;
509 let cardinality = values
510 .next()
511 .ok_or_else(|| "manifest abort acknowledgment lacks cardinality".to_owned())?
512 .parse::<u32>()
513 .map_err(|error| format!("manifest abort cardinality is invalid: {error}"))?;
514 let artifact = values
515 .next()
516 .ok_or_else(|| "manifest abort acknowledgment lacks artifact".to_owned())?;
517 let temporary = values
518 .next()
519 .ok_or_else(|| "manifest abort acknowledgment lacks temporary path".to_owned())?;
520 let committed = values
521 .next()
522 .ok_or_else(|| "manifest abort acknowledgment lacks committed path".to_owned())?;
523 let parse_bool = |field: &str, value: &str| match value {
524 "true" => Ok(true),
525 "false" => Ok(false),
526 _ => Err(format!("manifest abort {field} is not a bool: {value}")),
527 };
528 let rename_performed = parse_bool(
529 "rename_performed",
530 values
531 .next()
532 .ok_or_else(|| "manifest abort acknowledgment lacks rename result".to_owned())?,
533 )?;
534 let new_segment_final = parse_bool(
535 "new_segment_final",
536 values.next().ok_or_else(|| {
537 "manifest abort acknowledgment lacks final-segment result".to_owned()
538 })?,
539 )?;
540 let directory_sync_returned = parse_bool(
541 "directory_sync_returned",
542 values.next().ok_or_else(|| {
543 "manifest abort acknowledgment lacks directory-sync result".to_owned()
544 })?,
545 )?;
546 if campaign != "storage-durability"
547 || operation != "publication"
548 || cardinality != 1
549 || artifact != crate::manifest::io::MANIFEST_FILE
550 || temporary != crate::manifest::io::MANIFEST_TEMP_FILE
551 || committed != crate::manifest::io::MANIFEST_FILE
552 || !new_segment_final
553 || directory_sync_returned
554 {
555 return Err("manifest abort acknowledgment fixed facts differ".to_owned());
556 }
557 let (fault, site, expected_rename) = match (fault, site) {
558 ("manifest-pre-rename-crash", "ManifestCommit.BeforeRename") => (
559 "manifest-pre-rename-crash",
560 StorageReceiptSite::ManifestCommitBeforeRename,
561 false,
562 ),
563 ("manifest-post-rename-crash", "ManifestCommit.AfterRename") => (
564 "manifest-post-rename-crash",
565 StorageReceiptSite::ManifestCommitAfterRename,
566 true,
567 ),
568 _ => return Err("manifest abort fault and production site differ".to_owned()),
569 };
570 if rename_performed != expected_rename {
571 return Err("manifest abort rename result differs from production site".to_owned());
572 }
573 Ok(Self {
574 operation: "publication",
575 fault,
576 site,
577 plan: StorageFaultPlan::new(op_index, artifact),
578 observed: StorageReceiptObserved::ManifestRename {
579 temporary: temporary.to_owned(),
580 committed: committed.to_owned(),
581 rename_performed,
582 new_segment_final,
583 directory_sync_returned,
584 },
585 cardinality,
586 })
587 }
588}
589
590#[cfg(any(test, feature = "test-support"))]
591#[derive(Debug)]
592struct StorageFaultState {
593 fault: StorageTestFault,
594 plan: StorageFaultPlan,
595 consumed: bool,
596 receipt: Option<StorageFaultReceipt>,
597 cleanup_report: Option<StorageCleanupReport>,
598 #[cfg(unix)]
599 child_abort_ack: Option<std::os::unix::net::UnixStream>,
600}
601
602#[cfg(any(test, feature = "test-support"))]
604#[derive(Clone, Debug, Eq, PartialEq)]
605#[doc(hidden)]
606pub struct StorageCleanupReport {
607 reclaimed_bytes: u64,
608 deleted_paths: Vec<String>,
609 retained_eligible_paths: Vec<String>,
610 directory_synced: bool,
611}
612
613#[cfg(any(test, feature = "test-support"))]
614impl StorageCleanupReport {
615 #[must_use]
617 pub const fn reclaimed_bytes(&self) -> u64 {
618 self.reclaimed_bytes
619 }
620
621 #[must_use]
623 pub fn deleted_paths(&self) -> &[String] {
624 &self.deleted_paths
625 }
626
627 #[must_use]
629 pub fn retained_eligible_paths(&self) -> &[String] {
630 &self.retained_eligible_paths
631 }
632
633 #[must_use]
635 pub const fn directory_synced(&self) -> bool {
636 self.directory_synced
637 }
638}
639
640#[cfg(any(test, feature = "test-support"))]
642#[derive(Clone, Debug)]
643#[doc(hidden)]
644pub struct StorageFaultController {
645 state: Arc<Mutex<StorageFaultState>>,
646}
647
648#[cfg(any(test, feature = "test-support"))]
649impl StorageFaultController {
650 #[must_use]
652 pub fn new(fault: StorageTestFault, plan: StorageFaultPlan) -> Self {
653 Self {
654 state: Arc::new(Mutex::new(StorageFaultState {
655 fault,
656 plan,
657 consumed: false,
658 receipt: None,
659 cleanup_report: None,
660 #[cfg(unix)]
661 child_abort_ack: None,
662 })),
663 }
664 }
665
666 #[cfg(unix)]
668 #[must_use]
669 pub fn with_child_abort_ack(self, acknowledgment: std::os::unix::net::UnixStream) -> Self {
670 if let Ok(mut state) = self.state.lock() {
671 state.child_abort_ack = Some(acknowledgment);
672 }
673 self
674 }
675
676 pub fn take_receipt(&self) -> Option<StorageFaultReceipt> {
678 self.state
679 .lock()
680 .ok()
681 .and_then(|mut state| state.receipt.take())
682 }
683
684 pub fn take_cleanup_report(&self) -> Option<StorageCleanupReport> {
686 self.state
687 .lock()
688 .ok()
689 .and_then(|mut state| state.cleanup_report.take())
690 }
691
692 fn record_cleanup_report(&self, report: crate::manifest::io::OrphanCleanupReport) {
693 let normalize = |paths: Vec<PathBuf>| {
694 paths
695 .into_iter()
696 .filter_map(|path| {
697 path.file_name()
698 .and_then(|name| name.to_str())
699 .map(str::to_owned)
700 })
701 .collect::<Vec<_>>()
702 };
703 if let Ok(mut state) = self.state.lock() {
704 state.cleanup_report = Some(StorageCleanupReport {
705 reclaimed_bytes: report.reclaimed_bytes,
706 deleted_paths: normalize(report.deleted_paths),
707 retained_eligible_paths: normalize(report.retained_eligible_paths),
708 directory_synced: report.directory_synced,
709 });
710 }
711 }
712
713 fn armed_fault(&self) -> Option<StorageTestFault> {
714 self.state
715 .lock()
716 .ok()
717 .and_then(|state| (!state.consumed).then(|| state.fault.clone()))
718 }
719
720 fn planned_segment(&self) -> Option<crate::segment::SegmentId> {
721 self.state
722 .lock()
723 .ok()
724 .and_then(|state| state.plan.segment())
725 }
726
727 fn emit(&self, site: StorageReceiptSite, observed: StorageReceiptObserved) -> bool {
728 let Ok(mut state) = self.state.lock() else {
729 return false;
730 };
731 if state.consumed {
732 return false;
733 }
734 state.consumed = true;
735 state.receipt = Some(StorageFaultReceipt {
736 operation: state.fault.operation(),
737 fault: state.fault.key(),
738 site,
739 plan: state.plan.clone(),
740 observed,
741 cardinality: 1,
742 });
743 true
744 }
745
746 fn abort_after_manifest_receipt(&self) -> ! {
747 #[cfg(unix)]
748 if let Ok(mut state) = self.state.lock() {
749 let wire_receipt = state
750 .receipt
751 .as_ref()
752 .and_then(storage_manifest_abort_wire_receipt);
753 if let (Some(wire_receipt), Some(acknowledgment)) =
754 (wire_receipt, state.child_abort_ack.as_mut())
755 {
756 use std::io::Write as _;
757 let _ = acknowledgment.write_all(wire_receipt.as_bytes());
758 let _ = acknowledgment.flush();
759 }
760 }
761 std::process::abort()
762 }
763}
764
765#[cfg(any(test, feature = "test-support"))]
766thread_local! {
767 static STORAGE_OPEN_CONTROLLER: std::cell::RefCell<
768 Option<std::sync::Weak<Mutex<StorageFaultState>>>
769 > = const { std::cell::RefCell::new(None) };
770}
771
772#[cfg(any(test, feature = "test-support"))]
773fn with_storage_open_controller<R>(
774 controller: Option<&StorageFaultController>,
775 operation: impl FnOnce() -> R,
776) -> R {
777 STORAGE_OPEN_CONTROLLER.with(|slot| {
778 let previous = slot.replace(controller.map(|controller| Arc::downgrade(&controller.state)));
779 let result = operation();
780 let _ = slot.replace(previous);
781 result
782 })
783}
784
785#[cfg(any(test, feature = "test-support"))]
786fn storage_open_controller() -> Option<StorageFaultController> {
787 STORAGE_OPEN_CONTROLLER.with(|slot| {
788 slot.borrow()
789 .as_ref()
790 .and_then(std::sync::Weak::upgrade)
791 .map(|state| StorageFaultController { state })
792 })
793}
794
795#[cfg(any(test, feature = "test-support"))]
796pub(crate) fn record_storage_wal_recovery_fault(error: &crate::wal::WalRecoveryError) {
797 let Some(controller) = storage_open_controller() else {
798 return;
799 };
800 match (controller.armed_fault(), error) {
801 (
802 Some(StorageTestFault::TornWalHeader),
803 crate::wal::WalRecoveryError::InvalidHeader(reason),
804 ) => {
805 let _ = controller.emit(
806 StorageReceiptSite::WalOpenHeaderValidation,
807 StorageReceiptObserved::WalHeader {
808 artifact: "wal.ze".to_owned(),
809 reason: *reason,
810 },
811 );
812 }
813 (
814 Some(StorageTestFault::TornWalBody),
815 crate::wal::WalRecoveryError::CorruptAt {
816 offset,
817 reason:
818 reason @ crate::wal::replay::CorruptionReason::Record {
819 error:
820 crate::wal::record::RecordError::HeaderTruncated { .. }
821 | crate::wal::record::RecordError::BodyTruncated { .. }
822 | crate::wal::record::RecordError::LengthOverflow { .. },
823 ..
824 },
825 },
826 ) => record_storage_wal_record_fault(
827 &controller,
828 StorageReceiptSite::WalOpenRecordValidation,
829 *offset,
830 *reason,
831 ),
832 (
833 Some(StorageTestFault::TornWalChecksum),
834 crate::wal::WalRecoveryError::CorruptAt {
835 offset,
836 reason:
837 reason @ crate::wal::replay::CorruptionReason::Record {
838 error: crate::wal::record::RecordError::ChecksumMismatch { .. },
839 ..
840 },
841 },
842 ) => record_storage_wal_record_fault(
843 &controller,
844 StorageReceiptSite::WalOpenRecordChecksum,
845 *offset,
846 *reason,
847 ),
848 _ => {}
849 }
850}
851
852#[cfg(any(test, feature = "test-support"))]
853fn record_storage_wal_record_fault(
854 controller: &StorageFaultController,
855 site: StorageReceiptSite,
856 offset: usize,
857 reason: crate::wal::replay::CorruptionReason,
858) {
859 let Ok(offset) = u64::try_from(offset) else {
860 return;
861 };
862 let _ = controller.emit(
863 site,
864 StorageReceiptObserved::WalRecord {
865 artifact: "wal.ze".to_owned(),
866 offset,
867 reason,
868 },
869 );
870}
871
872#[cfg(any(test, feature = "test-support"))]
873pub(crate) fn record_storage_manifest_format_fault(
874 error: &crate::format::frame::FormatError,
875 actual_family: Option<u16>,
876) {
877 let Some(controller) = storage_open_controller() else {
878 return;
879 };
880 if matches!(
881 controller.armed_fault(),
882 Some(StorageTestFault::WrongManifestObject)
883 ) && error.check() == crate::format::frame::FormatCheck::Family
884 {
885 let _ = controller.emit(
886 StorageReceiptSite::ManifestOpenFamilyValidation,
887 StorageReceiptObserved::Format {
888 artifact: storage_artifact_name(error.artifact()).to_owned(),
889 check: error.check(),
890 expected_family: Some(crate::format::FormatFamily::Manifest.id()),
891 actual_family,
892 expected_id: None,
893 actual_id: None,
894 },
895 );
896 }
897}
898
899#[cfg(any(test, feature = "test-support"))]
900pub(crate) fn record_storage_segment_format_fault(
901 error: &crate::format::frame::FormatError,
902 actual_family: Option<u16>,
903 expected_id: crate::segment::SegmentId,
904) {
905 let Some(controller) = storage_open_controller() else {
906 return;
907 };
908 if matches!(
909 controller.armed_fault(),
910 Some(StorageTestFault::WrongSegmentObject)
911 ) && error.check() == crate::format::frame::FormatCheck::Family
912 {
913 let _ = controller.emit(
914 StorageReceiptSite::SegmentOpenFamilyValidation,
915 StorageReceiptObserved::Format {
916 artifact: storage_artifact_name(error.artifact()).to_owned(),
917 check: error.check(),
918 expected_family: Some(crate::format::FormatFamily::Segment.id()),
919 actual_family,
920 expected_id: Some(expected_id),
921 actual_id: None,
922 },
923 );
924 }
925}
926
927#[cfg(any(test, feature = "test-support"))]
928pub(crate) fn record_storage_segment_identity_fault(
929 artifact: &str,
930 expected: crate::segment::SegmentId,
931 actual: crate::segment::SegmentId,
932) {
933 let Some(controller) = storage_open_controller() else {
934 return;
935 };
936 if matches!(
937 controller.armed_fault(),
938 Some(StorageTestFault::WrongSegmentObject)
939 ) {
940 let _ = controller.emit(
941 StorageReceiptSite::SegmentOpenObjectIdentity,
942 StorageReceiptObserved::Format {
943 artifact: storage_artifact_name(artifact).to_owned(),
944 check: crate::format::frame::FormatCheck::ObjectIdentity,
945 expected_family: None,
946 actual_family: None,
947 expected_id: Some(expected),
948 actual_id: Some(actual),
949 },
950 );
951 }
952}
953
954#[cfg(any(test, feature = "test-support"))]
955fn storage_manifest_abort_wire_receipt(receipt: &StorageFaultReceipt) -> Option<String> {
956 let StorageReceiptObserved::ManifestRename {
957 temporary,
958 committed,
959 rename_performed,
960 new_segment_final,
961 directory_sync_returned,
962 } = &receipt.observed
963 else {
964 return None;
965 };
966 Some(format!(
967 "campaign={}|operation={}|fault={}|site={}|op_index={}|cardinality={}|artifact={}|temporary={temporary}|committed={committed}|rename_performed={rename_performed}|new_segment_final={new_segment_final}|directory_sync_returned={directory_sync_returned}\n",
968 receipt.campaign(),
969 receipt.operation,
970 receipt.fault,
971 receipt.site.key(),
972 receipt.plan.op_index,
973 receipt.cardinality,
974 receipt.plan.artifact,
975 ))
976}
977
978#[cfg(any(test, feature = "test-support"))]
979fn storage_artifact_name(artifact: &str) -> &str {
980 Path::new(artifact)
981 .file_name()
982 .and_then(|name| name.to_str())
983 .unwrap_or(artifact)
984}
985
986#[cfg(any(test, feature = "test-support"))]
987type StorageSegmentControllerRegistry = Mutex<
988 Vec<(
989 PathBuf,
990 crate::segment::SegmentId,
991 std::sync::Weak<Mutex<StorageFaultState>>,
992 )>,
993>;
994
995#[cfg(any(test, feature = "test-support"))]
996static STORAGE_SEGMENT_CONTROLLERS: std::sync::OnceLock<StorageSegmentControllerRegistry> =
997 std::sync::OnceLock::new();
998
999#[cfg(any(test, feature = "test-support"))]
1000fn register_storage_segment_controller(
1001 store_directory: &Path,
1002 segment: crate::segment::SegmentId,
1003 controller: &StorageFaultController,
1004) {
1005 let registry = STORAGE_SEGMENT_CONTROLLERS.get_or_init(|| Mutex::new(Vec::new()));
1006 if let Ok(mut entries) = registry.lock() {
1007 entries.retain(|(_, _, state)| state.strong_count() != 0);
1008 entries.push((
1009 store_directory.to_path_buf(),
1010 segment,
1011 Arc::downgrade(&controller.state),
1012 ));
1013 }
1014}
1015
1016#[cfg(any(test, feature = "test-support"))]
1017pub(crate) fn record_storage_segment_checksum_fault(
1018 store_directory: &Path,
1019 segment: crate::segment::SegmentId,
1020 kind: crate::segment::layout::RegionKind,
1021 chunk: u32,
1022 expected: u64,
1023 actual: u64,
1024) {
1025 let Some(registry) = STORAGE_SEGMENT_CONTROLLERS.get() else {
1026 return;
1027 };
1028 let states = match registry.lock() {
1029 Ok(mut entries) => {
1030 entries.retain(|(_, _, state)| state.strong_count() != 0);
1031 entries
1032 .iter()
1033 .filter(|(directory, registered, _)| {
1034 directory == store_directory && *registered == segment
1035 })
1036 .filter_map(|(_, _, state)| state.upgrade())
1037 .collect::<Vec<_>>()
1038 }
1039 Err(_) => return,
1040 };
1041 for state in states {
1042 let controller = StorageFaultController { state };
1043 if matches!(
1044 controller.armed_fault(),
1045 Some(StorageTestFault::CorruptSegmentRegion)
1046 ) {
1047 let _ = controller.emit(
1048 StorageReceiptSite::SegmentReadRegionChecksum,
1049 StorageReceiptObserved::SegmentChecksum {
1050 artifact: segment.file_name(),
1051 segment,
1052 region_kind: kind.id(),
1053 chunk,
1054 expected_checksum: expected,
1055 actual_checksum: actual,
1056 },
1057 );
1058 }
1059 }
1060}
1061
1062#[cfg(any(test, feature = "test-support"))]
1063struct StorageFaultVfs {
1064 inner: Arc<dyn crate::vfs::Vfs>,
1065 controller: StorageFaultController,
1066}
1067
1068#[cfg(any(test, feature = "test-support"))]
1069struct StorageFaultFile {
1070 inner: Box<dyn crate::vfs::VfsFile>,
1071 path: PathBuf,
1072 controller: StorageFaultController,
1073}
1074
1075#[cfg(any(test, feature = "test-support"))]
1076impl StorageFaultFile {
1077 fn after_append(&self, encoded_len: u64, first_seq: u64, last_seq: u64) -> std::io::Result<()> {
1078 if matches!(
1079 self.controller.armed_fault(),
1080 Some(StorageTestFault::PostCommitError)
1081 ) && self.path.file_name().is_some_and(|name| name == "wal.ze")
1082 && self.controller.emit(
1083 StorageReceiptSite::WalCommitAppendAfterInnerSuccess,
1084 StorageReceiptObserved::WalAppend {
1085 artifact: "wal.ze".to_owned(),
1086 encoded_len,
1087 first_seq,
1088 last_seq,
1089 inner_append_completed: true,
1090 caller_saw_error: true,
1091 },
1092 )
1093 {
1094 return Err(std::io::Error::other(
1095 "scheduled post-commit error after inner append",
1096 ));
1097 }
1098 Ok(())
1099 }
1100}
1101
1102#[cfg(any(test, feature = "test-support"))]
1103impl crate::vfs::VfsFile for StorageFaultFile {
1104 fn append(&mut self, bytes: &[u8]) -> std::io::Result<()> {
1105 let record_facts = storage_wal_append_facts(bytes)?;
1106 self.inner.append(bytes)?;
1107 record_facts.map_or(Ok(()), |(encoded_len, first_seq, last_seq)| {
1108 self.after_append(encoded_len, first_seq, last_seq)
1109 })
1110 }
1111
1112 fn append_vectored(&mut self, buffers: &mut [IoSlice<'_>]) -> std::io::Result<()> {
1113 let mut aggregate = None::<(u64, u64, u64)>;
1114 for bytes in buffers.iter() {
1115 if let Some((encoded_len, first_seq, last_seq)) = storage_wal_append_facts(bytes)? {
1116 aggregate = Some(match aggregate {
1117 None => (encoded_len, first_seq, last_seq),
1118 Some((total, first, _)) => (
1119 total.checked_add(encoded_len).ok_or_else(|| {
1120 std::io::Error::other("WAL record append length overflow")
1121 })?,
1122 first,
1123 last_seq,
1124 ),
1125 });
1126 }
1127 }
1128 self.inner.append_vectored(buffers)?;
1129 aggregate.map_or(Ok(()), |(encoded_len, first_seq, last_seq)| {
1130 self.after_append(encoded_len, first_seq, last_seq)
1131 })
1132 }
1133
1134 fn sync(&self, kind: crate::vfs::SyncKind) -> std::io::Result<()> {
1135 self.inner.sync(kind)
1136 }
1137}
1138
1139#[cfg(any(test, feature = "test-support"))]
1140impl crate::vfs::Vfs for StorageFaultVfs {
1141 fn segment_data_read_counter(&self) -> Option<Arc<AtomicU64>> {
1142 self.inner.segment_data_read_counter()
1143 }
1144
1145 fn ensure_directory(&self, path: &Path, create: bool) -> std::io::Result<bool> {
1146 self.inner.ensure_directory(path, create)
1147 }
1148
1149 fn open(&self, path: &Path) -> std::io::Result<u64> {
1150 self.inner.open(path)
1151 }
1152
1153 fn open_for_map(&self, path: &Path) -> std::io::Result<std::fs::File> {
1154 self.inner.open_for_map(path)
1155 }
1156
1157 fn read(&self, path: &Path) -> std::io::Result<Vec<u8>> {
1158 self.inner.read(path)
1159 }
1160
1161 fn read_range(&self, path: &Path, offset: u64, length: usize) -> std::io::Result<Vec<u8>> {
1162 self.inner.read_range(path, offset, length)
1163 }
1164
1165 fn write(&self, path: &Path, bytes: &[u8]) -> std::io::Result<()> {
1166 self.inner.write(path, bytes)
1167 }
1168
1169 fn open_append(&self, path: &Path) -> std::io::Result<Box<dyn crate::vfs::VfsFile>> {
1170 Ok(Box::new(StorageFaultFile {
1171 inner: self.inner.open_append(path)?,
1172 path: path.to_path_buf(),
1173 controller: self.controller.clone(),
1174 }))
1175 }
1176
1177 fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
1178 let is_manifest = from
1179 .file_name()
1180 .is_some_and(|name| name == crate::manifest::io::MANIFEST_TEMP_FILE)
1181 && to
1182 .file_name()
1183 .is_some_and(|name| name == crate::manifest::io::MANIFEST_FILE);
1184 if is_manifest
1185 && matches!(
1186 self.controller.armed_fault(),
1187 Some(StorageTestFault::ManifestPreRename)
1188 )
1189 && self.controller.emit(
1190 StorageReceiptSite::ManifestCommitBeforeRename,
1191 StorageReceiptObserved::ManifestRename {
1192 temporary: crate::manifest::io::MANIFEST_TEMP_FILE.to_owned(),
1193 committed: crate::manifest::io::MANIFEST_FILE.to_owned(),
1194 rename_performed: false,
1195 new_segment_final: storage_planned_segment_is_final(
1196 &*self.inner,
1197 from,
1198 self.controller.planned_segment(),
1199 ),
1200 directory_sync_returned: false,
1201 },
1202 )
1203 {
1204 self.controller.abort_after_manifest_receipt();
1205 }
1206 self.inner.rename(from, to)?;
1207 if is_manifest
1208 && matches!(
1209 self.controller.armed_fault(),
1210 Some(StorageTestFault::ManifestPostRename)
1211 )
1212 && self.controller.emit(
1213 StorageReceiptSite::ManifestCommitAfterRename,
1214 StorageReceiptObserved::ManifestRename {
1215 temporary: crate::manifest::io::MANIFEST_TEMP_FILE.to_owned(),
1216 committed: crate::manifest::io::MANIFEST_FILE.to_owned(),
1217 rename_performed: true,
1218 new_segment_final: storage_planned_segment_is_final(
1219 &*self.inner,
1220 to,
1221 self.controller.planned_segment(),
1222 ),
1223 directory_sync_returned: false,
1224 },
1225 )
1226 {
1227 self.controller.abort_after_manifest_receipt();
1228 }
1229 Ok(())
1230 }
1231
1232 fn sync(&self, path: &Path, kind: crate::vfs::SyncKind) -> std::io::Result<()> {
1233 self.inner.sync(path, kind)
1234 }
1235
1236 fn list(&self, directory: &Path) -> std::io::Result<Vec<PathBuf>> {
1237 let mut paths = self.inner.list(directory)?;
1238 if let Some(StorageTestFault::ListOmission { file_name }) = self.controller.armed_fault()
1239 && let Some(position) = paths.iter().position(|path| {
1240 path.file_name()
1241 .is_some_and(|name| name == file_name.as_str())
1242 })
1243 {
1244 let _ = paths.remove(position);
1245 let _ = self.controller.emit(
1246 StorageReceiptSite::OrphanCleanupList,
1247 StorageReceiptObserved::Omission {
1248 artifact: file_name,
1249 deletion_observed: false,
1250 },
1251 );
1252 }
1253 Ok(paths)
1254 }
1255
1256 fn delete(&self, path: &Path) -> std::io::Result<()> {
1257 if let Some(StorageTestFault::DeleteOmission { file_name }) = self.controller.armed_fault()
1258 && path
1259 .file_name()
1260 .is_some_and(|name| name == file_name.as_str())
1261 && self.controller.emit(
1262 StorageReceiptSite::OrphanCleanupDelete,
1263 StorageReceiptObserved::Omission {
1264 artifact: file_name,
1265 deletion_observed: false,
1266 },
1267 )
1268 {
1269 return Ok(());
1270 }
1271 self.inner.delete(path)
1272 }
1273}
1274
1275#[cfg(any(test, feature = "test-support"))]
1276fn storage_wal_append_facts(bytes: &[u8]) -> std::io::Result<Option<(u64, u64, u64)>> {
1277 const RECORD_HEADER_LEN: usize = 14;
1278 const RECORD_CHECKSUM_LEN: usize = 8;
1279 let has_wal_header = bytes
1280 .get(..8)
1281 .is_some_and(|magic| magic == crate::format::frame::FILE_MAGIC)
1282 && bytes
1283 .get(8..10)
1284 .is_some_and(|family| family == crate::format::FormatFamily::Wal.id().to_le_bytes());
1285 let mut offset = if has_wal_header {
1286 crate::wal::header::WAL_HEADER_LEN
1287 } else {
1288 0
1289 };
1290 if offset == bytes.len() {
1291 return Ok(None);
1292 }
1293 let mut encoded_len = 0_u64;
1294 let mut first_seq = None;
1295 let mut last_seq = None;
1296 while offset < bytes.len() {
1297 let header_end = offset
1298 .checked_add(RECORD_HEADER_LEN)
1299 .ok_or_else(|| std::io::Error::other("WAL record header offset overflow"))?;
1300 let header = bytes
1301 .get(offset..header_end)
1302 .ok_or_else(|| std::io::Error::other("WAL append has a torn record header"))?;
1303 let payload_len = usize::try_from(u32::from_le_bytes(
1304 header
1305 .get(..4)
1306 .ok_or_else(|| std::io::Error::other("WAL payload length is absent"))?
1307 .try_into()
1308 .map_err(|_| std::io::Error::other("WAL payload length width changed"))?,
1309 ))
1310 .map_err(|_| std::io::Error::other("WAL payload length does not fit usize"))?;
1311 let seq = u64::from_le_bytes(
1312 header
1313 .get(4..12)
1314 .ok_or_else(|| std::io::Error::other("WAL sequence is absent"))?
1315 .try_into()
1316 .map_err(|_| std::io::Error::other("WAL sequence width changed"))?,
1317 );
1318 let record_len = RECORD_HEADER_LEN
1319 .checked_add(payload_len)
1320 .and_then(|length| length.checked_add(RECORD_CHECKSUM_LEN))
1321 .ok_or_else(|| std::io::Error::other("WAL record length overflow"))?;
1322 offset = offset
1323 .checked_add(record_len)
1324 .ok_or_else(|| std::io::Error::other("WAL record end overflow"))?;
1325 if offset > bytes.len() {
1326 return Err(std::io::Error::other("WAL append has a torn record body"));
1327 }
1328 encoded_len = encoded_len
1329 .checked_add(
1330 u64::try_from(record_len)
1331 .map_err(|_| std::io::Error::other("WAL record length does not fit u64"))?,
1332 )
1333 .ok_or_else(|| std::io::Error::other("WAL record append length overflow"))?;
1334 first_seq.get_or_insert(seq);
1335 last_seq = Some(seq);
1336 }
1337 match (first_seq, last_seq) {
1338 (Some(first), Some(last)) => Ok(Some((encoded_len, first, last))),
1339 _ => Ok(None),
1340 }
1341}
1342
1343#[cfg(any(test, feature = "test-support"))]
1344fn storage_planned_segment_is_final(
1345 vfs: &dyn crate::vfs::Vfs,
1346 artifact: &Path,
1347 planned: Option<crate::segment::SegmentId>,
1348) -> bool {
1349 let Some(file_name) = planned.map(crate::segment::SegmentId::file_name) else {
1350 return false;
1351 };
1352 artifact.parent().is_some_and(|directory| {
1353 vfs.list(directory).is_ok_and(|paths| {
1354 paths.iter().any(|path| {
1355 path.file_name()
1356 .and_then(|name| name.to_str())
1357 .is_some_and(|name| name == file_name)
1358 })
1359 })
1360 })
1361}
1362
1363#[cfg(any(test, feature = "test-support"))]
1365#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1366#[doc(hidden)]
1367pub enum HybridLegTestFault {
1368 Panic(crate::fusion::FusionLeg),
1370}
1371
1372#[allow(clippy::panic)]
1373fn maybe_trigger_hybrid_leg_panic(armed: bool, detail: &'static str) {
1374 #[cfg(any(test, feature = "test-support"))]
1375 if armed {
1376 std::panic::panic_any(detail);
1377 }
1378 #[cfg(not(any(test, feature = "test-support")))]
1379 let _ = (armed, detail);
1380}
1381
1382#[cfg(any(test, feature = "test-support"))]
1384#[derive(Clone, Debug, Eq, PartialEq)]
1385#[doc(hidden)]
1386pub struct HybridExecutionReceipt {
1387 pub vector_thread: std::thread::ThreadId,
1389 pub lexical_thread_name: Option<String>,
1391 pub vector_completed: bool,
1393 pub lexical_completed: bool,
1395}
1396
1397pub const DEFAULT_READER_DRAIN_TIMEOUT: Duration = Duration::from_millis(250);
1399
1400#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1402pub enum StoreState {
1403 Open,
1405 Closing,
1407 Closed,
1409}
1410
1411#[derive(Clone, Debug, Eq, PartialEq)]
1413pub struct OpenOptions {
1414 access_mode: AccessMode,
1415 durability_mode: DurabilityMode,
1416 commit_tier: CommitTier,
1417 reader_drain_timeout: Duration,
1418 max_resident_bytes: u64,
1419 max_temp_bytes: u64,
1420 epoch: Option<crate::epoch::StoreEpoch>,
1421 tokenizer: Option<crate::fts::tokenizer::TokenizerConfig>,
1422 schema: Option<crate::meta::Schema>,
1423}
1424
1425#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1427pub enum AccessMode {
1428 ReadWrite,
1430 ReadOnly,
1432}
1433
1434impl OpenOptions {
1435 #[must_use]
1437 pub const fn new() -> Self {
1438 Self {
1439 access_mode: AccessMode::ReadWrite,
1440 durability_mode: DurabilityMode::Derived,
1441 commit_tier: CommitTier::Ordered,
1442 reader_drain_timeout: DEFAULT_READER_DRAIN_TIMEOUT,
1443 max_resident_bytes: u64::MAX,
1444 max_temp_bytes: u64::MAX,
1445 epoch: None,
1446 tokenizer: None,
1447 schema: None,
1448 }
1449 }
1450
1451 #[must_use]
1453 pub const fn read_only() -> Self {
1454 Self {
1455 access_mode: AccessMode::ReadOnly,
1456 durability_mode: DurabilityMode::Derived,
1457 commit_tier: CommitTier::Ordered,
1458 reader_drain_timeout: DEFAULT_READER_DRAIN_TIMEOUT,
1459 max_resident_bytes: u64::MAX,
1460 max_temp_bytes: u64::MAX,
1461 epoch: None,
1462 tokenizer: None,
1463 schema: None,
1464 }
1465 }
1466
1467 #[must_use]
1469 pub const fn with_durability(mut self, mode: DurabilityMode, tier: CommitTier) -> Self {
1470 self.durability_mode = mode;
1471 self.commit_tier = tier;
1472 self
1473 }
1474
1475 #[must_use]
1477 pub const fn with_reader_drain_timeout(mut self, timeout: Duration) -> Self {
1478 self.reader_drain_timeout = timeout;
1479 self
1480 }
1481
1482 #[must_use]
1484 pub const fn with_max_resident_bytes(mut self, bytes: u64) -> Self {
1485 self.max_resident_bytes = bytes;
1486 self
1487 }
1488
1489 #[must_use]
1493 pub fn with_schema(mut self, schema: crate::meta::Schema) -> Self {
1494 self.schema = Some(schema);
1495 self
1496 }
1497
1498 #[must_use]
1500 pub const fn with_max_temp_bytes(mut self, bytes: u64) -> Self {
1501 self.max_temp_bytes = bytes;
1502 self
1503 }
1504
1505 #[must_use]
1511 pub fn with_epoch(mut self, epoch: crate::epoch::StoreEpoch) -> Self {
1512 self.epoch = Some(epoch);
1513 self
1514 }
1515
1516 #[must_use]
1521 pub fn with_tokenizer(mut self, tokenizer: crate::fts::tokenizer::TokenizerConfig) -> Self {
1522 self.tokenizer = Some(tokenizer);
1523 self
1524 }
1525}
1526
1527impl Default for OpenOptions {
1528 fn default() -> Self {
1529 Self::new()
1530 }
1531}
1532
1533#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1535pub struct GraphSearchOptions {
1536 profile: crate::graph::search::GraphSearchProfile,
1537 ef: Option<usize>,
1538 seed: u64,
1539 #[cfg(any(test, feature = "test-support"))]
1540 cancel_after_hops: Option<usize>,
1541}
1542
1543impl GraphSearchOptions {
1544 #[must_use]
1546 pub const fn new(profile: crate::graph::search::GraphSearchProfile) -> Self {
1547 Self {
1548 profile,
1549 ef: None,
1550 seed: 0,
1551 #[cfg(any(test, feature = "test-support"))]
1552 cancel_after_hops: None,
1553 }
1554 }
1555
1556 #[must_use]
1558 pub const fn with_ef(mut self, ef: usize) -> Self {
1559 self.ef = Some(ef);
1560 self
1561 }
1562
1563 #[must_use]
1565 pub const fn with_seed(mut self, seed: u64) -> Self {
1566 self.seed = seed;
1567 self
1568 }
1569
1570 #[cfg(any(test, feature = "test-support"))]
1572 #[must_use]
1573 pub const fn with_cancel_after_hops(mut self, hops: usize) -> Self {
1574 self.cancel_after_hops = Some(hops);
1575 self
1576 }
1577
1578 pub(crate) const fn profile(self) -> crate::graph::search::GraphSearchProfile {
1579 self.profile
1580 }
1581
1582 pub(crate) const fn ef(self) -> Option<usize> {
1583 self.ef
1584 }
1585
1586 pub(crate) const fn seed(self) -> u64 {
1587 self.seed
1588 }
1589
1590 #[cfg(any(test, feature = "test-support"))]
1591 pub(crate) const fn cancel_after_hops(self) -> Option<usize> {
1592 self.cancel_after_hops
1593 }
1594}
1595
1596impl Default for GraphSearchOptions {
1597 fn default() -> Self {
1598 Self::new(crate::graph::search::GraphSearchProfile::SiftClass)
1599 }
1600}
1601
1602pub(crate) fn auto_graph_search_options(
1603 snapshot: &PublishedSnapshot,
1604) -> Result<GraphSearchOptions, QueryError> {
1605 snapshot
1606 .graph_profile()
1607 .map(|profile| GraphSearchOptions::new(profile.search_profile()))
1608 .map_err(crate::graph::search::GraphSearchError::Profile)
1609 .map_err(QueryError::Graph)
1610}
1611
1612#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1614pub enum SearchTier {
1615 #[default]
1617 Auto,
1618 Exact,
1622 Scan,
1624 Graph(GraphSearchOptions),
1626}
1627
1628#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1630pub struct SearchOptions {
1631 scan: crate::scan::ScanOptions,
1632 tier: Option<SearchTier>,
1633 scan_rescore: Option<ScanRescoreOptions>,
1634}
1635
1636#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1637enum GraphBoundMode {
1638 Shared,
1639 #[cfg(test)]
1640 Independent,
1641}
1642
1643impl SearchOptions {
1644 #[must_use]
1646 pub const fn new(scan: crate::scan::ScanOptions) -> Self {
1647 Self {
1648 scan,
1649 tier: None,
1650 scan_rescore: None,
1651 }
1652 }
1653
1654 #[must_use]
1656 pub const fn with_tier(mut self, tier: SearchTier) -> Self {
1657 self.tier = Some(tier);
1658 self.scan_rescore = None;
1659 self
1660 }
1661
1662 #[must_use]
1666 pub const fn with_scan_rescore(mut self, options: ScanRescoreOptions) -> Self {
1667 self.tier = Some(SearchTier::Scan);
1668 self.scan_rescore = Some(options);
1669 self
1670 }
1671
1672 #[must_use]
1674 pub const fn scan_rescore(self) -> Option<ScanRescoreOptions> {
1675 self.scan_rescore
1676 }
1677
1678 #[must_use]
1680 pub const fn scan(self) -> crate::scan::ScanOptions {
1681 self.scan
1682 }
1683
1684 #[must_use]
1687 pub const fn tier(self) -> SearchTier {
1688 match self.tier {
1689 Some(tier) => tier,
1690 None => SearchTier::Auto,
1691 }
1692 }
1693
1694 pub(crate) const fn explicit_tier(self) -> Option<SearchTier> {
1695 self.tier
1696 }
1697}
1698
1699impl From<crate::scan::ScanOptions> for SearchOptions {
1700 fn from(scan: crate::scan::ScanOptions) -> Self {
1701 Self::new(scan)
1702 }
1703}
1704
1705#[derive(Debug)]
1707pub enum StoreError {
1708 Io {
1710 path: PathBuf,
1712 source: std::io::Error,
1714 },
1715 NotDirectory {
1717 path: PathBuf,
1719 },
1720 StoreBusy {
1722 path: PathBuf,
1724 },
1725 Lock(StoreLockError),
1727 Durability(DurabilityPolicyError),
1729 Kernel(crate::kernels::KernelInitError),
1731 Manifest(crate::manifest::ManifestError),
1733 Tokenizer(crate::fts::tokenizer::TokenizerError),
1735 EpochMismatch(crate::epoch::EpochMismatch),
1737 EpochUndeclared,
1739 EpochUnstamped,
1741 SchemaMismatch {
1743 persisted: crate::meta::Schema,
1745 declared: crate::meta::Schema,
1747 },
1748 Segment(crate::segment::SegmentError),
1750 GraphUnavailable {
1752 segment_id: crate::segment::SegmentId,
1754 },
1755 Wal(crate::wal::WalReadError),
1757 WalRecovery(crate::wal::WalRecoveryError),
1759 WalRecord {
1761 seq: crate::wal::LogSeq,
1763 source: crate::wal::VisibleRecordError,
1765 },
1766 WalMutation {
1768 seq: crate::wal::LogSeq,
1770 op: u16,
1772 source: crate::ingest::wal_payload::PayloadError,
1774 },
1775 WalRevisionOrder {
1777 seq: crate::wal::LogSeq,
1779 doc_id: crate::ingest::DocId,
1781 current: crate::ingest::Revision,
1783 attempted: crate::ingest::Revision,
1785 },
1786 UnsupportedWalMutation {
1788 seq: crate::wal::LogSeq,
1790 op: u16,
1792 },
1793 WalVector {
1795 seq: crate::wal::LogSeq,
1797 source: crate::quant::QuantError,
1799 },
1800 WalWrite(crate::wal::WalWriteError),
1802 WalRetire(crate::wal::WalRetireError),
1804 Statistics {
1806 component: &'static str,
1808 source: std::io::Error,
1810 },
1811 BudgetExceeded {
1813 needed: u64,
1815 budget: u64,
1817 component: &'static str,
1819 },
1820 AllocationFailed {
1822 needed: u64,
1824 component: &'static str,
1826 },
1827 DimensionMismatch {
1829 expected: usize,
1831 actual: usize,
1833 },
1834 ActiveRowOverflow,
1836 EmptyActiveSegment,
1838 SealCancelled,
1840 ReadOnly,
1842 SealedTombstoneRecoveryRequired,
1844 ForeignPreparedSegment,
1846 GenerationOverflow,
1848 ConcurrentActiveMutation {
1850 expected_generation: u64,
1852 actual_generation: u64,
1854 },
1855 PartitionBytesOverflow,
1857 PurgeRecovery {
1859 detail: String,
1861 },
1862 Closing,
1864 Closed,
1866 ReadCancelled,
1868 BackgroundStart {
1870 source: std::io::Error,
1872 },
1873 BackgroundHandshake,
1875 BackgroundThreadPanicked,
1877 QueryPoolStart {
1879 source: std::io::Error,
1881 },
1882 QueryPoolHandshake,
1884 QueryPoolThreadPanicked,
1886 QueryPoolCapacity {
1888 source: String,
1890 },
1891 Synchronization {
1893 component: &'static str,
1895 },
1896}
1897
1898impl std::fmt::Display for StoreError {
1899 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1900 match self {
1901 Self::Io { path, source } => {
1902 write!(formatter, "store I/O {}: {source}", path.display())
1903 }
1904 Self::NotDirectory { path } => {
1905 write!(
1906 formatter,
1907 "store path is not a directory: {}",
1908 path.display()
1909 )
1910 }
1911 Self::StoreBusy { path } => {
1912 write!(formatter, "store already has a writer: {}", path.display())
1913 }
1914 Self::Lock(error) => error.fmt(formatter),
1915 Self::Durability(error) => error.fmt(formatter),
1916 Self::Kernel(error) => error.fmt(formatter),
1917 Self::Manifest(error) => error.fmt(formatter),
1918 Self::Tokenizer(error) => write!(formatter, "store tokenizer: {error}"),
1919 Self::EpochMismatch(error) => error.fmt(formatter),
1920 Self::EpochUndeclared => {
1921 formatter.write_str("store epoch is persisted but the caller declared none")
1922 }
1923 Self::EpochUnstamped => formatter
1924 .write_str("store manifest predates epoch identity and cannot adopt a declaration"),
1925 Self::SchemaMismatch {
1926 persisted,
1927 declared,
1928 } => write!(
1929 formatter,
1930 "declared store schema {declared:?} does not match persisted schema {persisted:?}"
1931 ),
1932 Self::Segment(error) => error.fmt(formatter),
1933 Self::GraphUnavailable { segment_id } => {
1934 write!(formatter, "sealed segment {segment_id} has no graph region")
1935 }
1936 Self::Wal(error) => error.fmt(formatter),
1937 Self::WalRecovery(error) => error.fmt(formatter),
1938 Self::WalRecord { seq, source } => {
1939 write!(
1940 formatter,
1941 "WAL sequence {} retained record: {source}",
1942 seq.get()
1943 )
1944 }
1945 Self::WalMutation { seq, op, source } => write!(
1946 formatter,
1947 "WAL sequence {} operation {op} payload: {source}",
1948 seq.get()
1949 ),
1950 Self::WalRevisionOrder {
1951 seq,
1952 doc_id,
1953 current,
1954 attempted,
1955 } => write!(
1956 formatter,
1957 "WAL sequence {} document {} revision {} does not follow {}",
1958 seq.get(),
1959 doc_id.get(),
1960 attempted.get(),
1961 current.get()
1962 ),
1963 Self::UnsupportedWalMutation { seq, op } => write!(
1964 formatter,
1965 "WAL sequence {} operation {op} has no active-segment recovery path",
1966 seq.get()
1967 ),
1968 Self::WalVector { seq, source } => {
1969 write!(formatter, "WAL sequence {} vector: {source}", seq.get())
1970 }
1971 Self::WalWrite(error) => error.fmt(formatter),
1972 Self::WalRetire(error) => error.fmt(formatter),
1973 Self::Statistics { component, source } => {
1974 write!(formatter, "store statistics {component}: {source}")
1975 }
1976 Self::BudgetExceeded {
1977 needed,
1978 budget,
1979 component,
1980 } => write!(
1981 formatter,
1982 "store {component} allocation needs {needed} bytes, budget is {budget} bytes"
1983 ),
1984 Self::AllocationFailed { needed, component } => write!(
1985 formatter,
1986 "store {component} allocator rejected {needed} bytes"
1987 ),
1988 Self::DimensionMismatch { expected, actual } => write!(
1989 formatter,
1990 "active vector dimension {actual} does not match {expected}"
1991 ),
1992 Self::ActiveRowOverflow => {
1993 formatter.write_str("active segment row or byte geometry overflow")
1994 }
1995 Self::EmptyActiveSegment => formatter.write_str("active segment is empty"),
1996 Self::SealCancelled => formatter.write_str("seal was cancelled before commit"),
1997 Self::ReadOnly => formatter.write_str("store handle is read-only"),
1998 Self::SealedTombstoneRecoveryRequired => formatter.write_str(
1999 "store requires writable sealed-tombstone recovery before reads are safe",
2000 ),
2001 Self::ForeignPreparedSegment => {
2002 formatter.write_str("prepared segment belongs to another store")
2003 }
2004 Self::GenerationOverflow => formatter.write_str("store snapshot generation overflow"),
2005 Self::ConcurrentActiveMutation {
2006 expected_generation,
2007 actual_generation,
2008 } => write!(
2009 formatter,
2010 "active generation changed from {expected_generation} to {actual_generation} while preparing a commit"
2011 ),
2012 Self::PartitionBytesOverflow => {
2013 formatter.write_str("partition reclaimed-byte count overflow")
2014 }
2015 Self::PurgeRecovery { detail } => {
2016 write!(formatter, "physical purge recovery failed: {detail}")
2017 }
2018 Self::Closing => formatter.write_str("store is closing"),
2019 Self::Closed => formatter.write_str("store is closed"),
2020 Self::ReadCancelled => formatter.write_str("store close cancelled the admitted read"),
2021 Self::BackgroundStart { source } => {
2022 write!(
2023 formatter,
2024 "store lifecycle thread could not start: {source}"
2025 )
2026 }
2027 Self::BackgroundHandshake => {
2028 formatter.write_str("store lifecycle thread startup handshake failed")
2029 }
2030 Self::BackgroundThreadPanicked => {
2031 formatter.write_str("store lifecycle thread panicked")
2032 }
2033 Self::QueryPoolStart { source } => {
2034 write!(formatter, "store query worker could not start: {source}")
2035 }
2036 Self::QueryPoolHandshake => {
2037 formatter.write_str("store query worker startup handshake failed")
2038 }
2039 Self::QueryPoolThreadPanicked => formatter.write_str("store query worker panicked"),
2040 Self::QueryPoolCapacity { source } => {
2041 write!(formatter, "store query worker capacity failed: {source}")
2042 }
2043 Self::Synchronization { component } => {
2044 write!(
2045 formatter,
2046 "store lifecycle synchronization poisoned: {component}"
2047 )
2048 }
2049 }
2050 }
2051}
2052
2053#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
2058pub enum StoreErrorKind {
2059 Io,
2061 InvalidArgument,
2063 StoreBusy,
2065 Unsupported,
2067 Corrupt,
2069 BudgetExceeded,
2071 OutOfMemory,
2073 DimensionMismatch,
2075 EpochMismatch,
2077 EpochUndeclared,
2079 EpochUnstamped,
2081 Internal,
2083 EmptyBatch,
2085 Cancelled,
2087 ReadOnly,
2089 Closing,
2091 Closed,
2093 Panic,
2095 Synchronization,
2097}
2098
2099impl StoreError {
2100 #[must_use]
2102 pub fn kind(&self) -> StoreErrorKind {
2103 match self {
2104 Self::Io { .. }
2105 | Self::Lock(_)
2106 | Self::Statistics { .. }
2107 | Self::BackgroundStart { .. }
2108 | Self::QueryPoolStart { .. }
2109 | Self::WalWrite(_)
2110 | Self::WalRetire(_) => StoreErrorKind::Io,
2111 Self::NotDirectory { .. } | Self::Tokenizer(_) | Self::SchemaMismatch { .. } => {
2112 StoreErrorKind::InvalidArgument
2113 }
2114 Self::StoreBusy { .. } => StoreErrorKind::StoreBusy,
2115 Self::Durability(_)
2116 | Self::Kernel(_)
2117 | Self::GraphUnavailable { .. }
2118 | Self::UnsupportedWalMutation { .. } => StoreErrorKind::Unsupported,
2119 Self::Manifest(_)
2120 | Self::Segment(_)
2121 | Self::Wal(_)
2122 | Self::WalRecovery(_)
2123 | Self::WalRecord { .. }
2124 | Self::WalMutation { .. }
2125 | Self::WalRevisionOrder { .. }
2126 | Self::WalVector { .. }
2127 | Self::PurgeRecovery { .. } => StoreErrorKind::Corrupt,
2128 Self::BudgetExceeded { .. } => StoreErrorKind::BudgetExceeded,
2129 Self::AllocationFailed { .. } => StoreErrorKind::OutOfMemory,
2130 Self::DimensionMismatch { .. } => StoreErrorKind::DimensionMismatch,
2131 Self::EpochMismatch(_) => StoreErrorKind::EpochMismatch,
2132 Self::EpochUndeclared => StoreErrorKind::EpochUndeclared,
2133 Self::EpochUnstamped => StoreErrorKind::EpochUnstamped,
2134 Self::ActiveRowOverflow
2135 | Self::GenerationOverflow
2136 | Self::ConcurrentActiveMutation { .. }
2137 | Self::PartitionBytesOverflow
2138 | Self::ForeignPreparedSegment
2139 | Self::BackgroundHandshake
2140 | Self::QueryPoolHandshake
2141 | Self::QueryPoolCapacity { .. } => StoreErrorKind::Internal,
2142 Self::EmptyActiveSegment => StoreErrorKind::EmptyBatch,
2143 Self::SealCancelled | Self::ReadCancelled => StoreErrorKind::Cancelled,
2144 Self::ReadOnly | Self::SealedTombstoneRecoveryRequired => StoreErrorKind::ReadOnly,
2145 Self::Closing => StoreErrorKind::Closing,
2146 Self::Closed => StoreErrorKind::Closed,
2147 Self::BackgroundThreadPanicked | Self::QueryPoolThreadPanicked => StoreErrorKind::Panic,
2148 Self::Synchronization { .. } => StoreErrorKind::Synchronization,
2149 }
2150 }
2151}
2152
2153impl std::error::Error for StoreError {
2154 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
2155 match self {
2156 Self::Io { source, .. } => Some(source),
2157 Self::Lock(error) => Some(error),
2158 Self::Durability(error) => Some(error),
2159 Self::Kernel(error) => Some(error),
2160 Self::Manifest(error) => Some(error),
2161 Self::Tokenizer(error) => Some(error),
2162 Self::EpochMismatch(error) => Some(error),
2163 Self::Segment(error) => Some(error),
2164 Self::Wal(error) => Some(error),
2165 Self::WalRecovery(error) => Some(error),
2166 Self::WalRecord { source, .. } => Some(source),
2167 Self::WalMutation { source, .. } => Some(source),
2168 Self::WalVector { source, .. } => Some(source),
2169 Self::WalWrite(error) => Some(error),
2170 Self::WalRetire(error) => Some(error),
2171 Self::Statistics { source, .. } => Some(source),
2172 Self::BackgroundStart { source } => Some(source),
2173 Self::QueryPoolStart { source } => Some(source),
2174 Self::NotDirectory { .. }
2175 | Self::StoreBusy { .. }
2176 | Self::EpochUndeclared
2177 | Self::EpochUnstamped
2178 | Self::SchemaMismatch { .. }
2179 | Self::GraphUnavailable { .. }
2180 | Self::WalRevisionOrder { .. }
2181 | Self::UnsupportedWalMutation { .. }
2182 | Self::BudgetExceeded { .. }
2183 | Self::AllocationFailed { .. }
2184 | Self::DimensionMismatch { .. }
2185 | Self::ActiveRowOverflow
2186 | Self::EmptyActiveSegment
2187 | Self::SealCancelled
2188 | Self::ReadOnly
2189 | Self::SealedTombstoneRecoveryRequired
2190 | Self::ForeignPreparedSegment
2191 | Self::GenerationOverflow
2192 | Self::ConcurrentActiveMutation { .. }
2193 | Self::PartitionBytesOverflow
2194 | Self::PurgeRecovery { .. }
2195 | Self::Closing
2196 | Self::Closed
2197 | Self::ReadCancelled
2198 | Self::BackgroundHandshake
2199 | Self::BackgroundThreadPanicked
2200 | Self::QueryPoolHandshake
2201 | Self::QueryPoolThreadPanicked
2202 | Self::QueryPoolCapacity { .. }
2203 | Self::Synchronization { .. } => None,
2204 }
2205 }
2206}
2207
2208pub struct Store {
2210 pub(crate) directory: PathBuf,
2211 pub(crate) vfs: Arc<dyn crate::vfs::Vfs>,
2212 pub(crate) clock: Arc<dyn MonotonicClock>,
2213 pub(crate) state: Mutex<StoreState>,
2214 pub(crate) state_changed: Condvar,
2215 pub(crate) background: Mutex<Option<BackgroundThread>>,
2216 pub(crate) query_pool: Mutex<Option<Arc<pool::QueryPool>>>,
2217 pub(crate) lexical_worker: Mutex<Option<Arc<pool::LexicalWorker>>>,
2218 pub(crate) snapshot: RwLock<Option<Arc<PublishedSnapshot>>>,
2219 pub(crate) active: Mutex<Option<crate::ingest::ActiveState>>,
2222 pub(crate) wal_writer: Mutex<Option<crate::ingest::StoreWal>>,
2223 pub(crate) writer_lock: Mutex<Option<StoreLock>>,
2224 pub(crate) maintenance: Mutex<()>,
2225 pub(crate) health_state: Mutex<crate::diag::HealthState>,
2226 pub(crate) durability_policy: DurabilityPolicy,
2227 pub(crate) reader_drain_timeout: Duration,
2228 pub(crate) accounting: Arc<stats::Accounting>,
2229 pub(crate) active_queries: AtomicU64,
2230 #[cfg(any(test, feature = "test-support"))]
2231 text_materialization_work: materialize::TestMaterializationWork,
2232 pub(crate) epoch: Option<crate::epoch::StoreEpoch>,
2233 pub(crate) epoch_alias: crate::epoch::EpochAliasCell,
2234 pub(crate) tokenizer: crate::fts::tokenizer::Analyzer,
2235 pub(crate) schema: crate::meta::Schema,
2236 lexical_index_cache: LexicalIndexCache,
2240 #[cfg(any(test, feature = "test-support"))]
2241 pub(crate) ingest_retention_fault_controller:
2242 Option<crate::ingest::IngestRetentionFaultController>,
2243 #[cfg(any(test, feature = "test-support"))]
2244 hybrid_leg_fault: Mutex<Option<HybridLegTestFault>>,
2245 #[cfg(any(test, feature = "test-support"))]
2246 hybrid_execution_receipt: Mutex<Option<HybridExecutionReceipt>>,
2247 #[cfg(any(test, feature = "test-support"))]
2248 pub(crate) metadata_test_controller: Option<Arc<crate::planner::MetadataTestController>>,
2249 #[cfg(any(test, feature = "test-support"))]
2250 pub(crate) vector_fault_controller: Option<crate::scan::vector_fault::VectorFaultController>,
2251 #[cfg(any(test, feature = "test-support"))]
2252 pub(crate) kernel_fault_controller: Option<crate::kernels::vector_fault::KernelFaultController>,
2253 #[cfg(any(test, feature = "test-support"))]
2254 pub(crate) vector_seal_scheme: Option<crate::quant::QuantScheme>,
2255 #[cfg(test)]
2256 pub(crate) teardown_probe: Arc<close::TeardownProbe>,
2257}
2258
2259fn acquire_writer_lock(
2260 path: &Path,
2261 access_mode: AccessMode,
2262) -> Result<Option<StoreLock>, StoreError> {
2263 Ok(match access_mode {
2264 AccessMode::ReadWrite => Some(StoreLock::acquire(path).map_err(|error| match error {
2265 StoreLockError::Io { path, source }
2266 if source.kind() == std::io::ErrorKind::WouldBlock =>
2267 {
2268 StoreError::StoreBusy { path }
2269 }
2270 StoreLockError::Io { path, source } => {
2271 StoreError::Lock(StoreLockError::Io { path, source })
2272 }
2273 })?),
2274 AccessMode::ReadOnly => None,
2275 })
2276}
2277
2278fn resolve_open_schema(
2279 manifest_exists: bool,
2280 snapshot: &PublishedSnapshot,
2281 declared: Option<&crate::meta::Schema>,
2282) -> Result<crate::meta::Schema, StoreError> {
2283 if manifest_exists {
2284 let persisted = snapshot.schema().clone();
2285 if let Some(declared) = declared
2286 && declared != &persisted
2287 {
2288 return Err(StoreError::SchemaMismatch {
2289 persisted,
2290 declared: declared.clone(),
2291 });
2292 }
2293 Ok(snapshot.schema().clone())
2294 } else {
2295 Ok(declared
2296 .cloned()
2297 .unwrap_or_else(crate::meta::Schema::timestamp_only))
2298 }
2299}
2300
2301fn validate_epoch_identity(
2302 persisted: Option<crate::epoch::EpochIdentity>,
2303 declared: Option<crate::epoch::EpochIdentity>,
2304 manifest_exists: bool,
2305 access_mode: AccessMode,
2306) -> Result<(), StoreError> {
2307 match (persisted, declared) {
2308 (Some(expected), Some(declared)) if expected != declared => {
2309 Err(StoreError::EpochMismatch(crate::epoch::EpochMismatch {
2310 expected,
2311 declared,
2312 }))
2313 }
2314 (Some(_), None) => Err(StoreError::EpochUndeclared),
2315 (None, Some(_)) if manifest_exists || access_mode == AccessMode::ReadOnly => {
2316 Err(StoreError::EpochUnstamped)
2317 }
2318 (Some(_), Some(_)) | (None, Some(_)) | (None, None) => Ok(()),
2319 }
2320}
2321
2322fn validate_tokenizer_epoch(
2323 expected: Option<crate::epoch::EpochIdentity>,
2324 tokenizer: &crate::fts::tokenizer::Analyzer,
2325) -> Result<(), StoreError> {
2326 if let Some(expected) = expected
2327 && tokenizer.epoch() != expected.tokenizer
2328 {
2329 return Err(StoreError::EpochMismatch(crate::epoch::EpochMismatch {
2330 expected,
2331 declared: crate::epoch::EpochIdentity {
2332 embedding: expected.embedding,
2333 tokenizer: tokenizer.epoch(),
2334 },
2335 }));
2336 }
2337 Ok(())
2338}
2339
2340fn cleanup_open_orphans(
2341 store: &Store,
2342) -> Result<crate::manifest::io::OrphanCleanupReport, StoreError> {
2343 let reachable_segments = store
2344 .snapshot
2345 .read()
2346 .map_err(|_| StoreError::Synchronization {
2347 component: "published snapshot",
2348 })?
2349 .as_ref()
2350 .map(|snapshot| {
2351 snapshot
2352 .all_segments()
2353 .iter()
2354 .map(|segment| store.directory.join(segment.meta().id.file_name()))
2355 .collect::<HashSet<_>>()
2356 })
2357 .ok_or(StoreError::Closed)?;
2358 crate::manifest::io::cleanup_store_orphans(
2359 store.vfs.as_ref(),
2360 &store.directory,
2361 &reachable_segments,
2362 store.durability_policy,
2363 )
2364 .map_err(StoreError::Manifest)
2365}
2366
2367impl Store {
2368 pub fn open(path: impl AsRef<Path>, options: OpenOptions) -> Result<Self, StoreError> {
2370 Self::open_with_infrastructure(
2371 path,
2372 options,
2373 Arc::new(crate::vfs::StdVfs),
2374 Arc::new(SystemMonotonicClock),
2375 #[cfg(any(test, feature = "test-support"))]
2376 None,
2377 #[cfg(any(test, feature = "test-support"))]
2378 None,
2379 #[cfg(any(test, feature = "test-support"))]
2380 None,
2381 #[cfg(any(test, feature = "test-support"))]
2382 None,
2383 #[cfg(any(test, feature = "test-support"))]
2384 None,
2385 #[cfg(any(test, feature = "test-support"))]
2386 None,
2387 #[cfg(any(test, feature = "test-support"))]
2388 None,
2389 )
2390 }
2391
2392 #[allow(
2393 clippy::too_many_arguments,
2394 reason = "test-support controllers are explicit optional infrastructure dependencies"
2395 )]
2396 fn open_with_infrastructure(
2397 path: impl AsRef<Path>,
2398 options: OpenOptions,
2399 vfs: Arc<dyn crate::vfs::Vfs>,
2400 clock: Arc<dyn MonotonicClock>,
2401 #[cfg(any(test, feature = "test-support"))] hybrid_leg_fault: Option<HybridLegTestFault>,
2402 #[cfg(any(test, feature = "test-support"))] storage_fault_controller: Option<
2403 StorageFaultController,
2404 >,
2405 #[cfg(any(test, feature = "test-support"))] ingest_retention_fault_controller: Option<
2406 crate::ingest::IngestRetentionFaultController,
2407 >,
2408 #[cfg(any(test, feature = "test-support"))] metadata_test_controller: Option<
2409 Arc<crate::planner::MetadataTestController>,
2410 >,
2411 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
2412 crate::scan::vector_fault::VectorFaultController,
2413 >,
2414 #[cfg(any(test, feature = "test-support"))] kernel_fault_controller: Option<
2415 crate::kernels::vector_fault::KernelFaultController,
2416 >,
2417 #[cfg(any(test, feature = "test-support"))] vector_seal_scheme: Option<
2418 crate::quant::QuantScheme,
2419 >,
2420 ) -> Result<Self, StoreError> {
2421 let path = path.as_ref();
2422 #[cfg(any(test, feature = "test-support"))]
2423 match kernel_fault_controller.as_ref() {
2424 Some(controller) => {
2425 controller
2426 .initialize_for_store()
2427 .map_err(StoreError::Kernel)?;
2428 }
2429 None => {
2430 crate::kernels::initialize().map_err(StoreError::Kernel)?;
2431 }
2432 }
2433 #[cfg(not(any(test, feature = "test-support")))]
2434 crate::kernels::initialize().map_err(StoreError::Kernel)?;
2435 #[cfg(any(test, feature = "test-support"))]
2436 let vfs = match storage_fault_controller.as_ref() {
2437 Some(controller) => Arc::new(StorageFaultVfs {
2438 inner: vfs,
2439 controller: controller.clone(),
2440 }) as Arc<dyn crate::vfs::Vfs>,
2441 None => vfs,
2442 };
2443 let durability_policy = DurabilityPolicy::new(options.durability_mode, options.commit_tier)
2444 .map_err(StoreError::Durability)?;
2445 let accounting = Arc::new(stats::Accounting::new(
2446 options.max_resident_bytes,
2447 options.max_temp_bytes,
2448 ));
2449 let is_directory = vfs
2450 .ensure_directory(path, options.access_mode == AccessMode::ReadWrite)
2451 .map_err(|source| StoreError::Io {
2452 path: path.to_path_buf(),
2453 source,
2454 })?;
2455 if !is_directory {
2456 return Err(StoreError::NotDirectory {
2457 path: path.to_path_buf(),
2458 });
2459 }
2460 let writer_lock = acquire_writer_lock(path, options.access_mode)?;
2461 let manifest_path = path.join(crate::manifest::io::MANIFEST_FILE);
2462 let manifest_exists = match vfs.open(&manifest_path) {
2463 Ok(_) => true,
2464 Err(source) if source.kind() == std::io::ErrorKind::NotFound => false,
2465 Err(source) => {
2466 return Err(StoreError::Io {
2467 path: manifest_path,
2468 source,
2469 });
2470 }
2471 };
2472 let snapshot = PublishedSnapshot::load_for_open_on_vfs(path, &accounting, vfs.as_ref())?;
2477 let schema = resolve_open_schema(manifest_exists, &snapshot, options.schema.as_ref())?;
2478 let persisted_epoch = snapshot.epoch_alias();
2479 let declared_epoch = options
2480 .epoch
2481 .as_ref()
2482 .map(crate::epoch::StoreEpoch::identity);
2483 validate_epoch_identity(
2484 persisted_epoch,
2485 declared_epoch,
2486 manifest_exists,
2487 options.access_mode,
2488 )?;
2489 let tokenizer = crate::fts::tokenizer::Analyzer::new(
2490 options
2491 .tokenizer
2492 .clone()
2493 .unwrap_or_else(crate::fts::tokenizer::TokenizerConfig::text_default),
2494 )
2495 .map_err(StoreError::Tokenizer)?;
2496 validate_tokenizer_epoch(persisted_epoch.or(declared_epoch), &tokenizer)?;
2497 let absorbed_through = snapshot.absorbed_through();
2498 let wal_path = path.join("wal.ze");
2499 let (active, recovered_wal, sealed_tombstones) = crate::ingest::ActiveState::recover(
2500 vfs.as_ref(),
2501 &wal_path,
2502 snapshot.generation(),
2503 absorbed_through,
2504 &accounting,
2505 &schema,
2506 &tokenizer,
2507 )?;
2508 if options.access_mode == AccessMode::ReadWrite && manifest_exists {
2509 match durability_policy.directory_sync() {
2513 SyncRequirement::Skip => {}
2514 SyncRequirement::Sync(kind) => {
2515 vfs.sync(path, kind).map_err(|source| StoreError::Io {
2516 path: path.to_path_buf(),
2517 source,
2518 })?;
2519 }
2520 }
2521 }
2522 if options.access_mode == AccessMode::ReadWrite
2523 && !manifest_exists
2524 && (options.epoch.is_some() || options.schema.is_some())
2525 {
2526 crate::manifest::io::commit_manifest(
2527 vfs.as_ref(),
2528 path,
2529 &crate::manifest::Manifest {
2530 generation: active.generation,
2531 log_seq: 0,
2532 segments: Vec::new(),
2533 epochs: options
2534 .epoch
2535 .as_ref()
2536 .cloned()
2537 .map(crate::manifest::EpochMeta::from)
2538 .into_iter()
2539 .collect(),
2540 epoch_alias: options
2541 .epoch
2542 .as_ref()
2543 .map(crate::epoch::StoreEpoch::identity),
2544 schema: schema.clone(),
2545 },
2546 durability_policy,
2547 )
2548 .map_err(StoreError::Manifest)?;
2549 }
2550 let wal_writer = match options.access_mode {
2551 AccessMode::ReadWrite => Some(match recovered_wal {
2552 Some(recovered) => crate::ingest::StoreWal::resume(
2553 vfs.as_ref(),
2554 &wal_path,
2555 recovered,
2556 durability_policy,
2557 absorbed_through,
2558 &accounting,
2559 )?,
2560 None => crate::ingest::StoreWal::create(
2561 Arc::clone(&vfs),
2562 path,
2563 &wal_path,
2564 durability_policy,
2565 &accounting,
2566 )?,
2567 }),
2568 AccessMode::ReadOnly => {
2569 drop(recovered_wal);
2570 None
2571 }
2572 };
2573 let background = match options.access_mode {
2574 AccessMode::ReadWrite => Some(BackgroundThread::start()?),
2575 AccessMode::ReadOnly => None,
2576 };
2577 #[cfg(test)]
2578 let (snapshot, background, teardown_probe) = {
2579 let mut snapshot = snapshot;
2580 let mut background = background;
2581 let teardown_probe = Arc::new(close::TeardownProbe::new());
2582 snapshot.set_teardown_probe(Arc::clone(&teardown_probe));
2583 if let Some(background) = background.as_mut() {
2584 background.set_teardown_probe(Arc::clone(&teardown_probe));
2585 }
2586 (snapshot, background, teardown_probe)
2587 };
2588 let store = Self {
2589 directory: path.to_path_buf(),
2590 vfs,
2591 clock,
2592 state: Mutex::new(StoreState::Open),
2593 state_changed: Condvar::new(),
2594 background: Mutex::new(background),
2595 query_pool: Mutex::new(None),
2596 lexical_worker: Mutex::new(None),
2597 snapshot: RwLock::new(None),
2598 active: Mutex::new(Some(active)),
2599 wal_writer: Mutex::new(wal_writer),
2600 writer_lock: Mutex::new(writer_lock),
2601 maintenance: Mutex::new(()),
2602 health_state: Mutex::new(crate::diag::HealthState::default()),
2603 durability_policy,
2604 reader_drain_timeout: options.reader_drain_timeout,
2605 accounting,
2606 active_queries: AtomicU64::new(0),
2607 #[cfg(any(test, feature = "test-support"))]
2608 text_materialization_work: materialize::TestMaterializationWork::default(),
2609 epoch_alias: crate::epoch::EpochAliasCell::new(persisted_epoch.or(declared_epoch)),
2610 epoch: options.epoch,
2611 tokenizer,
2612 schema,
2613 lexical_index_cache: LexicalIndexCache::new(),
2614 #[cfg(any(test, feature = "test-support"))]
2615 ingest_retention_fault_controller,
2616 #[cfg(any(test, feature = "test-support"))]
2617 hybrid_leg_fault: Mutex::new(hybrid_leg_fault),
2618 #[cfg(any(test, feature = "test-support"))]
2619 hybrid_execution_receipt: Mutex::new(None),
2620 #[cfg(any(test, feature = "test-support"))]
2621 metadata_test_controller,
2622 #[cfg(any(test, feature = "test-support"))]
2623 vector_fault_controller,
2624 #[cfg(any(test, feature = "test-support"))]
2625 kernel_fault_controller,
2626 #[cfg(any(test, feature = "test-support"))]
2627 vector_seal_scheme,
2628 #[cfg(test)]
2629 teardown_probe,
2630 };
2631 store.publish_snapshot(snapshot)?;
2632 store
2633 .recover_pending_physical_purge()
2634 .map_err(|error| StoreError::PurgeRecovery {
2635 detail: error.to_string(),
2636 })?;
2637 store.recover_sealed_tombstones(&sealed_tombstones)?;
2638 if options.access_mode == AccessMode::ReadWrite {
2639 let cleanup_report = cleanup_open_orphans(&store)?;
2640 #[cfg(any(test, feature = "test-support"))]
2641 if let Some(controller) = storage_fault_controller.as_ref() {
2642 controller.record_cleanup_report(cleanup_report);
2643 }
2644 #[cfg(not(any(test, feature = "test-support")))]
2645 let _ = cleanup_report;
2646 }
2647 #[cfg(any(test, feature = "test-support"))]
2648 if let Some(controller) = storage_fault_controller.as_ref() {
2649 let published = store
2650 .snapshot
2651 .read()
2652 .map_err(|_| StoreError::Synchronization {
2653 component: "published snapshot",
2654 })?;
2655 let snapshot = published.as_ref().ok_or(StoreError::Closed)?;
2656 for segment in snapshot.all_segments() {
2657 register_storage_segment_controller(
2658 &store.directory,
2659 segment.meta().id,
2660 controller,
2661 );
2662 }
2663 }
2664 Ok(store)
2665 }
2666
2667 #[cfg(any(test, feature = "test-support"))]
2669 #[doc(hidden)]
2670 pub fn open_with_test_dependencies(
2671 path: impl AsRef<Path>,
2672 options: OpenOptions,
2673 dependencies: StoreTestDependencies,
2674 ) -> Result<Self, StoreError> {
2675 let controller = dependencies.storage_fault_controller.clone();
2676 with_storage_open_controller(controller.as_ref(), || {
2677 Self::open_with_infrastructure(
2678 path,
2679 options,
2680 dependencies.vfs,
2681 dependencies.clock,
2682 dependencies.hybrid_leg_fault,
2683 dependencies.storage_fault_controller,
2684 dependencies.ingest_retention_fault_controller,
2685 dependencies.metadata_test_controller,
2686 dependencies.vector_fault_controller,
2687 dependencies.kernel_fault_controller,
2688 dependencies.vector_seal_scheme,
2689 )
2690 })
2691 }
2692
2693 #[cfg(any(test, feature = "test-support"))]
2695 #[doc(hidden)]
2696 pub fn take_hybrid_execution_receipt(&self) -> Option<HybridExecutionReceipt> {
2697 self.hybrid_execution_receipt
2698 .lock()
2699 .ok()
2700 .and_then(|mut receipt| receipt.take())
2701 }
2702
2703 pub(crate) fn epoch_registry(
2711 &self,
2712 prior: &[crate::manifest::EpochMeta],
2713 ) -> Vec<crate::manifest::EpochMeta> {
2714 if prior.is_empty() {
2715 self.epoch
2716 .as_ref()
2717 .map(crate::manifest::EpochMeta::from)
2718 .into_iter()
2719 .collect()
2720 } else {
2721 prior.to_vec()
2722 }
2723 }
2724
2725 #[must_use]
2728 pub fn epoch_identity(&self) -> Option<crate::epoch::EpochIdentity> {
2729 self.epoch_alias.load()
2730 }
2731
2732 pub fn state(&self) -> Result<StoreState, StoreError> {
2734 self.state
2735 .lock()
2736 .map(|state| *state)
2737 .map_err(|_| StoreError::Synchronization { component: "state" })
2738 }
2739
2740 pub fn snapshot(&self) -> Result<SnapshotLease, StoreError> {
2742 let state = self
2743 .state
2744 .lock()
2745 .map_err(|_| StoreError::Synchronization { component: "state" })?;
2746 match *state {
2747 StoreState::Open => {}
2748 StoreState::Closing => return Err(StoreError::Closing),
2749 StoreState::Closed => return Err(StoreError::Closed),
2750 }
2751 let active = self
2752 .active
2753 .lock()
2754 .map_err(|_| StoreError::Synchronization {
2755 component: "active segment",
2756 })?;
2757 let generation = active.as_ref().ok_or(StoreError::Closed)?.generation;
2758 let snapshot = self
2759 .snapshot
2760 .read()
2761 .map_err(|_| StoreError::Synchronization {
2762 component: "published snapshot",
2763 })?
2764 .as_ref()
2765 .cloned()
2766 .ok_or(StoreError::Closed)?;
2767 drop(active);
2768 drop(state);
2769 Ok(SnapshotLease::new_at(snapshot, generation))
2770 }
2771
2772 #[cfg(any(test, feature = "test-support"))]
2778 #[doc(hidden)]
2779 pub fn test_metadata_row_values(
2780 &self,
2781 row_id: crate::ingest::GlobalRowId,
2782 ) -> Result<Vec<(crate::meta::ColumnId, Option<crate::meta::PredicateValue>)>, StoreError> {
2783 match row_id.source() {
2784 crate::ingest::RowSource::Active => {
2785 let row = usize::try_from(row_id.local_row())
2786 .map_err(|_| StoreError::ActiveRowOverflow)?;
2787 let active = self
2788 .active
2789 .lock()
2790 .map_err(|_| StoreError::Synchronization {
2791 component: "active segment",
2792 })?;
2793 let segment = &active.as_ref().ok_or(StoreError::Closed)?.segment;
2794 let timestamp = segment
2795 .timestamps()
2796 .get(row)
2797 .copied()
2798 .ok_or(StoreError::ActiveRowOverflow)?;
2799 let mut values = segment.column_values(row)?;
2800 values.push((
2801 crate::meta::TIMESTAMP_COLUMN,
2802 crate::meta::PredicateValue::I64(timestamp),
2803 ));
2804 self.schema
2805 .columns()
2806 .iter()
2807 .map(|definition| {
2808 let value = values
2809 .iter()
2810 .find(|(column, _)| *column == definition.id())
2811 .map(|(_, value)| value.clone());
2812 Ok((definition.id(), value))
2813 })
2814 .collect()
2815 }
2816 crate::ingest::RowSource::Sealed(segment_id) => {
2817 let snapshot = self.snapshot()?;
2818 let segment = snapshot
2819 .segments()
2820 .iter()
2821 .find(|segment| segment.meta().id == segment_id)
2822 .ok_or_else(|| {
2823 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
2824 "public metadata row names absent segment {segment_id}"
2825 )))
2826 })?;
2827 let columns = segment.columns().map_err(StoreError::Segment)?;
2828 if row_id.local_row() >= columns.row_count() {
2829 return Err(StoreError::Segment(crate::segment::SegmentError::Geometry(
2830 format!(
2831 "public metadata row {} is outside {} rows",
2832 row_id.local_row(),
2833 columns.row_count()
2834 ),
2835 )));
2836 }
2837 columns
2838 .schema()
2839 .columns()
2840 .iter()
2841 .map(|definition| {
2842 let column = columns.column(definition.id()).ok_or_else(|| {
2843 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
2844 "public metadata row is missing column {}",
2845 definition.id().get()
2846 )))
2847 })?;
2848 let value = match column {
2849 crate::meta::Column::U64(column) => column
2850 .get(row_id.local_row())
2851 .map(crate::meta::PredicateValue::U64),
2852 crate::meta::Column::I64(column) => column
2853 .get(row_id.local_row())
2854 .map(crate::meta::PredicateValue::I64),
2855 crate::meta::Column::F64(column) => column
2856 .get(row_id.local_row())
2857 .map(crate::meta::PredicateValue::F64),
2858 crate::meta::Column::Bool(column) => column
2859 .get(row_id.local_row())
2860 .map(crate::meta::PredicateValue::Bool),
2861 crate::meta::Column::DictionaryString(column) => column
2862 .get(row_id.local_row())
2863 .map(|value| crate::meta::PredicateValue::String(value.to_owned())),
2864 crate::meta::Column::RawString(column) => column
2865 .get(row_id.local_row())
2866 .map(|value| crate::meta::PredicateValue::String(value.to_owned())),
2867 };
2868 Ok((definition.id(), value))
2869 })
2870 .collect()
2871 }
2872 }
2873 }
2874
2875 #[cfg(any(test, feature = "test-support"))]
2877 #[doc(hidden)]
2878 pub fn test_metadata_evaluate_source(
2879 &self,
2880 source: crate::ingest::RowSource,
2881 predicate: &crate::meta::Predicate,
2882 ) -> Result<(u32, Vec<u32>, Vec<u32>), StoreError> {
2883 let evaluate = |columns: &crate::meta::ColumnStore,
2884 alive: &crate::meta::AliveSet|
2885 -> Result<(u32, Vec<u32>, Vec<u32>), StoreError> {
2886 let matches = crate::meta::evaluate(predicate, columns, alive).map_err(|error| {
2887 StoreError::Segment(crate::segment::SegmentError::Columns(format!(
2888 "test metadata source evaluation: {error}"
2889 )))
2890 })?;
2891 Ok((
2892 columns.row_count(),
2893 alive.iter_alive().collect(),
2894 matches.iter().collect(),
2895 ))
2896 };
2897 match source {
2898 crate::ingest::RowSource::Active => {
2899 let active = self
2900 .active
2901 .lock()
2902 .map_err(|_| StoreError::Synchronization {
2903 component: "active segment",
2904 })?;
2905 let segment = &active.as_ref().ok_or(StoreError::Closed)?.segment;
2906 let mut builder = crate::meta::ColumnStoreBuilder::new(self.schema.clone());
2907 for (row, timestamp) in segment.timestamps().iter().copied().enumerate() {
2908 let values = segment.column_values(row)?;
2909 builder
2910 .push_row(timestamp, &crate::ingest::column_inputs(&values))
2911 .map_err(|error| {
2912 StoreError::Segment(crate::segment::SegmentError::Columns(
2913 error.to_string(),
2914 ))
2915 })?;
2916 }
2917 let columns = builder.finish().map_err(|error| {
2918 StoreError::Segment(crate::segment::SegmentError::Columns(error.to_string()))
2919 })?;
2920 let alive = segment.alive()?;
2921 evaluate(&columns, &alive)
2922 }
2923 crate::ingest::RowSource::Sealed(segment_id) => {
2924 let snapshot = self.snapshot()?;
2925 let segment = snapshot
2926 .segments()
2927 .iter()
2928 .find(|segment| segment.meta().id == segment_id)
2929 .ok_or_else(|| {
2930 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
2931 "metadata source evaluation names absent segment {segment_id}"
2932 )))
2933 })?;
2934 let columns = segment.columns().map_err(StoreError::Segment)?;
2935 let alive = segment.alive().map_err(StoreError::Segment)?;
2936 evaluate(&columns, &alive)
2937 }
2938 }
2939 }
2940
2941 pub(crate) fn publish_snapshot(&self, snapshot: PublishedSnapshot) -> Result<(), StoreError> {
2942 let state = self
2943 .state
2944 .lock()
2945 .map_err(|_| StoreError::Synchronization { component: "state" })?;
2946 match *state {
2947 StoreState::Open => {}
2948 StoreState::Closing => return Err(StoreError::Closing),
2949 StoreState::Closed => return Err(StoreError::Closed),
2950 }
2951 let generation = snapshot.generation();
2952 let mut active = self
2953 .active
2954 .lock()
2955 .map_err(|_| StoreError::Synchronization {
2956 component: "active segment",
2957 })?;
2958 let active_generation = &mut active.as_mut().ok_or(StoreError::Closed)?.generation;
2959 *active_generation = (*active_generation).max(generation);
2960 let mut published = self
2961 .snapshot
2962 .write()
2963 .map_err(|_| StoreError::Synchronization {
2964 component: "published snapshot",
2965 })?;
2966 *published = Some(Arc::new(snapshot));
2967 drop(active);
2968 drop(state);
2969 Ok(())
2970 }
2971
2972 pub fn top_k_with_options(
2977 &self,
2978 request: crate::scan::ScanRequest<'_>,
2979 k: usize,
2980 options: crate::scan::ScanOptions,
2981 control: QueryControl,
2982 ) -> Result<crate::scan::ScanOutcome, QueryError> {
2983 let control = control.with_clock(Arc::clone(&self.clock));
2984 let state = self
2985 .state
2986 .lock()
2987 .map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
2988 match *state {
2989 StoreState::Open => {}
2990 StoreState::Closing => return Err(QueryError::Store(StoreError::Closing)),
2991 StoreState::Closed => return Err(QueryError::Store(StoreError::Closed)),
2992 }
2993 let snapshot = self
2994 .snapshot
2995 .read()
2996 .map_err(|_| {
2997 QueryError::Store(StoreError::Synchronization {
2998 component: "published snapshot",
2999 })
3000 })?
3001 .as_ref()
3002 .cloned()
3003 .ok_or(QueryError::Store(StoreError::Closed))?;
3004 let mut pool_slot = self.query_pool.lock().map_err(|_| {
3005 QueryError::Store(StoreError::Synchronization {
3006 component: "query pool",
3007 })
3008 })?;
3009 if pool_slot.is_none() {
3010 let capacity = crate::scan::physical_thread_capacity().map_err(|error| {
3011 QueryError::Store(StoreError::QueryPoolCapacity {
3012 source: error.to_string(),
3013 })
3014 })?;
3015 let pool =
3016 pool::QueryPool::start(capacity, &self.accounting).map_err(QueryError::Store)?;
3017 *pool_slot = Some(Arc::new(pool));
3018 }
3019 let pool = pool_slot.as_ref().cloned().ok_or({
3020 QueryError::Store(StoreError::Synchronization {
3021 component: "query pool initialization",
3022 })
3023 })?;
3024 self.active_queries.fetch_add(1, Ordering::Relaxed);
3025 let active = ActiveQuery {
3026 count: &self.active_queries,
3027 };
3028 drop(pool_slot);
3029 drop(state);
3030
3031 let result = pool
3032 .execute(request, k, options, control, SnapshotLease::new(snapshot))
3033 .map(|mut outcome| {
3034 outcome.candidates.truncate(k);
3035 outcome
3036 });
3037 drop(active);
3038 result
3039 }
3040
3041 #[cfg(any(test, feature = "test-support"))]
3046 pub fn panic_query_worker_for_test(&self) -> Result<(), StoreError> {
3047 let pool = self
3048 .query_pool
3049 .lock()
3050 .map_err(|_| StoreError::Synchronization {
3051 component: "query pool",
3052 })?
3053 .as_ref()
3054 .cloned()
3055 .ok_or(StoreError::Synchronization {
3056 component: "query pool initialization",
3057 })?;
3058 pool.panic_one_and_join()
3059 }
3060
3061 pub fn search(
3070 &self,
3071 request: crate::ingest::SearchRequest<'_>,
3072 k: usize,
3073 options: impl Into<SearchOptions>,
3074 control: QueryControl,
3075 ) -> Result<crate::ingest::SearchOutcome, QueryError> {
3076 self.search_with_graph_bound_mode(
3077 request,
3078 k,
3079 options.into(),
3080 control,
3081 GraphBoundMode::Shared,
3082 )
3083 }
3084
3085 fn admit_lexical_query(&self) -> Result<AdmittedLexicalQuery<'_>, QueryError> {
3086 #[cfg(any(test, feature = "test-support"))]
3087 self.text_materialization_work
3088 .admissions
3089 .fetch_add(1, Ordering::Relaxed);
3090 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3091 let lock_started = std::time::Instant::now();
3092 let state = self
3093 .state
3094 .lock()
3095 .map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
3096 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3097 self.text_materialization_work
3098 .observe_lexical_lock(lock_started);
3099 match *state {
3100 StoreState::Open => {}
3101 StoreState::Closing => return Err(QueryError::Store(StoreError::Closing)),
3102 StoreState::Closed => return Err(QueryError::Store(StoreError::Closed)),
3103 }
3104 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3105 let lock_started = std::time::Instant::now();
3106 let active_guard = self.active.lock().map_err(|_| {
3107 QueryError::Store(StoreError::Synchronization {
3108 component: "active segment",
3109 })
3110 })?;
3111 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3112 self.text_materialization_work
3113 .observe_lexical_lock(lock_started);
3114 let active_state = active_guard
3115 .as_ref()
3116 .ok_or(QueryError::Store(StoreError::Closed))?;
3117 let generation = active_state.generation;
3118 let active = Arc::clone(&active_state.segment);
3119 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3120 let lock_started = std::time::Instant::now();
3121 let snapshot = self
3122 .snapshot
3123 .read()
3124 .map_err(|_| {
3125 QueryError::Store(StoreError::Synchronization {
3126 component: "published snapshot",
3127 })
3128 })
3129 .inspect(|_| {
3130 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3131 self.text_materialization_work
3132 .observe_lexical_lock(lock_started);
3133 })?
3134 .as_ref()
3135 .cloned()
3136 .ok_or(QueryError::Store(StoreError::Closed))?;
3137 self.active_queries.fetch_add(1, Ordering::Relaxed);
3138 let active_query = ActiveQuery {
3139 count: &self.active_queries,
3140 };
3141 drop(active_guard);
3142 drop(state);
3143 Ok(AdmittedLexicalQuery {
3144 generation,
3145 active,
3146 snapshot: snapshot::ReadSnapshot::new(snapshot),
3147 active_query,
3148 })
3149 }
3150
3151 pub fn search_lexical(
3153 &self,
3154 query: &crate::fts::search::TermQuery,
3155 k: usize,
3156 control: QueryControl,
3157 ) -> Result<crate::ingest::StoreLexicalSearchOutcome, crate::ingest::StoreLexicalError> {
3158 self.search_lexical_then(query, k, control, |outcome, _, _, _, _, _| Ok(outcome))
3159 }
3160
3161 fn search_lexical_then<R>(
3162 &self,
3163 query: &crate::fts::search::TermQuery,
3164 k: usize,
3165 control: QueryControl,
3166 finish: impl FnOnce(
3167 crate::ingest::StoreLexicalSearchOutcome,
3168 &PublishedSnapshot,
3169 &crate::ingest::ActiveSegment,
3170 &LexicalAssembly,
3171 &[crate::fts::search::ScoredDoc],
3172 &QueryCancellation<'_>,
3173 ) -> Result<R, crate::ingest::StoreLexicalError>,
3174 ) -> Result<R, crate::ingest::StoreLexicalError> {
3175 let control = control.with_clock(Arc::clone(&self.clock));
3176 let started = self.clock.now();
3177 let admission_started = timing_start(self.clock.as_ref());
3178 let AdmittedLexicalQuery {
3179 generation,
3180 active,
3181 snapshot,
3182 active_query,
3183 } = self.admit_lexical_query()?;
3184 let admission_time = timing_elapsed(self.clock.as_ref(), admission_started);
3185 let lexical_started = timing_start(self.clock.as_ref());
3186
3187 let lease = SnapshotLease::new_at(Arc::clone(&snapshot), generation);
3188 let cancellation = QueryCancellation::new(&control, &lease);
3189 cancellation.check_graph().map_err(QueryError::Scan)?;
3190 let assembly = assemble_lexical_index(
3191 LexicalInputs {
3192 generation,
3193 cache: &self.lexical_index_cache,
3194 snapshot: &snapshot,
3195 active: &active,
3196 accounting: &self.accounting,
3197 },
3198 true,
3199 Some(&cancellation),
3200 )
3201 .map_err(map_store_lexical_assembly_error)?;
3202 let mut work = crate::fts::control::WorkCheck::new(|| {
3203 cancellation
3204 .check_graph()
3205 .map_err(QueryError::Scan)
3206 .map_err(crate::ingest::StoreLexicalError::from)
3207 });
3208 work.check_now()?;
3209 let mut allow_lists = Vec::with_capacity(assembly.alive_sets.len());
3210 for alive in &assembly.alive_sets {
3211 work.step()?;
3212 allow_lists.push(alive.alive_bitmap());
3213 }
3214 let lexical = crate::planner::search_lexical_filtered_refs_controlled(
3215 &assembly.index,
3216 query,
3217 k,
3218 crate::fts::bm25::Bm25Params::beir(),
3219 &allow_lists,
3220 None,
3221 || {
3222 cancellation
3223 .check_graph()
3224 .map_err(QueryError::Scan)
3225 .map_err(crate::ingest::StoreLexicalError::from)
3226 },
3227 )?;
3228 cancellation.check_graph().map_err(QueryError::Scan)?;
3229 let mut candidates = Vec::with_capacity(lexical.result.hits.len());
3230 for hit in &lexical.result.hits {
3231 work.step()?;
3232 let document =
3233 structured_lexical_document(&snapshot, &active, &assembly.sources, hit.doc, true)
3234 .map_err(map_store_lexical_document_error)?
3235 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
3236 candidates.push(crate::ingest::LexicalCandidate {
3237 document,
3238 score: hit.score,
3239 });
3240 }
3241 let mut diagnostics =
3242 crate::diag::QueryDiagnostics::lexical(crate::diag::LexicalDiagnostics {
3243 snapshot_generation: generation,
3244 indexed_through_seq: active
3245 .indexed_through_seq()
3246 .max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
3247 requested_k: k,
3248 returned: candidates.len(),
3249 counters: lexical.result.counters,
3250 elapsed: self.clock.now().saturating_duration_since(started),
3251 });
3252 diagnostics.counters.lexical_cache_hits = usize::from(assembly.cache_hit);
3253 diagnostics.counters.lexical_cache_builds = usize::from(!assembly.cache_hit);
3254 diagnostics.timings =
3255 crate::diag::QUERY_TIMING_ENABLED.then(|| crate::diag::QueryTimings {
3256 admission: admission_time,
3257 lexical: timing_elapsed(self.clock.as_ref(), lexical_started),
3258 ..Default::default()
3259 });
3260 let outcome = crate::ingest::StoreLexicalSearchOutcome {
3261 candidates,
3262 generation,
3263 diagnostics,
3264 };
3265 work.check_now()?;
3266 let result = finish(
3267 outcome,
3268 &snapshot,
3269 &active,
3270 &assembly,
3271 &lexical.result.hits,
3272 &cancellation,
3273 );
3274 drop(active_query);
3275 result
3276 }
3277
3278 pub fn stored_text(
3284 &self,
3285 version: crate::ingest::DocumentVersion,
3286 ) -> Result<Option<String>, crate::ingest::StoreLexicalError> {
3287 let AdmittedLexicalQuery {
3288 active,
3289 snapshot,
3290 active_query,
3291 ..
3292 } = self.admit_lexical_query()?;
3293 #[cfg(any(test, feature = "test-support"))]
3294 self.text_materialization_work
3295 .reverse_version_lookups
3296 .fetch_add(1, Ordering::Relaxed);
3297 if let Some((row, active_version, _)) = active.existing(version.doc_id())
3298 && active_version == version
3299 {
3300 let text = active
3301 .text(row)
3302 .map_err(crate::ingest::StoreLexicalError::from)?;
3303 drop(active_query);
3304 return Ok(text.map(str::to_owned));
3305 }
3306 for segment in snapshot.segments() {
3307 #[cfg(any(test, feature = "test-support"))]
3308 self.text_materialization_work
3309 .reverse_version_lookups
3310 .fetch_add(1, Ordering::Relaxed);
3311 let Some(row) = segment
3312 .query_row_for_document_version(version)
3313 .map_err(crate::ingest::StoreLexicalError::from)?
3314 else {
3315 continue;
3316 };
3317 let text = segment
3318 .query_stored_text()
3319 .map_err(StoreError::Segment)
3320 .map_err(crate::ingest::StoreLexicalError::from)?
3321 .and_then(|rows| rows.row(row).flatten())
3322 .map(str::to_owned);
3323 drop(active_query);
3324 return Ok(text);
3325 }
3326 drop(active_query);
3327 Ok(None)
3328 }
3329
3330 pub fn search_lexical_structured(
3333 &self,
3334 query: &crate::fts::query::LexicalQuery,
3335 k: usize,
3336 snippet_bytes: usize,
3337 control: QueryControl,
3338 ) -> Result<crate::ingest::StoreStructuredLexicalSearchOutcome, crate::ingest::StoreLexicalError>
3339 {
3340 let control = control.with_clock(Arc::clone(&self.clock));
3341 let started = self.clock.now();
3342 let admission_started = timing_start(self.clock.as_ref());
3343 let AdmittedLexicalQuery {
3344 generation,
3345 active,
3346 snapshot,
3347 active_query,
3348 } = self.admit_lexical_query()?;
3349 let admission_time = timing_elapsed(self.clock.as_ref(), admission_started);
3350 let lexical_started = timing_start(self.clock.as_ref());
3351
3352 let lease = SnapshotLease::new_at(Arc::clone(&snapshot), generation);
3353 let cancellation = QueryCancellation::new(&control, &lease);
3354 cancellation.check_graph().map_err(QueryError::Scan)?;
3355 let assembly = assemble_lexical_index(
3356 LexicalInputs {
3357 generation,
3358 cache: &self.lexical_index_cache,
3359 snapshot: &snapshot,
3360 active: &active,
3361 accounting: &self.accounting,
3362 },
3363 true,
3364 Some(&cancellation),
3365 )
3366 .map_err(map_store_lexical_assembly_error)?;
3367 let index = &assembly.index;
3368 let sources = &assembly.sources;
3369 let mut expansions = assembly
3370 .expand(query, &self.accounting, &cancellation)
3371 .map_err(|error| match error {
3372 expansion::ExpansionError::Query(error) => {
3373 crate::ingest::StoreLexicalError::Query(error)
3374 }
3375 expansion::ExpansionError::Shape(error) => {
3376 crate::ingest::StoreLexicalError::Structured(error)
3377 }
3378 })?;
3379 let allow_lists = assembly
3380 .alive_sets
3381 .iter()
3382 .map(|alive| alive.alive_bitmap())
3383 .collect::<Vec<_>>();
3384 let mut prepared_memory =
3385 stats::AccountedCounter::new(&self.accounting, stats::AllocationComponent::Temporary)
3386 .map_err(QueryError::Store)?;
3387 let mut work = crate::fts::control::WorkCheck::new(|| cancellation.check_graph());
3388 let fields = query
3389 .fields_controlled(&mut work)
3390 .map_err(QueryError::Scan)?;
3391 prepared_memory
3392 .set(
3393 crate::fts::query::PreparedWeightedQuery::allocation_bytes_controlled(
3394 &expansions,
3395 &fields,
3396 &mut work,
3397 )
3398 .map_err(QueryError::Scan)?
3399 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?,
3400 )
3401 .map_err(QueryError::Store)?;
3402 let mut prepared = if index.segments().is_empty() || expansions.is_empty() {
3403 None
3404 } else {
3405 Some(
3406 crate::fts::query::PreparedWeightedQuery::new_controlled(
3407 index,
3408 std::mem::take(&mut expansions),
3409 fields,
3410 || cancellation.check_graph(),
3411 )
3412 .map_err(|error| match error {
3413 crate::fts::search::ControlledSearchError::Index(error) => {
3414 crate::ingest::StoreLexicalError::from(
3415 crate::planner::LexicalFilterError::from(error),
3416 )
3417 }
3418 crate::fts::search::ControlledSearchError::Control(error) => {
3419 crate::ingest::StoreLexicalError::from(QueryError::Scan(error))
3420 }
3421 })?,
3422 )
3423 };
3424 let mut scratch_memory =
3425 stats::AccountedCounter::new(&self.accounting, stats::AllocationComponent::Temporary)
3426 .map_err(QueryError::Store)?;
3427 let mut phrase = prepared_lexical::PhraseEligibility::new(
3428 index,
3429 Some(query),
3430 &self.accounting,
3431 &cancellation,
3432 )?;
3433 let result = match &prepared {
3434 None => crate::fts::prune::weighted::WeightedResult::default(),
3435 Some(prepared) => crate::fts::prune::weighted::search(
3436 index,
3437 prepared,
3438 k,
3439 &allow_lists,
3440 || {
3441 cancellation
3442 .check_graph()
3443 .map_err(QueryError::Scan)
3444 .map_err(crate::ingest::StoreLexicalError::from)
3445 },
3446 |doc| {
3447 phrase
3448 .matches(doc, &cancellation)
3449 .map_err(crate::ingest::StoreLexicalError::from)
3450 },
3451 |bytes| {
3452 prepared_lexical::reserve_weighted_scratch(&mut scratch_memory, bytes)
3453 .map_err(crate::ingest::StoreLexicalError::from)
3454 },
3455 )
3456 .map_err(|error| match error {
3457 crate::fts::search::ControlledSearchError::Index(error) => {
3458 crate::ingest::StoreLexicalError::from(
3459 crate::planner::LexicalFilterError::from(error),
3460 )
3461 }
3462 crate::fts::search::ControlledSearchError::Control(error) => error,
3463 })?,
3464 };
3465 let mut counters = result.counters;
3466 let scored = result.hits;
3467 let mut provenance_memory =
3468 stats::AccountedCounter::new(&self.accounting, stats::AllocationComponent::Temporary)
3469 .map_err(QueryError::Store)?;
3470 let mut candidates = Vec::with_capacity(scored.len());
3471 let mut provenance_work = crate::fts::control::WorkCheck::new(|| {
3472 cancellation.check_graph().map_err(QueryError::Scan)
3473 });
3474 for (doc, score) in scored {
3475 let mut provenance = Vec::new();
3476 let segment = index
3477 .segments()
3478 .get(doc.segment as usize)
3479 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
3480 let scoring = prepared
3481 .as_ref()
3482 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
3483 for expansion in scoring.expansions() {
3484 cancellation.check_graph().map_err(QueryError::Scan)?;
3485 let fields = &scoring.scoring().query().fields;
3486 prepared_lexical::reserve_weighted_scratch(
3487 &mut provenance_memory,
3488 crate::fts::sealed::TermStream::allocation_bytes(
3489 segment,
3490 &expansion.term,
3491 fields,
3492 ),
3493 )?;
3494 if let Some(mut stream) = crate::fts::sealed::TermStream::open_sized_controlled(
3495 segment,
3496 &expansion.term,
3497 fields,
3498 &mut provenance_work,
3499 )? {
3500 stream.seek_controlled(doc.row, &mut provenance_work)?;
3501 if stream.current_row() == Some(doc.row) {
3502 provenance.push(expansion.clone());
3503 counters.postings_decoded = counters.postings_decoded.saturating_add(1);
3504 }
3505 counters.blocks_decoded = counters
3506 .blocks_decoded
3507 .saturating_add(stream.blocks_decoded());
3508 counters.blocks_skipped = counters
3509 .blocks_skipped
3510 .saturating_add(stream.blocks_skipped());
3511 }
3512 provenance_memory.set(0).map_err(QueryError::Store)?;
3513 }
3514 cancellation.check_graph().map_err(QueryError::Scan)?;
3515 let (text, document) = structured_lexical_row(&snapshot, &active, sources, doc)?;
3516 let terms = provenance
3517 .iter()
3518 .map(|entry| entry.term.clone())
3519 .collect::<Vec<_>>();
3520 let snippet = crate::fts::snippet::best_window(
3521 &self.tokenizer,
3522 text,
3523 &terms,
3524 snippet_bytes,
3525 true,
3526 )?
3527 .ok_or(crate::ingest::StoreLexicalError::MissingSnippetMatch {
3528 segment: doc.segment,
3529 row: doc.row,
3530 })?;
3531 let snippet_text = snippet
3532 .text(text)
3533 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
3534 candidates.push(crate::ingest::ExplainedLexicalCandidate {
3535 document,
3536 score,
3537 provenance,
3538 snippet: crate::fts::query::OwnedLexicalSnippet {
3539 text: snippet_text.to_owned(),
3540 source: snippet.window,
3541 highlights: snippet.highlights,
3542 },
3543 });
3544 }
3545 cancellation.check_graph().map_err(QueryError::Scan)?;
3546 let mut diagnostics =
3547 crate::diag::QueryDiagnostics::lexical(crate::diag::LexicalDiagnostics {
3548 snapshot_generation: generation,
3549 indexed_through_seq: active
3550 .indexed_through_seq()
3551 .max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
3552 requested_k: k,
3553 returned: candidates.len(),
3554 counters,
3555 elapsed: self.clock.now().saturating_duration_since(started),
3556 });
3557 diagnostics.counters.lexical_cache_hits = usize::from(assembly.cache_hit);
3558 diagnostics.counters.lexical_cache_builds = usize::from(!assembly.cache_hit);
3559 diagnostics.timings =
3560 crate::diag::QUERY_TIMING_ENABLED.then(|| crate::diag::QueryTimings {
3561 admission: admission_time,
3562 lexical: timing_elapsed(self.clock.as_ref(), lexical_started),
3563 ..Default::default()
3564 });
3565 drop(active_query);
3566 Ok(crate::ingest::StoreStructuredLexicalSearchOutcome {
3567 candidates,
3568 expansions: prepared
3569 .as_mut()
3570 .map_or(expansions, |query| query.take_expansions()),
3571 generation,
3572 diagnostics,
3573 })
3574 }
3575
3576 pub fn search_hybrid(
3579 &self,
3580 vector_query: crate::ingest::SearchRequest<'_>,
3581 lexical_query: &crate::fts::search::TermQuery,
3582 hybrid_query: &crate::fusion::HybridQuery,
3583 options: impl Into<SearchOptions>,
3584 control: QueryControl,
3585 ) -> Result<crate::ingest::StoreHybridSearchOutcome, crate::fusion::FusionError> {
3586 self.search_hybrid_inner(
3587 vector_query,
3588 PinnedLexicalQuery::Term(lexical_query),
3589 hybrid_query,
3590 options.into(),
3591 control,
3592 )
3593 }
3594
3595 pub fn search_hybrid_structured(
3598 &self,
3599 vector_query: crate::ingest::SearchRequest<'_>,
3600 lexical_query: &crate::fts::query::LexicalQuery,
3601 hybrid_query: &crate::fusion::HybridQuery,
3602 options: impl Into<SearchOptions>,
3603 control: QueryControl,
3604 ) -> Result<crate::ingest::StoreHybridSearchOutcome, crate::fusion::FusionError> {
3605 self.search_hybrid_inner(
3606 vector_query,
3607 PinnedLexicalQuery::Structured(lexical_query),
3608 hybrid_query,
3609 options.into(),
3610 control,
3611 )
3612 }
3613
3614 fn search_hybrid_inner(
3615 &self,
3616 vector_query: crate::ingest::SearchRequest<'_>,
3617 lexical_query: PinnedLexicalQuery<'_>,
3618 hybrid_query: &crate::fusion::HybridQuery,
3619 options: SearchOptions,
3620 control: QueryControl,
3621 ) -> Result<crate::ingest::StoreHybridSearchOutcome, crate::fusion::FusionError> {
3622 self.search_hybrid_then(
3623 vector_query,
3624 lexical_query,
3625 hybrid_query,
3626 options,
3627 control,
3628 false,
3629 |outcome, _, _, _, _| Ok(outcome),
3630 )
3631 }
3632
3633 #[allow(clippy::too_many_arguments)] fn search_hybrid_then<R>(
3635 &self,
3636 vector_query: crate::ingest::SearchRequest<'_>,
3637 lexical_query: PinnedLexicalQuery<'_>,
3638 hybrid_query: &crate::fusion::HybridQuery,
3639 options: SearchOptions,
3640 control: QueryControl,
3641 capture_rows: bool,
3642 finish: impl FnOnce(
3643 crate::ingest::StoreHybridSearchOutcome,
3644 &PublishedSnapshot,
3645 &crate::ingest::ActiveSegment,
3646 &QueryCancellation<'_>,
3647 Option<&materialize::HybridAddresses>,
3648 ) -> Result<R, crate::fusion::FusionError>,
3649 ) -> Result<R, crate::fusion::FusionError> {
3650 self.search_hybrid_prepared_then(
3651 || Ok::<_, std::convert::Infallible>(vector_query),
3652 lexical_query,
3653 hybrid_query,
3654 options,
3655 control,
3656 capture_rows,
3657 finish,
3658 )
3659 .map_err(|error| match error {
3660 HybridPreparationError::Preparation(never) => match never {},
3661 HybridPreparationError::Search(error) => error,
3662 })
3663 }
3664
3665 #[allow(clippy::too_many_arguments)] fn search_hybrid_prepared_then<'vector, E, R>(
3667 &self,
3668 prepare_vector: impl FnOnce() -> Result<crate::ingest::SearchRequest<'vector>, E>,
3669 lexical_query: PinnedLexicalQuery<'_>,
3670 hybrid_query: &crate::fusion::HybridQuery,
3671 options: SearchOptions,
3672 control: QueryControl,
3673 capture_rows: bool,
3674 finish: impl FnOnce(
3675 crate::ingest::StoreHybridSearchOutcome,
3676 &PublishedSnapshot,
3677 &crate::ingest::ActiveSegment,
3678 &QueryCancellation<'_>,
3679 Option<&materialize::HybridAddresses>,
3680 ) -> Result<R, crate::fusion::FusionError>,
3681 ) -> Result<R, HybridPreparationError<E>> {
3682 let control = control.with_clock(Arc::clone(&self.clock));
3683 let started = self.clock.now();
3684 let admission_started = timing_start(self.clock.as_ref());
3685 let mut options = options;
3686 let requested_tier = options.explicit_tier();
3687 let admitted = self
3703 .admit_vector_search_for(options, requested_tier.is_none())
3704 .map_err(crate::fusion::FusionError::from)?;
3705 options = admitted.execution_options;
3706 let tier_resolution = if requested_tier.is_none() {
3707 if snapshot_has_graph(&admitted.snapshot) {
3708 Some(crate::diag::HybridTierResolution::Auto)
3709 } else {
3710 options = options.with_tier(SearchTier::Exact);
3711 Some(crate::diag::HybridTierResolution::Exact)
3712 }
3713 } else {
3714 None
3715 };
3716 let corpus_rows = hybrid::corpus_rows(&admitted.snapshot, &admitted.active_segment)?;
3717 let mut width = if hybrid_query.max_rounds == 0 {
3722 corpus_rows
3723 } else {
3724 hybrid::hybrid_window(hybrid_query.k, corpus_rows)?.width
3725 };
3726 let panic_vector = self.consume_hybrid_test_fault(crate::fusion::FusionLeg::Vector);
3727 let panic_lexical = self.consume_hybrid_test_fault(crate::fusion::FusionLeg::Lexical);
3728 let lexical_worker = self.ensure_lexical_worker()?;
3729 let mut timings = crate::diag::QueryTimings {
3730 admission: timing_elapsed(self.clock.as_ref(), admission_started),
3731 ..Default::default()
3732 };
3733 let lexical_inputs = LexicalInputs {
3734 generation: admitted.generation,
3735 cache: &self.lexical_index_cache,
3736 snapshot: &admitted.snapshot,
3737 active: &admitted.active_segment,
3738 accounting: &self.accounting,
3739 };
3740 #[cfg(any(test, feature = "test-support"))]
3741 let preparation_observer = crate::fts::preparation_observer::current();
3742 let mut lexical_preparation = prepared_lexical::PreparedLexicalQuery::new(lexical_query);
3743 #[cfg(any(test, feature = "test-support"))]
3744 let fresh_preparation = hybrid::fresh_round_test_control();
3745 let run_lexical_leg = |preparation: &mut prepared_lexical::PreparedLexicalQuery<'_>,
3746 bound,
3747 queued| {
3748 #[cfg(any(test, feature = "test-support"))]
3749 if fresh_preparation {
3750 preparation.reset_for_test();
3751 }
3752 #[cfg(any(test, feature = "test-support"))]
3753 let _preparation_scope =
3754 crate::fts::preparation_observer::enter(preparation_observer.clone());
3755 let queue_time = timing_elapsed(self.clock.as_ref(), queued);
3756 let lexical_started = timing_start(self.clock.as_ref());
3757 let name = std::thread::current().name().map(str::to_owned);
3758 let result = (|| {
3759 maybe_trigger_hybrid_leg_panic(panic_lexical, "injected lexical hybrid leg panic");
3760 let lease =
3761 SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
3762 let cancellation = QueryCancellation::new(&control, &lease);
3763 cancellation
3764 .check_graph()
3765 .map_err(QueryError::Scan)
3766 .map_err(crate::fusion::FusionError::from)?;
3767 match lexical_query {
3768 PinnedLexicalQuery::Term(_) => {
3769 exact_lexical_leg(lexical_inputs, bound, &cancellation, preparation)
3770 }
3771 PinnedLexicalQuery::Structured(query) => exact_structured_lexical_leg(
3772 lexical_inputs,
3773 &self.tokenizer,
3774 query,
3775 bound,
3776 &cancellation,
3777 preparation,
3778 ),
3779 }
3780 })();
3781 (
3782 name,
3783 result,
3784 queue_time,
3785 timing_elapsed(self.clock.as_ref(), lexical_started),
3786 )
3787 };
3788 let caller_chose_tier = requested_tier.is_some();
3789 let run_vector_leg = |vector_query: crate::ingest::SearchRequest<'_>,
3790 vector_preparation: &mut prepared::PreparedVectorQuery<'_>,
3791 k: usize| {
3792 let graph_available = snapshot_has_graph(&admitted.snapshot);
3793 let scan_reason_override = if caller_chose_tier || !graph_available {
3794 None
3795 } else if k > corpus_rows {
3796 Some(crate::planner::ScanReason::FullMaterialization)
3797 } else if graph_round_is_worth_it(k, corpus_rows) {
3798 None
3799 } else {
3800 Some(crate::planner::ScanReason::WideningCap {
3801 ef: graph_round_ef(k),
3802 rows: corpus_rows,
3803 })
3804 };
3805 search_pinned(
3806 admitted.pool.as_deref(),
3807 &admitted.snapshot,
3808 &admitted.active_segment,
3809 &self.accounting,
3810 admitted.generation,
3811 self.epoch_identity(),
3812 vector_query,
3813 k,
3814 if caller_chose_tier || graph_round_is_worth_it(k, corpus_rows) {
3826 options
3827 } else {
3828 options.with_tier(SearchTier::Exact)
3829 },
3830 control.clone(),
3831 GraphBoundMode::Shared,
3832 requested_tier,
3833 scan_reason_override,
3834 started,
3835 self.clock.as_ref(),
3836 vector_preparation,
3837 #[cfg(any(test, feature = "test-support"))]
3838 self.vector_fault_controller.as_ref(),
3839 )
3840 .map_err(crate::fusion::FusionError::from)
3841 };
3842 let lexical_queued = timing_start(self.clock.as_ref());
3843 let (prepared_vector_result, lexical_result) = lexical_worker.run_scoped(
3844 || {
3845 run_lexical_leg(
3846 &mut lexical_preparation,
3847 width.saturating_add(1),
3848 lexical_queued,
3849 )
3850 },
3851 || {
3852 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3853 let vector_query =
3854 prepare_vector().map_err(HybridPreparationError::Preparation)?;
3855 let mut vector_preparation = prepared::PreparedVectorQuery::new(
3856 vector_query.vector(),
3857 self.epoch_identity(),
3858 &self.accounting,
3859 );
3860 maybe_trigger_hybrid_leg_panic(
3861 panic_vector,
3862 "injected vector hybrid leg panic",
3863 );
3864 let result = run_vector_leg(
3865 vector_query,
3866 &mut vector_preparation,
3867 width.saturating_add(1),
3868 );
3869 Ok((vector_query, vector_preparation, result))
3870 }))
3871 .unwrap_or(Err(HybridPreparationError::Search(
3872 crate::fusion::FusionError::LegPanic {
3873 leg: crate::fusion::FusionLeg::Vector,
3874 detail: "vector hybrid leg panicked",
3875 },
3876 )))
3877 },
3878 )?;
3879 let (lexical_thread_name, lexical_result, queue_time, lexical_time) = lexical_result
3880 .unwrap_or_else(|error| {
3881 (
3882 Some("zeppelin-fts".to_owned()),
3883 Err(error),
3884 std::time::Duration::ZERO,
3885 std::time::Duration::ZERO,
3886 )
3887 });
3888 #[cfg(any(test, feature = "test-support"))]
3889 if let Ok(mut receipt) = self.hybrid_execution_receipt.lock() {
3890 *receipt = Some(HybridExecutionReceipt {
3891 vector_thread: std::thread::current().id(),
3892 lexical_thread_name,
3893 vector_completed: true,
3894 lexical_completed: true,
3895 });
3896 }
3897 #[cfg(not(any(test, feature = "test-support")))]
3898 let _ = lexical_thread_name;
3899 let (vector_query, mut vector_preparation, vector_result) = match prepared_vector_result {
3903 Ok(prepared) => prepared,
3904 Err(error) => {
3905 let lease =
3906 SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
3907 QueryCancellation::new(&control, &lease)
3908 .check_graph()
3909 .map_err(map_scan_error)?;
3910 if let Err(lexical_error) = lexical_result
3911 && matches!(
3912 lexical_error,
3913 crate::fusion::FusionError::ReadCancelled { .. }
3914 | crate::fusion::FusionError::Timeout { .. }
3915 | crate::fusion::FusionError::Cancelled { .. }
3916 )
3917 {
3918 return Err(HybridPreparationError::Search(lexical_error));
3919 }
3920 return Err(error);
3921 }
3922 };
3923 let (
3924 mut vector_outcome,
3925 (mut lexical_hits, mut lexical_assembly, mut lexical_counters, mut lexical_cache_hit),
3926 ) = resolve_hybrid_leg_results(vector_result, lexical_result)?;
3927
3928 timings.lexical_queue = queue_time;
3929 timings.lexical = lexical_time;
3930 let anchors = hybrid::HybridAnchors {
3931 vector_ceiling: vector_outcome.vector_ceiling,
3932 lexical_maximum: lexical_hits.first().map_or(0.0, |hit| hit.bm25),
3933 };
3934 let fusion_lease =
3935 SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
3936 let fusion_cancellation = QueryCancellation::new(&control, &fusion_lease);
3937 let mut work = hybrid::HybridWork::new();
3938 let mut cross_scores = hybrid::CrossScoreCache::new(vector_outcome.epoch);
3939 let mut round_buffers = hybrid::HybridRoundBuffers::new();
3940 let mut fusion_scratch =
3941 crate::fusion::StoreFusionScratch::new(Arc::clone(&self.accounting));
3942 let mut rescore_counters = crate::diag::ScanRescoreCounters::default();
3943 let mut rounds = 1_usize;
3944 let mut budget_exhausted = false;
3945 let (fused, report_window, addresses) = loop {
3946 if let Some(vector_timings) = vector_outcome.diagnostics.timings {
3947 timings.vector += vector_timings.vector;
3948 }
3949 let fusion_started = timing_start(self.clock.as_ref());
3950 let round = hybrid::build_round(
3951 &admitted.snapshot,
3952 &admitted.active_segment,
3953 &lexical_assembly,
3954 &lexical_preparation,
3955 &self.tokenizer,
3956 lexical_query,
3957 vector_query.vector(),
3958 &vector_outcome.candidates,
3959 anchors,
3960 &lexical_hits,
3961 width,
3962 &fusion_cancellation,
3963 capture_rows,
3964 &self.accounting,
3965 &mut cross_scores,
3966 &mut round_buffers,
3967 )?;
3968 let cross_fill_dims = u64::try_from(round.vector_scores_computed)
3972 .ok()
3973 .and_then(|rows| rows.checked_mul(vector_query.vector().len() as u64))
3974 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
3975 let cross_fill_bytes = cross_fill_dims
3976 .checked_mul(4)
3977 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
3978 vector_outcome.stats.dims_touched = vector_outcome
3979 .stats
3980 .dims_touched
3981 .checked_add(cross_fill_dims)
3982 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
3983 vector_outcome.stats.bytes_read = vector_outcome
3984 .stats
3985 .bytes_read
3986 .checked_add(cross_fill_bytes)
3987 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
3988 if cross_fill_dims != 0 {
3989 let caller = std::thread::current().id();
3990 if !vector_outcome.stats.worker_thread_ids.contains(&caller) {
3991 vector_outcome.stats.worker_thread_ids.push(caller);
3992 }
3993 vector_outcome.stats.threads_used = vector_outcome.stats.worker_thread_ids.len();
3994 }
3995 accumulate_search_counters(&mut lexical_counters, &round.lexical_counters);
3996 if let Some(counters) = vector_outcome.diagnostics.scan_rescore {
3997 rescored_scan::accumulate(&mut rescore_counters, counters)?;
3998 }
3999 work.complete_round(
4000 &vector_outcome.stats,
4001 vector_outcome.graph_stats,
4002 lexical_counters,
4003 )?;
4004 work.lexical_cache_hits += usize::from(lexical_cache_hit);
4005 work.lexical_cache_builds += usize::from(!lexical_cache_hit);
4006 work.plans.append(&mut vector_outcome.diagnostics.plan);
4007 work.cross_filled_vector = work
4008 .cross_filled_vector
4009 .checked_add(round.vector_scores_computed)
4010 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4011 work.cross_filled_lexical = work
4012 .cross_filled_lexical
4013 .checked_add(round.lexical_scores_computed)
4014 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4015 work.vector_candidates_produced = work
4016 .vector_candidates_produced
4017 .checked_add(vector_outcome.candidates.len())
4018 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4019 work.lexical_candidates_produced = work
4020 .lexical_candidates_produced
4021 .checked_add(lexical_hits.len())
4022 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4023 let provenance = hybrid::round_provenance(&vector_outcome, lexical_hits.len(), width);
4024 let fused = hybrid::fuse_round(
4025 hybrid_query,
4026 &round,
4027 provenance,
4028 requested_tier == Some(SearchTier::Exact),
4029 &mut fusion_scratch,
4030 )?;
4031 timings.fusion_cross_fill += timing_elapsed(self.clock.as_ref(), fusion_started);
4032 if fused.report.termination != crate::fusion::FusionTermination::WindowUnproven
4033 || width >= corpus_rows
4034 {
4035 break (
4036 fused,
4037 crate::diag::HybridReport {
4038 normalization_policy_version:
4039 crate::fusion::HYBRID_NORMALIZATION_POLICY_VERSION,
4040 lexical_full_materializations: 0,
4041 provenance,
4042 window: width,
4043 vector_returned: round.vector.len(),
4044 lexical_returned: round.lexical.len(),
4045 cross_filled_vector: round.cross_filled_vector,
4046 cross_filled_lexical: round.cross_filled_lexical,
4047 total_cross_filled_vector: work.cross_filled_vector,
4048 total_cross_filled_lexical: work.cross_filled_lexical,
4049 vector_candidates_produced: work.vector_candidates_produced,
4050 lexical_candidates_produced: work.lexical_candidates_produced,
4051 },
4052 round.addresses,
4053 );
4054 }
4055 width = if rounds >= hybrid_query.max_rounds {
4056 budget_exhausted = true;
4057 corpus_rows
4058 } else {
4059 width.saturating_mul(2).min(corpus_rows)
4060 };
4061 rounds = rounds.saturating_add(1);
4062 let lexical_queued = timing_start(self.clock.as_ref());
4063 let (vector_result, lexical_result) = lexical_worker.run_scoped(
4064 || {
4065 run_lexical_leg(
4066 &mut lexical_preparation,
4067 width.saturating_add(1),
4068 lexical_queued,
4069 )
4070 },
4071 || {
4072 run_vector_leg(
4073 vector_query,
4074 &mut vector_preparation,
4075 width.saturating_add(1),
4076 )
4077 },
4078 )?;
4079 let (_, lexical_result, queue_time, lexical_time) =
4080 lexical_result.unwrap_or_else(|error| {
4081 (
4082 Some("zeppelin-fts".to_owned()),
4083 Err(error),
4084 std::time::Duration::ZERO,
4085 std::time::Duration::ZERO,
4086 )
4087 });
4088 let (next_vector, (next_hits, next_sources, next_counters, next_cache_hit)) =
4089 resolve_hybrid_leg_results(vector_result, lexical_result)?;
4090 timings.lexical_queue += queue_time;
4091 timings.lexical += lexical_time;
4092 vector_outcome = next_vector;
4093 lexical_hits = next_hits;
4094 lexical_assembly = next_sources;
4095 lexical_counters = next_counters;
4096 lexical_cache_hit = next_cache_hit;
4097 };
4098
4099 let crate::ingest::SearchOutcome {
4100 candidates: _,
4101 vector_ceiling: _,
4102 stats: _,
4103 graph_stats: _,
4104 generation,
4105 epoch,
4106 diagnostics: vector_diagnostics,
4107 } = vector_outcome;
4108 let mut report = fused.report;
4109 report.rounds = rounds;
4110 if budget_exhausted {
4111 report.budget_exhausted = true;
4112 report.termination = crate::fusion::FusionTermination::BudgetFullMaterialization;
4113 }
4114 let mut diagnostics =
4115 crate::diag::QueryDiagnostics::hybrid(crate::diag::HybridDiagnostics {
4116 snapshot_generation: vector_diagnostics.snapshot_generation,
4117 indexed_through_seq: vector_diagnostics.indexed_through_seq,
4118 plan: work.plans,
4119 approximate: vector_diagnostics.approximate
4120 || report.termination
4121 == crate::fusion::FusionTermination::ApproximateCandidates,
4122 exact_rescore: vector_diagnostics.exact_rescore,
4123 requested_k: hybrid_query.k,
4124 returned: fused.hits.len(),
4125 scan: work.scan,
4126 graph: work.graph,
4127 lexical: work.lexical,
4128 report,
4129 hybrid: report_window,
4130 tier_resolution,
4131 epoch,
4132 elapsed: self.clock.now().saturating_duration_since(started),
4133 });
4134 diagnostics.counters.lexical_cache_hits = work.lexical_cache_hits;
4135 diagnostics.scan_rescore = options.scan_rescore().map(|_| rescore_counters);
4136 diagnostics.counters.lexical_cache_builds = work.lexical_cache_builds;
4137 diagnostics.timings = crate::diag::QUERY_TIMING_ENABLED.then_some(timings);
4138 let outcome = crate::ingest::StoreHybridSearchOutcome {
4139 hits: fused.hits,
4140 lexical_expansions: lexical_preparation.take_expansions()?,
4141 generation,
4142 diagnostics,
4143 };
4144 finish(
4145 outcome,
4146 &admitted.snapshot,
4147 &admitted.active_segment,
4148 &fusion_cancellation,
4149 addresses.as_ref(),
4150 )
4151 .map_err(HybridPreparationError::Search)
4152 }
4153
4154 #[cfg(any(test, feature = "test-support"))]
4160 #[must_use]
4161 pub fn lexical_index_cache_counters(&self) -> (u64, u64) {
4162 (
4163 self.lexical_index_cache.hits.load(Ordering::Relaxed),
4164 self.lexical_index_cache.builds.load(Ordering::Relaxed),
4165 )
4166 }
4167
4168 #[cfg(any(test, feature = "test-support"))]
4170 #[must_use]
4171 pub fn lexical_contribution_cache_counters(&self) -> (u64, u64) {
4172 (
4173 self.lexical_index_cache
4174 .contribution_hits
4175 .load(Ordering::Relaxed),
4176 self.lexical_index_cache
4177 .contribution_builds
4178 .load(Ordering::Relaxed),
4179 )
4180 }
4181
4182 fn consume_hybrid_test_fault(&self, leg: crate::fusion::FusionLeg) -> bool {
4183 #[cfg(any(test, feature = "test-support"))]
4184 {
4185 let Ok(mut fault) = self.hybrid_leg_fault.lock() else {
4186 return false;
4187 };
4188 if matches!(*fault, Some(HybridLegTestFault::Panic(candidate)) if candidate == leg) {
4189 *fault = None;
4190 return true;
4191 }
4192 }
4193 #[cfg(not(any(test, feature = "test-support")))]
4194 let _ = leg;
4195 false
4196 }
4197
4198 fn search_with_graph_bound_mode(
4199 &self,
4200 request: crate::ingest::SearchRequest<'_>,
4201 k: usize,
4202 options: SearchOptions,
4203 control: QueryControl,
4204 graph_bound_mode: GraphBoundMode,
4205 ) -> Result<crate::ingest::SearchOutcome, QueryError> {
4206 self.search_with_graph_bound_mode_then(
4207 request,
4208 k,
4209 options,
4210 control,
4211 graph_bound_mode,
4212 |outcome, _, _| Ok(outcome),
4213 )
4214 }
4215
4216 fn search_with_graph_bound_mode_then<R>(
4217 &self,
4218 request: crate::ingest::SearchRequest<'_>,
4219 k: usize,
4220 options: SearchOptions,
4221 control: QueryControl,
4222 graph_bound_mode: GraphBoundMode,
4223 finish: impl FnOnce(
4224 crate::ingest::SearchOutcome,
4225 &AdmittedVectorSearch<'_>,
4226 &QueryControl,
4227 ) -> Result<R, QueryError>,
4228 ) -> Result<R, QueryError> {
4229 let control = control.with_clock(Arc::clone(&self.clock));
4230 let started = self.clock.now();
4231 let admission_started = timing_start(self.clock.as_ref());
4232 let admitted = self.admit_vector_search(options)?;
4233 let admission_time = timing_elapsed(self.clock.as_ref(), admission_started);
4234 let mut vector_preparation = prepared::PreparedVectorQuery::new(
4235 request.vector(),
4236 self.epoch_identity(),
4237 &self.accounting,
4238 );
4239 let score = || {
4240 let mut outcome = search_pinned(
4241 admitted.pool.as_deref(),
4242 &admitted.snapshot,
4243 &admitted.active_segment,
4244 &self.accounting,
4245 admitted.generation,
4246 self.epoch_identity(),
4247 request,
4248 k,
4249 options,
4250 control.clone(),
4251 graph_bound_mode,
4252 options.explicit_tier(),
4253 None,
4254 started,
4255 self.clock.as_ref(),
4256 &mut vector_preparation,
4257 #[cfg(any(test, feature = "test-support"))]
4258 self.vector_fault_controller.as_ref(),
4259 )?;
4260 if let Some(timings) = &mut outcome.diagnostics.timings {
4261 timings.admission = admission_time;
4262 }
4263 finish(outcome, &admitted, &control)
4264 };
4265 #[cfg(any(test, feature = "test-support"))]
4266 let result = crate::kernels::vector_fault::run_store_scoring(
4267 self.kernel_fault_controller.as_ref(),
4268 score,
4269 );
4270 #[cfg(not(any(test, feature = "test-support")))]
4271 let result = score();
4272 #[cfg(any(test, feature = "test-support"))]
4273 {
4274 if let Some(controller) = self.vector_fault_controller.as_ref() {
4275 controller.finalize_search(result.is_ok());
4276 }
4277 }
4278 result
4279 }
4280
4281 fn ensure_query_pool(&self) -> Result<Arc<pool::QueryPool>, QueryError> {
4282 let mut pool_slot = self.query_pool.lock().map_err(|_| {
4283 QueryError::Store(StoreError::Synchronization {
4284 component: "query pool",
4285 })
4286 })?;
4287 if pool_slot.is_none() {
4288 let capacity = crate::scan::physical_thread_capacity().map_err(|error| {
4289 QueryError::Store(StoreError::QueryPoolCapacity {
4290 source: error.to_string(),
4291 })
4292 })?;
4293 let pool =
4294 pool::QueryPool::start(capacity, &self.accounting).map_err(QueryError::Store)?;
4295 *pool_slot = Some(Arc::new(pool));
4296 }
4297 let pool = pool_slot.as_ref().cloned().ok_or({
4298 QueryError::Store(StoreError::Synchronization {
4299 component: "query pool initialization",
4300 })
4301 })?;
4302 drop(pool_slot);
4303 Ok(pool)
4304 }
4305
4306 fn ensure_lexical_worker(
4307 &self,
4308 ) -> Result<Arc<pool::LexicalWorker>, crate::fusion::FusionError> {
4309 let mut worker_slot =
4310 self.lexical_worker
4311 .lock()
4312 .map_err(|_| crate::fusion::FusionError::LegThreadStart {
4313 leg: crate::fusion::FusionLeg::Lexical,
4314 detail: "pooled lexical worker synchronization was poisoned".to_owned(),
4315 })?;
4316 if worker_slot.is_none() {
4317 *worker_slot = Some(Arc::new(pool::LexicalWorker::start()?));
4318 }
4319 let worker = worker_slot.as_ref().cloned().ok_or_else(|| {
4320 crate::fusion::FusionError::LegThreadStart {
4321 leg: crate::fusion::FusionLeg::Lexical,
4322 detail: "pooled lexical worker initialization failed".to_owned(),
4323 }
4324 })?;
4325 drop(worker_slot);
4326 Ok(worker)
4327 }
4328
4329 fn admit_vector_search(
4330 &self,
4331 options: SearchOptions,
4332 ) -> Result<AdmittedVectorSearch<'_>, QueryError> {
4333 self.admit_vector_search_for(options, false)
4334 }
4335
4336 fn admit_vector_search_for(
4337 &self,
4338 options: SearchOptions,
4339 hybrid_default: bool,
4340 ) -> Result<AdmittedVectorSearch<'_>, QueryError> {
4341 #[cfg(any(test, feature = "test-support"))]
4342 self.text_materialization_work
4343 .admissions
4344 .fetch_add(1, Ordering::Relaxed);
4345 let state = self
4346 .state
4347 .lock()
4348 .map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
4349 match *state {
4350 StoreState::Open => {}
4351 StoreState::Closing => return Err(QueryError::Store(StoreError::Closing)),
4352 StoreState::Closed => return Err(QueryError::Store(StoreError::Closed)),
4353 }
4354 let active_guard = self.active.lock().map_err(|_| {
4355 QueryError::Store(StoreError::Synchronization {
4356 component: "active segment",
4357 })
4358 })?;
4359 let active_state = active_guard
4360 .as_ref()
4361 .ok_or(QueryError::Store(StoreError::Closed))?;
4362 let generation = active_state.generation;
4363 let active_segment = Arc::clone(&active_state.segment);
4364 let snapshot = self
4365 .snapshot
4366 .read()
4367 .map_err(|_| {
4368 QueryError::Store(StoreError::Synchronization {
4369 component: "published snapshot",
4370 })
4371 })?
4372 .as_ref()
4373 .cloned()
4374 .ok_or(QueryError::Store(StoreError::Closed))?;
4375 let execution_options = if hybrid_default && !snapshot_has_graph(&snapshot) {
4376 options.with_tier(SearchTier::Exact)
4377 } else {
4378 options
4379 };
4380 let needs_query_pool = !active_segment.is_empty()
4383 || (hybrid_default && !snapshot.segments().is_empty())
4384 || match execution_options.tier() {
4385 SearchTier::Auto => snapshot.segments().iter().any(|segment| {
4386 !segment.directory().iter().any(|entry| {
4387 entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id()
4388 })
4389 }),
4390 SearchTier::Exact => !snapshot.segments().is_empty(),
4391 SearchTier::Scan => true,
4392 SearchTier::Graph(_) => false,
4393 };
4394 let pool = if needs_query_pool {
4395 Some(self.ensure_query_pool()?)
4396 } else {
4397 None
4398 };
4399 self.active_queries.fetch_add(1, Ordering::Relaxed);
4400 let active_query = ActiveQuery {
4401 count: &self.active_queries,
4402 };
4403 drop(active_guard);
4404 drop(state);
4405 Ok(AdmittedVectorSearch {
4406 pool,
4407 execution_options,
4408 snapshot: snapshot::ReadSnapshot::new(snapshot),
4409 active_segment,
4410 generation,
4411 _active_query: active_query,
4412 })
4413 }
4414
4415 #[cfg(test)]
4416 fn search_independent_for_test(
4417 &self,
4418 request: crate::ingest::SearchRequest<'_>,
4419 k: usize,
4420 options: SearchOptions,
4421 control: QueryControl,
4422 ) -> Result<crate::ingest::SearchOutcome, QueryError> {
4423 self.search_with_graph_bound_mode(request, k, options, control, GraphBoundMode::Independent)
4424 }
4425}
4426
4427#[derive(Clone, Copy)]
4428pub(crate) enum StructuredLexicalSource {
4429 Sealed(usize),
4430 Active,
4431}
4432
4433struct LexicalAssembly {
4434 index: crate::fts::index::LexicalIndex,
4435 alive_sets: Vec<Arc<crate::meta::AliveSet>>,
4436 sources: Vec<StructuredLexicalSource>,
4437 contributions: Vec<Arc<CachedLexicalContribution>>,
4438 vocabulary: Mutex<Option<Arc<CachedVocabulary>>>,
4439 memory: stats::AccountedCounter,
4441}
4442
4443struct CachedVocabulary {
4444 view: crate::fts::vocabulary::Vocabulary,
4445 phonetic: Mutex<Option<Arc<expansion::CachedPhonetic>>>,
4446 _memory: stats::AccountedCounter,
4448}
4449
4450impl LexicalAssembly {
4451 fn vocabulary(
4452 &self,
4453 accounting: &Arc<stats::Accounting>,
4454 cancellation: &QueryCancellation<'_>,
4455 ) -> Result<Arc<CachedVocabulary>, QueryError> {
4456 let mut cached = cancellation.cache_lock(&self.vocabulary, "lexical vocabulary cache")?;
4457 if let Some(value) = cached.as_ref() {
4458 return Ok(Arc::clone(value));
4459 }
4460 let mut memory =
4461 stats::AccountedCounter::new(accounting, stats::AllocationComponent::Cache)
4462 .map_err(QueryError::Store)?;
4463 let mut temporary =
4464 stats::AccountedCounter::new(accounting, stats::AllocationComponent::Temporary)
4465 .map_err(QueryError::Store)?;
4466 let mut work = crate::fts::control::WorkCheck::new(|| {
4467 cancellation.check_graph().map_err(QueryError::Scan)
4468 });
4469 let view = crate::fts::vocabulary::Vocabulary::build_controlled(
4470 self.index.vocabulary_terms(),
4471 |scratch, owned| {
4472 cancellation.check_graph().map_err(QueryError::Scan)?;
4473 prepared_lexical::reserve_weighted_scratch(&mut temporary, scratch)?;
4474 let bytes = owned
4475 .and_then(|bytes| bytes.checked_add(std::mem::size_of::<CachedVocabulary>()))
4476 .ok_or(QueryError::Store(StoreError::BudgetExceeded {
4477 needed: u64::MAX,
4478 budget: u64::MAX,
4479 component: "lexical vocabulary cache",
4480 }))?;
4481 memory.set(bytes).map_err(QueryError::Store)
4482 },
4483 &mut work,
4484 )?;
4485 let value = Arc::new(CachedVocabulary {
4486 view,
4487 phonetic: Mutex::new(None),
4488 _memory: memory,
4489 });
4490 *cached = Some(Arc::clone(&value));
4491 Ok(value)
4492 }
4493
4494 fn owned_bytes(&self) -> Option<usize> {
4497 std::mem::size_of::<Self>()
4498 .checked_add(self.index.resident_spine_bytes()?)?
4499 .checked_add(
4500 self.alive_sets
4501 .capacity()
4502 .checked_mul(std::mem::size_of::<Arc<crate::meta::AliveSet>>())?,
4503 )?
4504 .checked_add(
4505 self.sources
4506 .capacity()
4507 .checked_mul(std::mem::size_of::<StructuredLexicalSource>())?,
4508 )?
4509 .checked_add(
4510 self.contributions
4511 .capacity()
4512 .checked_mul(std::mem::size_of::<Arc<CachedLexicalContribution>>())?,
4513 )
4514 }
4515}
4516
4517enum LexicalContributionKey {
4518 Sealed(crate::segment::SegmentMeta, (u64, u64)),
4521 Active(std::sync::Weak<crate::ingest::ActiveSegment>),
4522}
4523
4524struct CachedLexicalContribution {
4525 key: LexicalContributionKey,
4526 statistics: Arc<crate::fts::index::LiveSegmentStatistics>,
4527 active_alive: Option<Arc<crate::meta::AliveSet>>,
4530 _memory: stats::AccountedCounter,
4531}
4532
4533impl CachedLexicalContribution {
4534 fn build(
4535 key: LexicalContributionKey,
4536 postings: &crate::fts::sealed::SealedSegment,
4537 alive: &Arc<crate::meta::AliveSet>,
4538 ordinal: usize,
4539 accounting: &Arc<stats::Accounting>,
4540 work: &mut crate::fts::control::WorkCheck<impl FnMut() -> Result<(), LexicalAssemblyError>>,
4541 ) -> Result<Arc<Self>, LexicalAssemblyError> {
4542 work.check_now()?;
4543 let active = matches!(key, LexicalContributionKey::Active(_));
4544 let alive_bytes = if active {
4545 alive
4546 .resident_bytes()
4547 .and_then(|n| n.checked_add(std::mem::size_of::<crate::meta::AliveSet>()))
4548 .and_then(|n| n.checked_add(2 * std::mem::size_of::<usize>()))
4549 } else {
4550 Some(0)
4551 };
4552 let bytes =
4553 crate::fts::index::LiveSegmentStatistics::allocation_bytes(alive.alive_bitmap())
4554 .and_then(|n| n.checked_add(std::mem::size_of::<Self>()))
4555 .and_then(|n| n.checked_add(2 * std::mem::size_of::<usize>()))
4556 .and_then(|n| n.checked_add(alive_bytes?))
4557 .ok_or(LexicalAssemblyError::Store(StoreError::BudgetExceeded {
4558 needed: u64::MAX,
4559 budget: u64::MAX,
4560 component: "cache",
4561 }))?;
4562 let mut memory =
4563 stats::AccountedCounter::new(accounting, stats::AllocationComponent::Cache)
4564 .map_err(LexicalAssemblyError::Store)?;
4565 memory.set(bytes).map_err(LexicalAssemblyError::Store)?;
4566 let mut statistics = crate::fts::index::LiveSegmentStatistics::build_controlled(
4567 ordinal,
4568 postings,
4569 alive.alive_bitmap(),
4570 work,
4571 )?;
4572 if statistics.rows.cardinality() != 0
4573 && statistics.rows.cardinality() != u64::from(postings.row_count())
4574 {
4575 statistics.frequency_cache = Some(
4576 crate::fts::live_df::LiveFrequencyCache::new(accounting)
4577 .map_err(LexicalAssemblyError::Store)?,
4578 );
4579 }
4580 let statistics = Arc::new(statistics);
4581 work.check_now()?;
4582 Ok(Arc::new(Self {
4583 key,
4584 statistics,
4585 active_alive: active.then(|| Arc::clone(alive)),
4586 _memory: memory,
4587 }))
4588 }
4589}
4590
4591enum LexicalAssemblyError {
4592 Store(StoreError),
4593 Lexical(crate::fts::index::IndexError),
4594 MissingDocumentIdentity(crate::segment::SegmentId),
4595 Cancelled(crate::scan::ScanError),
4596}
4597
4598impl From<crate::fts::index::IndexError> for LexicalAssemblyError {
4599 fn from(error: crate::fts::index::IndexError) -> Self {
4600 Self::Lexical(error)
4601 }
4602}
4603
4604impl From<StoreError> for LexicalAssemblyError {
4605 fn from(error: StoreError) -> Self {
4606 Self::Store(error)
4607 }
4608}
4609
4610struct CachedLexicalAssembly {
4637 generation: u64,
4638 snapshot: std::sync::Weak<PublishedSnapshot>,
4639 active: std::sync::Weak<crate::ingest::ActiveSegment>,
4640 document_identity_verified: bool,
4645 assembly: Arc<LexicalAssembly>,
4646}
4647
4648impl CachedLexicalAssembly {
4649 fn matches(
4650 &self,
4651 snapshot: &Arc<PublishedSnapshot>,
4652 active: &Arc<crate::ingest::ActiveSegment>,
4653 ) -> bool {
4654 self.snapshot
4655 .upgrade()
4656 .is_some_and(|cached| Arc::ptr_eq(&cached, snapshot))
4657 && self
4658 .active
4659 .upgrade()
4660 .is_some_and(|cached| Arc::ptr_eq(&cached, active))
4661 }
4662}
4663
4664pub(crate) struct LexicalIndexCache {
4670 entry: Mutex<Option<CachedLexicalAssembly>>,
4671 #[cfg(any(test, feature = "test-support"))]
4672 hits: AtomicU64,
4673 #[cfg(any(test, feature = "test-support"))]
4674 builds: AtomicU64,
4675 #[cfg(any(test, feature = "test-support"))]
4676 contribution_hits: AtomicU64,
4677 #[cfg(any(test, feature = "test-support"))]
4678 contribution_builds: AtomicU64,
4679}
4680
4681impl LexicalIndexCache {
4682 pub(crate) const fn new() -> Self {
4683 Self {
4684 entry: Mutex::new(None),
4685 #[cfg(any(test, feature = "test-support"))]
4686 hits: AtomicU64::new(0),
4687 #[cfg(any(test, feature = "test-support"))]
4688 builds: AtomicU64::new(0),
4689 #[cfg(any(test, feature = "test-support"))]
4690 contribution_hits: AtomicU64::new(0),
4691 #[cfg(any(test, feature = "test-support"))]
4692 contribution_builds: AtomicU64::new(0),
4693 }
4694 }
4695
4696 #[cfg(any(test, feature = "test-support"))]
4697 fn record_hit(&self) {
4698 self.hits.fetch_add(1, Ordering::Relaxed);
4699 }
4700
4701 #[cfg(not(any(test, feature = "test-support")))]
4702 const fn record_hit(&self) {}
4703
4704 #[cfg(any(test, feature = "test-support"))]
4705 fn record_build(&self) {
4706 self.builds.fetch_add(1, Ordering::Relaxed);
4707 }
4708
4709 #[cfg(not(any(test, feature = "test-support")))]
4710 const fn record_build(&self) {}
4711
4712 fn record_contribution(&self, hit: bool) {
4713 #[cfg(any(test, feature = "test-support"))]
4714 if hit {
4715 self.contribution_hits.fetch_add(1, Ordering::Relaxed);
4716 } else {
4717 self.contribution_builds.fetch_add(1, Ordering::Relaxed);
4718 }
4719 #[cfg(not(any(test, feature = "test-support")))]
4720 let _ = hit;
4721 }
4722
4723 fn lock(
4724 &self,
4725 ) -> Result<std::sync::MutexGuard<'_, Option<CachedLexicalAssembly>>, LexicalAssemblyError>
4726 {
4727 self.entry.lock().map_err(|_| {
4728 LexicalAssemblyError::Store(StoreError::Synchronization {
4729 component: "lexical index cache",
4730 })
4731 })
4732 }
4733
4734 fn lock_controlled(
4735 &self,
4736 cancellation: Option<&QueryCancellation<'_>>,
4737 ) -> Result<std::sync::MutexGuard<'_, Option<CachedLexicalAssembly>>, LexicalAssemblyError>
4738 {
4739 let Some(cancellation) = cancellation else {
4740 return self.lock();
4741 };
4742 cancellation
4743 .check_graph()
4744 .map_err(LexicalAssemblyError::Cancelled)?;
4745 let guard = loop {
4746 match self.entry.try_lock() {
4747 Ok(guard) => break guard,
4748 Err(std::sync::TryLockError::Poisoned(_)) => {
4749 return Err(LexicalAssemblyError::Store(StoreError::Synchronization {
4750 component: "lexical index cache",
4751 }));
4752 }
4753 Err(std::sync::TryLockError::WouldBlock) => {}
4754 }
4755 cancellation
4756 .check_graph()
4757 .map_err(LexicalAssemblyError::Cancelled)?;
4758 std::thread::park_timeout(Duration::from_millis(1));
4759 };
4760 cancellation
4761 .check_graph()
4762 .map_err(LexicalAssemblyError::Cancelled)?;
4763 Ok(guard)
4764 }
4765}
4766
4767#[derive(Clone, Copy)]
4771struct LexicalInputs<'a> {
4772 generation: u64,
4773 cache: &'a LexicalIndexCache,
4774 snapshot: &'a Arc<PublishedSnapshot>,
4775 active: &'a Arc<crate::ingest::ActiveSegment>,
4776 accounting: &'a Arc<stats::Accounting>,
4777}
4778
4779struct LexicalAssemblyReceipt {
4780 assembly: Arc<LexicalAssembly>,
4781 cache_hit: bool,
4782}
4783
4784impl std::ops::Deref for LexicalAssemblyReceipt {
4785 type Target = LexicalAssembly;
4786 fn deref(&self) -> &Self::Target {
4787 &self.assembly
4788 }
4789}
4790
4791fn assemble_lexical_index(
4792 inputs: LexicalInputs<'_>,
4793 require_document_identity: bool,
4794 cancellation: Option<&QueryCancellation<'_>>,
4795) -> Result<LexicalAssemblyReceipt, LexicalAssemblyError> {
4796 let LexicalInputs {
4797 generation,
4798 cache,
4799 snapshot,
4800 active,
4801 accounting,
4802 } = inputs;
4803 let mut entry = cache.lock_controlled(cancellation)?;
4804 if let Some(cached) = entry.as_mut()
4805 && cached.matches(snapshot, active)
4806 {
4807 if require_document_identity && !cached.document_identity_verified {
4808 verify_document_identity(snapshot, cancellation)?;
4809 cached.document_identity_verified = true;
4810 }
4811 let assembly = Arc::clone(&cached.assembly);
4812 cache.record_hit();
4813 return Ok(LexicalAssemblyReceipt {
4814 assembly,
4815 cache_hit: true,
4816 });
4817 }
4818 let previous = entry.as_ref().map(|cached| Arc::clone(&cached.assembly));
4822 let contributions = previous
4823 .as_ref()
4824 .map_or(&[][..], |a| a.contributions.as_slice());
4825 let (mut assembly, owned_bytes) = build_lexical_assembly(
4826 snapshot,
4827 active,
4828 accounting,
4829 require_document_identity,
4830 cancellation,
4831 cache,
4832 contributions,
4833 )?;
4834 cache.record_build();
4835 assembly
4836 .memory
4837 .set(owned_bytes)
4838 .map_err(LexicalAssemblyError::Store)?;
4839 if let Some(cancellation) = cancellation {
4840 cancellation
4841 .check_graph()
4842 .map_err(LexicalAssemblyError::Cancelled)?;
4843 }
4844 let assembly = Arc::new(assembly);
4845 if entry
4846 .as_ref()
4847 .is_none_or(|cached| generation >= cached.generation)
4848 {
4849 *entry = Some(CachedLexicalAssembly {
4850 generation,
4851 snapshot: Arc::downgrade(snapshot),
4852 active: Arc::downgrade(active),
4853 document_identity_verified: require_document_identity,
4854 assembly: Arc::clone(&assembly),
4855 });
4856 }
4857 drop(entry);
4858 Ok(LexicalAssemblyReceipt {
4859 assembly,
4860 cache_hit: false,
4861 })
4862}
4863
4864fn check_segment_document_identity(
4865 segment: &crate::segment::reader::SegmentReader,
4866 postings: &crate::fts::sealed::SealedSegment,
4867) -> Result<(), LexicalAssemblyError> {
4868 if postings.row_count() != 0
4869 && segment
4870 .document_version(0)
4871 .map_err(StoreError::Segment)
4872 .map_err(LexicalAssemblyError::Store)?
4873 .is_none()
4874 {
4875 return Err(LexicalAssemblyError::MissingDocumentIdentity(
4876 segment.meta().id,
4877 ));
4878 }
4879 Ok(())
4880}
4881
4882fn verify_document_identity(
4888 snapshot: &PublishedSnapshot,
4889 cancellation: Option<&QueryCancellation<'_>>,
4890) -> Result<(), LexicalAssemblyError> {
4891 for segment in snapshot.segments() {
4892 if let Some(cancellation) = cancellation {
4893 cancellation
4894 .check_graph()
4895 .map_err(LexicalAssemblyError::Cancelled)?;
4896 }
4897 if let Some(postings) = segment
4898 .query_postings()
4899 .map_err(LexicalAssemblyError::Store)?
4900 {
4901 check_segment_document_identity(segment, &postings)?;
4902 }
4903 }
4904 Ok(())
4905}
4906
4907fn build_lexical_assembly(
4908 snapshot: &PublishedSnapshot,
4909 active: &Arc<crate::ingest::ActiveSegment>,
4910 accounting: &Arc<stats::Accounting>,
4911 require_document_identity: bool,
4912 cancellation: Option<&QueryCancellation<'_>>,
4913 cache: &LexicalIndexCache,
4914 previous: &[Arc<CachedLexicalContribution>],
4915) -> Result<(LexicalAssembly, usize), LexicalAssemblyError> {
4916 let mut work = crate::fts::control::WorkCheck::new(|| {
4917 cancellation
4918 .map_or(Ok(()), |control| control.check_graph())
4919 .map_err(LexicalAssemblyError::Cancelled)
4920 });
4921 work.check_now()?;
4922 let mut index = crate::fts::index::LexicalIndex::new();
4923 let mut alive_sets = Vec::new();
4924 let mut sources = Vec::new();
4925 let mut contributions = Vec::new();
4926 for (ordinal, segment) in snapshot.segments().iter().enumerate() {
4927 if let Some(cancellation) = cancellation {
4928 cancellation
4929 .check_graph()
4930 .map_err(LexicalAssemblyError::Cancelled)?;
4931 }
4932 if let Some(postings) = segment
4933 .query_postings()
4934 .map_err(LexicalAssemblyError::Store)?
4935 {
4936 if require_document_identity {
4937 check_segment_document_identity(segment, &postings)?;
4938 }
4939 let alive = segment.query_alive().map_err(LexicalAssemblyError::Store)?;
4940 let identity = segment
4941 .lexical_postings_identity()
4942 .map_err(StoreError::Segment)
4943 .map_err(LexicalAssemblyError::Store)?;
4944 let reused = previous
4945 .binary_search_by(|value| match &value.key {
4946 LexicalContributionKey::Sealed(meta, _) => meta.id.cmp(&segment.meta().id),
4947 LexicalContributionKey::Active(_) => std::cmp::Ordering::Greater,
4948 })
4949 .ok()
4950 .and_then(|slot| previous.get(slot))
4951 .filter(|value| {
4952 matches!(&value.key, LexicalContributionKey::Sealed(meta, stored)
4953 if meta.same_segment_file(segment.meta()) && *stored == identity)
4954 && value.statistics.rows == *alive.alive_bitmap()
4955 });
4956 cache.record_contribution(reused.is_some());
4957 let contribution = match reused {
4958 Some(value) => Arc::clone(value),
4959 None => CachedLexicalContribution::build(
4960 LexicalContributionKey::Sealed(segment.meta().clone(), identity),
4961 &postings,
4962 &alive,
4963 index.segments().len(),
4964 accounting,
4965 &mut work,
4966 )?,
4967 };
4968 index.push_shared_with_statistics(postings, Arc::clone(&contribution.statistics));
4969 contributions.push(contribution);
4970 alive_sets.push(alive);
4971 sources.push(StructuredLexicalSource::Sealed(ordinal));
4972 }
4973 }
4974 if active.has_text_controlled(&mut work)? {
4975 let sealed = active.sealed_lexical_controlled(accounting, &mut work)?;
4976 let reused = previous.last().filter(|value| {
4977 matches!(&value.key, LexicalContributionKey::Active(cached)
4978 if cached.upgrade().is_some_and(|cached| Arc::ptr_eq(&cached, active)))
4979 });
4980 cache.record_contribution(reused.is_some());
4981 let contribution = match reused {
4982 Some(value) => Arc::clone(value),
4983 None => {
4984 let alive = Arc::new(active.alive_controlled(&mut work)?);
4985 CachedLexicalContribution::build(
4986 LexicalContributionKey::Active(Arc::downgrade(active)),
4987 &sealed,
4988 &alive,
4989 index.segments().len(),
4990 accounting,
4991 &mut work,
4992 )?
4993 }
4994 };
4995 let alive = contribution
4996 .active_alive
4997 .as_ref()
4998 .ok_or(LexicalAssemblyError::Store(StoreError::Synchronization {
4999 component: "active lexical contribution",
5000 }))?;
5001 index.push_shared_with_statistics(sealed, Arc::clone(&contribution.statistics));
5002 alive_sets.push(Arc::clone(alive));
5003 contributions.push(contribution);
5004 sources.push(StructuredLexicalSource::Active);
5005 }
5006 crate::fts::control::sort_by(
5007 &mut contributions,
5008 |left, right| match (&left.key, &right.key) {
5009 (LexicalContributionKey::Sealed(a, _), LexicalContributionKey::Sealed(b, _)) => {
5010 a.id.cmp(&b.id)
5011 }
5012 (LexicalContributionKey::Active(_), LexicalContributionKey::Active(_)) => {
5013 std::cmp::Ordering::Equal
5014 }
5015 (LexicalContributionKey::Active(_), _) => std::cmp::Ordering::Greater,
5016 (_, LexicalContributionKey::Active(_)) => std::cmp::Ordering::Less,
5017 },
5018 &mut work,
5019 )?;
5020 let assembly = LexicalAssembly {
5021 index,
5022 alive_sets,
5023 sources,
5024 contributions,
5025 vocabulary: Mutex::new(None),
5026 memory: stats::AccountedCounter::new(accounting, stats::AllocationComponent::Cache)
5027 .map_err(LexicalAssemblyError::Store)?,
5028 };
5029 let owned_bytes =
5030 assembly
5031 .owned_bytes()
5032 .ok_or(LexicalAssemblyError::Store(StoreError::BudgetExceeded {
5033 needed: u64::MAX,
5034 budget: u64::MAX,
5035 component: "cache",
5036 }))?;
5037 work.check_now()?;
5038 Ok((assembly, owned_bytes))
5039}
5040
5041fn map_store_lexical_assembly_error(
5042 error: LexicalAssemblyError,
5043) -> crate::ingest::StoreLexicalError {
5044 match error {
5045 LexicalAssemblyError::Store(error) => QueryError::Store(error).into(),
5046 LexicalAssemblyError::Lexical(error) => {
5047 crate::planner::LexicalFilterError::from(error).into()
5048 }
5049 LexicalAssemblyError::MissingDocumentIdentity(segment_id) => {
5050 crate::ingest::StoreLexicalError::MissingDocumentIdentity { segment_id }
5051 }
5052 LexicalAssemblyError::Cancelled(error) => QueryError::Scan(error).into(),
5053 }
5054}
5055
5056fn map_fusion_lexical_assembly_error(error: LexicalAssemblyError) -> crate::fusion::FusionError {
5057 match error {
5058 LexicalAssemblyError::Store(error) => crate::fusion::FusionError::Leg {
5059 leg: crate::fusion::FusionLeg::Lexical,
5060 kind: crate::fusion::LegFailureKind::Store(error.kind()),
5061 detail: error.to_string(),
5062 },
5063 LexicalAssemblyError::Lexical(error) => crate::fusion::FusionError::Leg {
5064 leg: crate::fusion::FusionLeg::Lexical,
5065 kind: crate::fusion::LegFailureKind::Lexical,
5066 detail: error.to_string(),
5067 },
5068 LexicalAssemblyError::MissingDocumentIdentity(segment_id) => {
5069 crate::fusion::FusionError::Leg {
5070 leg: crate::fusion::FusionLeg::Lexical,
5071 kind: crate::fusion::LegFailureKind::Invariant,
5072 detail: format!("sealed lexical segment {segment_id} has no document identity"),
5073 }
5074 }
5075 LexicalAssemblyError::Cancelled(error) => {
5076 crate::fusion::FusionError::from(QueryError::Scan(error))
5077 }
5078 }
5079}
5080
5081fn accumulate_search_counters(
5082 total: &mut crate::fts::search::SearchCounters,
5083 delta: &crate::fts::search::SearchCounters,
5084) {
5085 total.docs_evaluated = total.docs_evaluated.saturating_add(delta.docs_evaluated);
5086 total.postings_decoded = total
5087 .postings_decoded
5088 .saturating_add(delta.postings_decoded);
5089 total.blocks_decoded = total.blocks_decoded.saturating_add(delta.blocks_decoded);
5090 total.blocks_skipped = total.blocks_skipped.saturating_add(delta.blocks_skipped);
5091}
5092
5093fn map_fusion_controlled_lexical_error(
5094 error: crate::fts::search::ControlledSearchError<crate::scan::ScanError>,
5095) -> crate::fusion::FusionError {
5096 match error {
5097 crate::fts::search::ControlledSearchError::Index(error) => {
5098 crate::fusion::FusionError::Leg {
5099 leg: crate::fusion::FusionLeg::Lexical,
5100 kind: crate::fusion::LegFailureKind::Lexical,
5101 detail: error.to_string(),
5102 }
5103 }
5104 crate::fts::search::ControlledSearchError::Control(error) => {
5105 crate::fusion::FusionError::from(QueryError::Scan(error))
5106 }
5107 }
5108}
5109
5110#[derive(Clone, Copy)]
5111enum PinnedLexicalQuery<'a> {
5112 Term(&'a crate::fts::search::TermQuery),
5113 Structured(&'a crate::fts::query::LexicalQuery),
5114}
5115
5116enum LexicalDocumentError {
5117 SourceOrdinalOutOfRange,
5118 RowOverflow,
5119 SealedSourceAbsent,
5120 Segment(crate::segment::SegmentError),
5121 MissingDocumentIdentity(crate::segment::SegmentId),
5122 ActiveRowAbsent,
5123}
5124
5125fn structured_lexical_document(
5126 snapshot: &PublishedSnapshot,
5127 active: &crate::ingest::ActiveSegment,
5128 sources: &[StructuredLexicalSource],
5129 doc: crate::fts::search::GlobalDocId,
5130 require_document_identity: bool,
5131) -> Result<Option<crate::ingest::DocumentVersion>, LexicalDocumentError> {
5132 let source = usize::try_from(doc.segment)
5133 .ok()
5134 .and_then(|slot| sources.get(slot))
5135 .ok_or(LexicalDocumentError::SourceOrdinalOutOfRange)?;
5136 let row = usize::try_from(doc.row).map_err(|_| LexicalDocumentError::RowOverflow)?;
5137 match source {
5138 StructuredLexicalSource::Sealed(ordinal) => {
5139 let segment = snapshot
5140 .segments()
5141 .get(*ordinal)
5142 .ok_or(LexicalDocumentError::SealedSourceAbsent)?;
5143 let document = segment
5144 .document_version(row)
5145 .map_err(LexicalDocumentError::Segment)?;
5146 if require_document_identity && document.is_none() {
5147 return Err(LexicalDocumentError::MissingDocumentIdentity(
5148 segment.meta().id,
5149 ));
5150 }
5151 Ok(document)
5152 }
5153 StructuredLexicalSource::Active => {
5154 let document = active.document(row);
5155 if require_document_identity && document.is_none() {
5156 return Err(LexicalDocumentError::ActiveRowAbsent);
5157 }
5158 Ok(document)
5159 }
5160 }
5161}
5162
5163fn map_store_lexical_document_error(
5164 error: LexicalDocumentError,
5165) -> crate::ingest::StoreLexicalError {
5166 match error {
5167 LexicalDocumentError::SourceOrdinalOutOfRange
5168 | LexicalDocumentError::RowOverflow
5169 | LexicalDocumentError::SealedSourceAbsent
5170 | LexicalDocumentError::ActiveRowAbsent => {
5171 QueryError::Store(StoreError::ActiveRowOverflow).into()
5172 }
5173 LexicalDocumentError::Segment(error) => {
5174 QueryError::Store(StoreError::Segment(error)).into()
5175 }
5176 LexicalDocumentError::MissingDocumentIdentity(segment_id) => {
5177 crate::ingest::StoreLexicalError::MissingDocumentIdentity { segment_id }
5178 }
5179 }
5180}
5181
5182fn map_fusion_lexical_document_error(
5183 error: LexicalDocumentError,
5184 invariant_kind: crate::fusion::LegFailureKind,
5185 segment_kind: crate::fusion::LegFailureKind,
5186) -> crate::fusion::FusionError {
5187 let (kind, detail) = match error {
5188 LexicalDocumentError::SourceOrdinalOutOfRange => (
5189 invariant_kind,
5190 "lexical source ordinal is out of range".to_owned(),
5191 ),
5192 LexicalDocumentError::RowOverflow => {
5193 (invariant_kind, "lexical row exceeds usize".to_owned())
5194 }
5195 LexicalDocumentError::SealedSourceAbsent => {
5196 (invariant_kind, "lexical sealed source is absent".to_owned())
5197 }
5198 LexicalDocumentError::Segment(error) => (segment_kind, error.to_string()),
5199 LexicalDocumentError::MissingDocumentIdentity(segment_id) => (
5200 invariant_kind,
5201 format!("sealed lexical segment {segment_id} has no document identity"),
5202 ),
5203 LexicalDocumentError::ActiveRowAbsent => (
5204 invariant_kind,
5205 "lexical active source row is absent".to_owned(),
5206 ),
5207 };
5208 crate::fusion::FusionError::Leg {
5209 leg: crate::fusion::FusionLeg::Lexical,
5210 kind,
5211 detail,
5212 }
5213}
5214
5215fn map_structured_leg_document_error(error: LexicalDocumentError) -> crate::fusion::FusionError {
5216 map_fusion_lexical_document_error(
5217 error,
5218 crate::fusion::LegFailureKind::Lexical,
5219 crate::fusion::LegFailureKind::Lexical,
5220 )
5221}
5222
5223fn map_term_leg_document_error(error: LexicalDocumentError) -> crate::fusion::FusionError {
5224 map_fusion_lexical_document_error(
5225 error,
5226 crate::fusion::LegFailureKind::Invariant,
5227 crate::fusion::LegFailureKind::Segment,
5228 )
5229}
5230
5231fn structured_lexical_row<'a>(
5232 snapshot: &'a PublishedSnapshot,
5233 active: &'a crate::ingest::ActiveSegment,
5234 sources: &[StructuredLexicalSource],
5235 doc: crate::fts::search::GlobalDocId,
5236) -> Result<(&'a str, crate::ingest::DocumentVersion), crate::ingest::StoreLexicalError> {
5237 let source = usize::try_from(doc.segment)
5238 .ok()
5239 .and_then(|slot| sources.get(slot))
5240 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
5241 let row =
5242 usize::try_from(doc.row).map_err(|_| QueryError::Store(StoreError::ActiveRowOverflow))?;
5243 match source {
5244 StructuredLexicalSource::Sealed(ordinal) => {
5245 let segment = snapshot
5246 .segments()
5247 .get(*ordinal)
5248 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
5249 let text = segment
5250 .query_stored_text()
5251 .map_err(StoreError::Segment)
5252 .map_err(QueryError::Store)?
5253 .and_then(|rows| rows.row(row).flatten())
5254 .ok_or(crate::ingest::StoreLexicalError::MissingStoredText {
5255 segment_id: segment.meta().id,
5256 row: doc.row,
5257 })?;
5258 let document = segment
5259 .document_version(row)
5260 .map_err(StoreError::Segment)
5261 .map_err(QueryError::Store)?
5262 .ok_or(crate::ingest::StoreLexicalError::MissingDocumentIdentity {
5263 segment_id: segment.meta().id,
5264 })?;
5265 Ok((text, document))
5266 }
5267 StructuredLexicalSource::Active => {
5268 let text = active
5269 .text(row)
5270 .map_err(QueryError::Store)?
5271 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
5272 let document = active
5273 .document(row)
5274 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
5275 Ok((text, document))
5276 }
5277 }
5278}
5279
5280type ExactLexicalLeg = (
5281 Vec<hybrid::LexicalHit>,
5282 Arc<LexicalAssembly>,
5283 crate::fts::search::SearchCounters,
5284 bool,
5285);
5286
5287fn resolve_hybrid_leg_results<Vector, Lexical>(
5288 vector: Result<Vector, crate::fusion::FusionError>,
5289 lexical: Result<Lexical, crate::fusion::FusionError>,
5290) -> Result<(Vector, Lexical), crate::fusion::FusionError> {
5291 fn control_priority(error: &crate::fusion::FusionError) -> Option<u8> {
5292 match error {
5293 crate::fusion::FusionError::ReadCancelled { .. } => Some(0),
5294 crate::fusion::FusionError::Timeout { .. } => Some(1),
5295 crate::fusion::FusionError::Cancelled { .. } => Some(2),
5296 _ => None,
5297 }
5298 }
5299
5300 match (vector, lexical) {
5301 (Ok(vector), Ok(lexical)) => Ok((vector, lexical)),
5302 (Err(vector), Ok(_)) => Err(vector),
5303 (Ok(_), Err(lexical)) => Err(lexical),
5304 (Err(vector), Err(lexical)) => {
5305 match (control_priority(&vector), control_priority(&lexical)) {
5306 (Some(vector_priority), Some(lexical_priority)) => {
5307 if vector_priority <= lexical_priority {
5308 Err(vector)
5309 } else {
5310 Err(lexical)
5311 }
5312 }
5313 (Some(_), None) => Err(vector),
5314 (None, Some(_)) => Err(lexical),
5315 (None, None) => Err(vector),
5316 }
5317 }
5318 }
5319}
5320
5321fn exact_structured_lexical_leg(
5322 inputs: LexicalInputs<'_>,
5323 _analyzer: &crate::fts::tokenizer::Analyzer,
5324 query: &crate::fts::query::LexicalQuery,
5325 bound: usize,
5326 cancellation: &QueryCancellation<'_>,
5327 preparation: &mut prepared_lexical::PreparedLexicalQuery<'_>,
5328) -> Result<ExactLexicalLeg, crate::fusion::FusionError> {
5329 let (snapshot, active) = (inputs.snapshot, inputs.active);
5330 let lexical_error = |detail: String| crate::fusion::FusionError::Leg {
5331 leg: crate::fusion::FusionLeg::Lexical,
5332 kind: crate::fusion::LegFailureKind::Lexical,
5333 detail,
5334 };
5335 let (prepared, cache_hit) = preparation.prepare_structured(inputs, cancellation)?;
5336 let assembly = &prepared.assembly;
5337 let index = &assembly.index;
5338 let sources = &assembly.sources;
5339 let Some(query_scoring) = prepared.query.as_ref() else {
5340 return Ok((
5341 Vec::new(),
5342 Arc::clone(assembly),
5343 crate::fts::search::SearchCounters::default(),
5344 cache_hit,
5345 ));
5346 };
5347 let allow_lists = assembly
5348 .alive_sets
5349 .iter()
5350 .map(|alive| alive.alive_bitmap())
5351 .collect::<Vec<_>>();
5352 let mut scratch_memory =
5353 stats::AccountedCounter::new(inputs.accounting, stats::AllocationComponent::Temporary)
5354 .map_err(QueryError::Store)?;
5355 let mut phrase = prepared_lexical::PhraseEligibility::new(
5356 index,
5357 Some(query),
5358 inputs.accounting,
5359 cancellation,
5360 )?;
5361 let result = crate::fts::prune::weighted::search(
5362 index,
5363 query_scoring,
5364 bound,
5365 &allow_lists,
5366 || {
5367 cancellation
5368 .check_graph()
5369 .map_err(QueryError::Scan)
5370 .map_err(crate::fusion::FusionError::from)
5371 },
5372 |doc| {
5373 phrase
5374 .matches(doc, cancellation)
5375 .map_err(crate::fusion::FusionError::from)
5376 },
5377 |bytes| {
5378 prepared_lexical::reserve_weighted_scratch(&mut scratch_memory, bytes)
5379 .map_err(crate::fusion::FusionError::from)
5380 },
5381 )
5382 .map_err(|error| match error {
5383 crate::fts::search::ControlledSearchError::Index(error) => lexical_error(error.to_string()),
5384 crate::fts::search::ControlledSearchError::Control(error) => error,
5385 })?;
5386 let counters = result.counters;
5387 let scored = result.hits;
5388 let mut joined = Vec::with_capacity(scored.len());
5389 for (doc, score) in scored {
5390 let document = structured_lexical_document(snapshot, active, sources, doc, false)
5391 .map_err(map_structured_leg_document_error)?;
5392 joined.push(hybrid::LexicalHit {
5393 doc,
5394 document,
5395 bm25: score,
5396 });
5397 }
5398 Ok((joined, Arc::clone(assembly), counters, cache_hit))
5399}
5400
5401fn exact_lexical_leg(
5402 inputs: LexicalInputs<'_>,
5403 bound: usize,
5404 cancellation: &QueryCancellation<'_>,
5405 preparation: &mut prepared_lexical::PreparedLexicalQuery<'_>,
5406) -> Result<ExactLexicalLeg, crate::fusion::FusionError> {
5407 let (snapshot, active) = (inputs.snapshot, inputs.active);
5408 let (prepared, cache_hit) = preparation.prepare_term(inputs, cancellation)?;
5409 let assembly = &prepared.assembly;
5410 let index = &assembly.index;
5411 let sources = &assembly.sources;
5412 if index.segments().is_empty() {
5413 return Ok((
5414 Vec::new(),
5415 Arc::clone(assembly),
5416 crate::fts::search::SearchCounters::default(),
5417 cache_hit,
5418 ));
5419 }
5420 let k = bound.min(usize::try_from(index.document_count()).unwrap_or(usize::MAX));
5421 let allow_lists = assembly
5422 .alive_sets
5423 .iter()
5424 .map(|alive| alive.alive_bitmap())
5425 .collect::<Vec<_>>();
5426 let result =
5427 exact_hybrid_lexical_search(index, prepared.scoring()?, k, &allow_lists, cancellation)?;
5428 let mut joined = Vec::with_capacity(result.hits.len());
5429 for hit in result.hits {
5430 let document = structured_lexical_document(snapshot, active, sources, hit.doc, false)
5431 .map_err(map_term_leg_document_error)?;
5432 joined.push(hybrid::LexicalHit {
5433 doc: hit.doc,
5434 document,
5435 bm25: hit.score,
5436 });
5437 }
5438 Ok((joined, Arc::clone(assembly), result.counters, cache_hit))
5439}
5440
5441fn exact_hybrid_lexical_search(
5442 index: &crate::fts::index::LexicalIndex,
5443 prepared: &crate::fts::search::PreparedTermQuery,
5444 k: usize,
5445 allow_lists: &[&crate::meta::DocBitmap],
5446 cancellation: &QueryCancellation<'_>,
5447) -> Result<crate::fts::search::SearchResult, crate::fusion::FusionError> {
5448 let allowed = allow_lists
5449 .iter()
5450 .map(|allow_list| allow_list.cardinality())
5451 .fold(0_u64, u64::saturating_add);
5452 if allowed.saturating_mul(crate::planner::LEXICAL_ALLOW_LIST_DIVISOR) <= index.document_count()
5453 {
5454 return crate::fts::search::search_allow_list_prepared_controlled(
5455 index,
5456 prepared,
5457 k,
5458 allow_lists,
5459 || cancellation.check_graph(),
5460 )
5461 .map_err(map_fusion_controlled_lexical_error);
5462 }
5463 crate::fts::prune::search_pruned_prepared_controlled(index, prepared, k, allow_lists, || {
5464 cancellation.check_graph()
5465 })
5466 .map_err(map_fusion_controlled_lexical_error)
5467}
5468
5469#[allow(clippy::too_many_arguments)]
5470fn search_pinned(
5471 pool: Option<&pool::QueryPool>,
5472 snapshot: &Arc<PublishedSnapshot>,
5473 active: &crate::ingest::ActiveSegment,
5474 accounting: &Arc<stats::Accounting>,
5475 generation: u64,
5476 epoch: Option<crate::epoch::EpochIdentity>,
5477 request: crate::ingest::SearchRequest<'_>,
5478 k: usize,
5479 options: SearchOptions,
5480 control: QueryControl,
5481 graph_bound_mode: GraphBoundMode,
5482 requested_tier: Option<SearchTier>,
5483 hybrid_scan_reason: Option<crate::planner::ScanReason>,
5484 started: std::time::Instant,
5485 clock: &dyn MonotonicClock,
5486 vector_preparation: &mut prepared::PreparedVectorQuery<'_>,
5487 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
5488 &crate::scan::vector_fault::VectorFaultController,
5489 >,
5490) -> Result<crate::ingest::SearchOutcome, QueryError> {
5491 use crate::ingest::{GraphSearchStats, RowSource, SearchOutcome};
5492 use crate::quant::{QuantError, prepare_bit4_query};
5493 use crate::scan::{ScanQuery, ScanRequest, ScanRows, ScanStats};
5494
5495 let vector_started = timing_start(clock);
5496 let scan_options = options.scan();
5497 let rescore_options = options.scan_rescore();
5498 let mut rescore_counters = crate::diag::ScanRescoreCounters::default();
5499 if let Some(options) = rescore_options {
5500 options.frontier(k, 0)?;
5501 }
5502 let query = request.vector();
5503 let quantized_query_validation = if query.is_empty() {
5504 Err(QuantError::EmptyVector)
5505 } else if query.len() > crate::kernels::MAX_DOT_I8_DIMENSION {
5506 Err(QuantError::DimensionTooLarge {
5507 actual: query.len(),
5508 maximum: crate::kernels::MAX_DOT_I8_DIMENSION,
5509 })
5510 } else if let Some((index, _)) = query
5511 .iter()
5512 .enumerate()
5513 .find(|(_, value)| !value.is_finite())
5514 {
5515 Err(QuantError::NonFinite { index })
5516 } else {
5517 Ok(())
5518 };
5519 quantized_query_validation
5520 .map_err(crate::scan::ScanError::Quant)
5521 .map_err(QueryError::Scan)?;
5522 vector_preparation.validate_binding(query, epoch)?;
5523 let bit4_query = std::cell::OnceCell::new();
5524 let int8_query = std::cell::OnceCell::new();
5525 let mut candidates = Vec::new();
5526 let mut dims_touched = 0_u64;
5527 let mut bytes_read = 0_u64;
5528 let mut worst_squared_l2: Option<f64> = None;
5529 let mut worst_exhaustive = true;
5530 let mut worker_thread_ids = Vec::new();
5531 let mut graph_stats = GraphSearchStats::default();
5532 let mut plans = Vec::new();
5533 let auto_graph_options = if matches!(options.tier(), SearchTier::Auto)
5534 && snapshot.segments().iter().any(|segment| {
5535 segment
5536 .directory()
5537 .iter()
5538 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
5539 }) {
5540 Some(auto_graph_search_options(snapshot)?)
5541 } else {
5542 None
5543 };
5544
5545 if matches!(options.tier(), SearchTier::Graph(_)) {
5546 for segment in snapshot.segments() {
5547 if !segment
5548 .directory()
5549 .iter()
5550 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
5551 {
5552 return Err(QueryError::Store(StoreError::GraphUnavailable {
5553 segment_id: segment.meta().id,
5554 }));
5555 }
5556 }
5557 }
5558 let full_precision = matches!(options.tier(), SearchTier::Exact | SearchTier::Graph(_))
5559 || (matches!(options.tier(), SearchTier::Auto)
5560 && auto_uses_full_precision(snapshot, active));
5561
5562 let exact_lane = if full_precision {
5564 pool.map_or(0, |pool| {
5565 pool.reserve_exact_lane(scan_options.thread_budget)
5566 })
5567 } else {
5568 0
5569 };
5570
5571 if !active.is_empty() {
5572 let alive = active.alive().map_err(QueryError::Store)?;
5573 let mut exact_score = full_precision;
5574 let mut approximate = false;
5575 let frontier = rescore_options
5576 .map(|options| options.frontier(k, alive.live_count()))
5577 .transpose()?
5578 .unwrap_or(k);
5579 let mut exact_memory = ExactScanMemory::new(accounting)?;
5580 let mut outcome = match options.tier() {
5581 SearchTier::Auto if full_precision => {
5582 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
5583 scan_active_squared_l2(
5584 active,
5585 &alive,
5586 request.vector(),
5587 frontier,
5588 ExactScanDispatch {
5589 pool,
5590 accounting,
5591 control: &control,
5592 lease: &lease,
5593 thread_budget: scan_options.thread_budget,
5594 lane: exact_lane,
5595 },
5596 &mut exact_memory,
5597 #[cfg(any(test, feature = "test-support"))]
5598 vector_fault_controller,
5599 #[cfg(any(test, feature = "test-support"))]
5600 vector_fault_tier(options.tier()),
5601 )?
5602 }
5603 SearchTier::Auto | SearchTier::Scan => {
5604 let query_pool = pool.ok_or(QueryError::Store(StoreError::Synchronization {
5605 component: "active scan-tier query pool",
5606 }))?;
5607 let bit4_query = bit4_query
5608 .get_or_init(|| prepare_bit4_query(request.vector(), 0))
5609 .as_ref()
5610 .map_err(|error| crate::scan::ScanError::Quant(error.clone()))
5611 .map_err(QueryError::Scan)?;
5612 execute_store_scan(
5613 query_pool,
5614 ScanRequest {
5615 query: ScanQuery::Bit4(bit4_query),
5616 rows: ScanRows::Bit4RowMajor {
5617 codes: active.codes(),
5618 factors: active.factors(),
5619 },
5620 row_mask: Some(alive.scan_mask()),
5621 },
5622 frontier,
5623 scan_options,
5624 control.clone(),
5625 SnapshotLease::new_at(Arc::clone(snapshot), generation),
5626 #[cfg(any(test, feature = "test-support"))]
5627 vector_fault_controller,
5628 #[cfg(any(test, feature = "test-support"))]
5629 crate::scan::vector_fault::VectorRowSource::Active,
5630 #[cfg(any(test, feature = "test-support"))]
5631 vector_fault_tier(options.tier()),
5632 )?
5633 }
5634 SearchTier::Exact | SearchTier::Graph(_) => {
5635 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
5636 scan_active_squared_l2(
5637 active,
5638 &alive,
5639 request.vector(),
5640 frontier,
5641 ExactScanDispatch {
5642 pool,
5643 accounting,
5644 control: &control,
5645 lease: &lease,
5646 thread_budget: scan_options.thread_budget,
5647 lane: exact_lane,
5648 },
5649 &mut exact_memory,
5650 #[cfg(any(test, feature = "test-support"))]
5651 vector_fault_controller,
5652 #[cfg(any(test, feature = "test-support"))]
5653 vector_fault_tier(options.tier()),
5654 )?
5655 }
5656 };
5657 if let Some(rescore) = rescore_options {
5658 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
5659 let cancellation = QueryCancellation::new(&control, &lease);
5660 let (counters, exhaustive, worst) = rescored_scan::rescore(
5661 &mut outcome.candidates,
5662 &mut outcome.stats,
5663 active.vectors(),
5664 query,
5665 k,
5666 alive.live_count(),
5667 rescore,
5668 &cancellation,
5669 accounting,
5670 &mut exact_memory,
5671 )?;
5672 outcome.worst_score = worst;
5673 exact_score = true;
5674 approximate = !exhaustive;
5675 rescored_scan::accumulate(&mut rescore_counters, counters)?;
5676 }
5677 fold_worst_squared_l2(&mut worst_squared_l2, &mut worst_exhaustive, &outcome);
5678 merge_store_outcome(
5679 outcome,
5680 RowSource::Active,
5681 |row| Ok(active.document(row)),
5682 &mut candidates,
5683 &mut dims_touched,
5684 &mut bytes_read,
5685 &mut worker_thread_ids,
5686 exact_score,
5687 #[cfg(any(test, feature = "test-support"))]
5688 vector_fault_controller,
5689 #[cfg(any(test, feature = "test-support"))]
5690 vector_fault_tier(options.tier()),
5691 )?;
5692 let mut plan = crate::planner::SegmentPlan::unfiltered_scan(
5693 RowSource::Active,
5694 crate::planner::SegmentTier::ActiveScan,
5695 alive.live_count(),
5696 crate::planner::ScanReason::ActiveSegment,
5697 );
5698 plan.approximate = approximate;
5699 plans.push(plan);
5700 if matches!(graph_bound_mode, GraphBoundMode::Shared) {
5701 retain_global_top_k(&mut candidates, k);
5702 }
5703 }
5704
5705 let mut ordered_segments = snapshot.segments().iter().collect::<Vec<_>>();
5706 if matches!(graph_bound_mode, GraphBoundMode::Shared) {
5707 ordered_segments.sort_unstable_by(|left, right| {
5711 right
5712 .meta()
5713 .row_count
5714 .cmp(&left.meta().row_count)
5715 .then_with(|| left.meta().id.cmp(&right.meta().id))
5716 });
5717 }
5718
5719 for segment in ordered_segments {
5720 let alive = segment.query_alive().map_err(QueryError::Store)?;
5721 let source = RowSource::Sealed(segment.meta().id);
5722 let graph_options = match options.tier() {
5727 SearchTier::Graph(graph_options) => Some(graph_options),
5728 SearchTier::Auto
5729 if segment.directory().iter().any(|entry| {
5730 entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id()
5731 }) =>
5732 {
5733 auto_graph_options
5734 }
5735 SearchTier::Auto | SearchTier::Exact | SearchTier::Scan => None,
5736 };
5737 if let Some(graph_options) = graph_options {
5738 worst_exhaustive = false;
5739 let competitive_distance = global_competitive_distance(&candidates, k);
5740 let MergedSegmentGraph {
5741 plan,
5742 traversed,
5743 _result,
5744 _scratch,
5745 _lease,
5746 } = traverse_segment_graph(
5747 segment,
5748 &alive,
5749 source,
5750 graph_options,
5751 graph_bound_mode,
5752 request,
5753 k,
5754 vector_preparation,
5755 accounting,
5756 snapshot,
5757 generation,
5758 &control,
5759 competitive_distance,
5760 #[cfg(any(test, feature = "test-support"))]
5761 vector_fault_controller,
5762 options.tier(),
5763 )?
5764 .merge_into(
5765 segment,
5766 &alive,
5767 source,
5768 &mut candidates,
5769 &mut dims_touched,
5770 &mut bytes_read,
5771 &mut worker_thread_ids,
5772 &mut graph_stats,
5773 #[cfg(any(test, feature = "test-support"))]
5774 vector_fault_controller,
5775 options.tier(),
5776 )?;
5777 plans.push(plan);
5778 if traversed && matches!(graph_bound_mode, GraphBoundMode::Shared) {
5779 retain_global_top_k(&mut candidates, k);
5780 }
5781 continue;
5782 }
5783
5784 let mut exact_memory = ExactScanMemory::new(accounting)?;
5785 let frontier = rescore_options
5786 .map(|options| options.frontier(k, alive.live_count()))
5787 .transpose()?
5788 .unwrap_or(k);
5789 let mut outcome = scan_sealed_segment(
5790 pool,
5791 snapshot,
5792 segment,
5793 &alive,
5794 generation,
5795 request,
5796 frontier,
5797 scan_options,
5798 &control,
5799 &bit4_query,
5800 &int8_query,
5801 full_precision,
5802 exact_lane,
5803 accounting,
5804 &mut exact_memory,
5805 source,
5806 options.tier(),
5807 #[cfg(any(test, feature = "test-support"))]
5808 vector_fault_controller,
5809 )?;
5810 let mut exact_score = full_precision || segment.meta().scheme == 0;
5811 let mut approximate = false;
5812 if let Some(rescore) = rescore_options {
5813 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
5814 let cancellation = QueryCancellation::new(&control, &lease);
5815 let vectors = exact_rescore_rows_for_search(
5816 segment,
5817 #[cfg(any(test, feature = "test-support"))]
5818 vector_fault_controller,
5819 #[cfg(any(test, feature = "test-support"))]
5820 vector_fault_tier(options.tier()),
5821 )?;
5822 let (counters, exhaustive, worst) = rescored_scan::rescore(
5823 &mut outcome.candidates,
5824 &mut outcome.stats,
5825 vectors,
5826 query,
5827 k,
5828 alive.live_count(),
5829 rescore,
5830 &cancellation,
5831 accounting,
5832 &mut exact_memory,
5833 )?;
5834 outcome.worst_score = worst;
5835 exact_score = true;
5836 approximate = !exhaustive;
5837 rescored_scan::accumulate(&mut rescore_counters, counters)?;
5838 }
5839 fold_worst_squared_l2(&mut worst_squared_l2, &mut worst_exhaustive, &outcome);
5840 merge_store_outcome(
5841 outcome,
5842 source,
5843 |row| {
5844 segment
5845 .document_version(row)
5846 .map_err(StoreError::Segment)
5847 .map_err(QueryError::Store)
5848 },
5849 &mut candidates,
5850 &mut dims_touched,
5851 &mut bytes_read,
5852 &mut worker_thread_ids,
5853 exact_score,
5854 #[cfg(any(test, feature = "test-support"))]
5855 vector_fault_controller,
5856 #[cfg(any(test, feature = "test-support"))]
5857 vector_fault_tier(options.tier()),
5858 )?;
5859 let tier = if segment
5860 .directory()
5861 .iter()
5862 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
5863 {
5864 crate::planner::SegmentTier::SealedGraph
5865 } else {
5866 crate::planner::SegmentTier::SealedScan
5867 };
5868 let mut plan = crate::planner::SegmentPlan::unfiltered_scan(
5869 source,
5870 tier,
5871 alive.live_count(),
5872 sealed_scan_reason(segment, requested_tier, hybrid_scan_reason),
5873 );
5874 plan.approximate = approximate;
5875 plans.push(plan);
5876 if matches!(graph_bound_mode, GraphBoundMode::Shared) {
5877 retain_global_top_k(&mut candidates, k);
5878 }
5879 }
5880
5881 retain_global_top_k(&mut candidates, k);
5882 let stats = ScanStats {
5883 dims_touched,
5884 bytes_read,
5885 threads_used: worker_thread_ids.len(),
5886 worker_thread_ids,
5887 };
5888 accounting.record_plans(&plans).map_err(QueryError::Store)?;
5889 let vector_ceiling = Some(exact_vector_ceiling(snapshot, active, query)?);
5892 let mut diagnostics = crate::diag::QueryDiagnostics::vector(crate::diag::VectorDiagnostics {
5893 snapshot_generation: generation,
5894 indexed_through_seq: active
5895 .indexed_through_seq()
5896 .max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
5897 approximate: plans.iter().any(|plan| plan.approximate),
5898 exact_rescore: candidates.iter().all(|candidate| candidate.exact_score()),
5899 requested_k: k,
5900 returned: candidates.len(),
5901 budget_exhausted: false,
5902 plan: plans,
5903 scan: stats.clone(),
5904 graph: graph_stats,
5905 epoch,
5906 elapsed: clock.now().saturating_duration_since(started),
5907 });
5908 diagnostics.scan_rescore = rescore_options.map(|_| rescore_counters);
5909 diagnostics.timings = crate::diag::QUERY_TIMING_ENABLED.then(|| crate::diag::QueryTimings {
5910 vector: timing_elapsed(clock, vector_started),
5911 ..Default::default()
5912 });
5913 Ok(SearchOutcome {
5914 candidates,
5915 vector_ceiling,
5920 stats,
5921 graph_stats,
5922 generation,
5923 epoch,
5924 diagnostics,
5925 })
5926}
5927
5928const GRAPH_WIDENING_CAP_NUMERATOR: usize = 1;
5936const GRAPH_WIDENING_CAP_DENOMINATOR: usize = 4;
5938
5939fn graph_round_is_worth_it(k: usize, corpus_rows: usize) -> bool {
5946 if corpus_rows == 0 {
5947 return true;
5948 }
5949 let ef = graph_round_ef(k);
5950 let cap =
5951 corpus_rows.saturating_mul(GRAPH_WIDENING_CAP_NUMERATOR) / GRAPH_WIDENING_CAP_DENOMINATOR;
5952 ef <= cap
5953}
5954
5955fn graph_round_ef(k: usize) -> usize {
5956 k.saturating_mul(4).max(140)
5957}
5958
5959fn sealed_scan_reason(
5960 segment: &crate::segment::reader::SegmentReader,
5961 requested_tier: Option<SearchTier>,
5962 hybrid_scan_reason: Option<crate::planner::ScanReason>,
5963) -> crate::planner::ScanReason {
5964 use crate::planner::{ExplicitScanTier, ScanReason};
5965
5966 let has_graph = segment
5967 .directory()
5968 .iter()
5969 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id());
5970 if has_graph && let Some(reason) = hybrid_scan_reason {
5971 return reason;
5972 }
5973 match requested_tier {
5974 Some(SearchTier::Exact) => return ScanReason::ExplicitTier(ExplicitScanTier::Exact),
5975 Some(SearchTier::Scan) => return ScanReason::ExplicitTier(ExplicitScanTier::Scan),
5976 Some(SearchTier::Auto | SearchTier::Graph(_)) | None => {}
5977 }
5978 let rows = segment.meta().row_count;
5979 let min_rows = crate::tier::thresholds::for_bucket(segment.meta().dims, segment.meta().scheme)
5980 .graph_min_rows;
5981 if rows < min_rows {
5982 ScanReason::BelowGraphThreshold { rows, min_rows }
5983 } else {
5984 ScanReason::GraphPending
5985 }
5986}
5987
5988fn snapshot_has_graph(snapshot: &PublishedSnapshot) -> bool {
5989 snapshot.segments().iter().any(|segment| {
5990 segment
5991 .directory()
5992 .iter()
5993 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
5994 })
5995}
5996
5997fn exact_vector_ceiling(
5998 snapshot: &PublishedSnapshot,
5999 active: &crate::ingest::ActiveSegment,
6000 query: &[f32],
6001) -> Result<f64, QueryError> {
6002 fold_exact_vector_ceiling(
6003 snapshot,
6004 active,
6005 query,
6006 remembered_segment_vector_ceiling_norm_range,
6007 )
6008}
6009
6010fn remembered_segment_vector_ceiling_norm_range(
6017 segment: &crate::segment::reader::SegmentReader,
6018) -> Result<crate::graph::search::GraphSegmentNormRange, QueryError> {
6019 if let Some(range) = segment.cached_vector_ceiling_norm_range() {
6020 return Ok(range);
6021 }
6022 let range = compute_segment_vector_ceiling_norm_range(segment)?;
6023 segment.remember_vector_ceiling_norm_range(range);
6024 Ok(range)
6025}
6026
6027fn fold_exact_vector_ceiling(
6033 snapshot: &PublishedSnapshot,
6034 active: &crate::ingest::ActiveSegment,
6035 query: &[f32],
6036 mut segment_range: impl FnMut(
6037 &crate::segment::reader::SegmentReader,
6038 )
6039 -> Result<crate::graph::search::GraphSegmentNormRange, QueryError>,
6040) -> Result<f64, QueryError> {
6041 let mut ceiling = 0.0_f64;
6042 if !active.is_empty() {
6043 ceiling = fold_vector_ceiling(
6046 ceiling,
6047 crate::graph::search::GraphSegmentNormRange::from_factors(active.factors()),
6048 query,
6049 )?;
6050 }
6051 for segment in snapshot.segments() {
6052 let range = segment_range(segment)?;
6053 ceiling = fold_vector_ceiling(ceiling, range, query)?;
6054 }
6055 Ok(ceiling)
6056}
6057
6058fn compute_segment_vector_ceiling_norm_range(
6060 segment: &crate::segment::reader::SegmentReader,
6061) -> Result<crate::graph::search::GraphSegmentNormRange, QueryError> {
6062 match segment.meta().scheme {
6063 4 => Ok(crate::graph::search::GraphSegmentNormRange::from_factors(
6064 segment
6065 .query_bit4_factors()
6066 .map_err(StoreError::Segment)
6067 .map_err(QueryError::Store)?,
6068 )),
6069 0 | 2 => Ok(
6070 crate::graph::search::GraphSegmentNormRange::from_exact_rows(
6071 segment
6072 .query_rescore_f32()
6073 .map_err(StoreError::Segment)
6074 .map_err(QueryError::Store)?,
6075 segment.meta().dims as usize,
6076 ),
6077 ),
6078 scheme => Err(QueryError::Store(StoreError::Segment(
6079 crate::segment::SegmentError::Geometry(format!(
6080 "vector norm enclosure does not support scheme {scheme}"
6081 )),
6082 ))),
6083 }
6084}
6085
6086fn fold_vector_ceiling(
6087 current: f64,
6088 range: crate::graph::search::GraphSegmentNormRange,
6089 query: &[f32],
6090) -> Result<f64, QueryError> {
6091 let ceiling = range.squared_l2_upper_bound(query);
6092 if !ceiling.is_finite() {
6093 return Err(QueryError::Store(StoreError::Segment(
6094 crate::segment::SegmentError::Geometry(
6095 "vector norm enclosure did not produce a finite squared-L2 ceiling".to_owned(),
6096 ),
6097 )));
6098 }
6099 Ok(current.max(ceiling))
6100}
6101
6102fn fold_worst_squared_l2(
6106 worst: &mut Option<f64>,
6107 exhaustive: &mut bool,
6108 outcome: &crate::scan::ScanOutcome,
6109) {
6110 match outcome.worst_score {
6111 Some(score) => {
6112 let squared_l2 = -f64::from(score);
6113 *worst = Some(worst.map_or(squared_l2, |current: f64| current.max(squared_l2)));
6114 }
6115 None => {
6116 if !outcome.candidates.is_empty() {
6117 *exhaustive = false;
6118 }
6119 }
6120 }
6121}
6122
6123#[allow(clippy::large_enum_variant)]
6125enum SegmentGraphResult<'a> {
6126 Pruned {
6127 plan: crate::planner::SegmentPlan,
6128 graph_stats: crate::ingest::GraphSearchStats,
6129 _lease: SnapshotLease,
6130 },
6131 Traversed {
6132 result: crate::graph::search::GraphSearchResult,
6133 plan: crate::planner::SegmentPlan,
6134 dims_touched: u64,
6135 bytes_read: u64,
6136 caller_thread: std::thread::ThreadId,
6137 graph_stats: crate::ingest::GraphSearchStats,
6138 _scratch: graph_cache::GraphScratchLease<'a>,
6139 _lease: SnapshotLease,
6140 },
6141}
6142
6143struct MergedSegmentGraph<'a> {
6144 plan: crate::planner::SegmentPlan,
6145 traversed: bool,
6146 _result: Option<crate::graph::search::GraphSearchResult>,
6147 _scratch: Option<graph_cache::GraphScratchLease<'a>>,
6148 _lease: SnapshotLease,
6149}
6150
6151impl<'a> SegmentGraphResult<'a> {
6152 #[allow(clippy::too_many_arguments)]
6153 fn merge_into(
6154 self,
6155 segment: &crate::segment::reader::SegmentReader,
6156 alive: &crate::meta::AliveSet,
6157 source: crate::ingest::RowSource,
6158 candidates: &mut Vec<crate::ingest::SearchCandidate>,
6159 dims_touched: &mut u64,
6160 bytes_read: &mut u64,
6161 worker_thread_ids: &mut Vec<std::thread::ThreadId>,
6162 graph_stats: &mut crate::ingest::GraphSearchStats,
6163 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
6164 &crate::scan::vector_fault::VectorFaultController,
6165 >,
6166 _tier: SearchTier,
6167 ) -> Result<MergedSegmentGraph<'a>, QueryError> {
6168 match self {
6169 Self::Pruned {
6170 plan,
6171 graph_stats: delta,
6172 _lease,
6173 } => {
6174 add_traversal_stats(graph_stats, delta)?;
6175 Ok(MergedSegmentGraph {
6176 plan,
6177 traversed: false,
6178 _result: None,
6179 _scratch: None,
6180 _lease,
6181 })
6182 }
6183 Self::Traversed {
6184 result,
6185 plan,
6186 dims_touched: traversal_dims_touched,
6187 bytes_read: traversal_bytes_read,
6188 caller_thread,
6189 graph_stats: delta,
6190 _scratch,
6191 _lease,
6192 } => {
6193 *dims_touched = dims_touched
6194 .checked_add(traversal_dims_touched)
6195 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6196 *bytes_read = bytes_read
6197 .checked_add(traversal_bytes_read)
6198 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6199 if !worker_thread_ids.contains(&caller_thread) {
6200 worker_thread_ids.push(caller_thread);
6201 }
6202 add_traversal_stats(graph_stats, delta)?;
6203 reserve_global_candidates(
6204 candidates,
6205 result
6206 .candidates()
6207 .iter()
6208 .filter(|candidate| alive.is_alive(candidate.row_id()))
6209 .count(),
6210 #[cfg(any(test, feature = "test-support"))]
6211 vector_fault_controller,
6212 #[cfg(any(test, feature = "test-support"))]
6213 vector_fault_tier(_tier),
6214 )?;
6215 for candidate in result.candidates() {
6216 if !alive.is_alive(candidate.row_id()) {
6217 continue;
6218 }
6219 let score = -(candidate.distance() as f32);
6220 if !score.is_finite() {
6221 return Err(QueryError::Graph(
6222 crate::graph::search::GraphSearchError::Geometry(format!(
6223 "graph result distance for row {} is not representable as f32",
6224 candidate.row_id()
6225 )),
6226 ));
6227 }
6228 candidates.push(crate::ingest::SearchCandidate::new(
6229 crate::ingest::GlobalRowId::new(source, candidate.row_id()),
6230 segment
6231 .document_version(candidate.row_id() as usize)
6232 .map_err(StoreError::Segment)
6233 .map_err(QueryError::Store)?,
6234 score,
6235 true,
6236 ));
6237 }
6238 Ok(MergedSegmentGraph {
6239 plan,
6240 traversed: true,
6241 _result: Some(result),
6242 _scratch: Some(_scratch),
6243 _lease,
6244 })
6245 }
6246 }
6247 }
6248}
6249
6250#[allow(clippy::too_many_arguments)]
6251fn traverse_segment_graph<'a>(
6252 segment: &'a crate::segment::reader::SegmentReader,
6253 alive: &crate::meta::AliveSet,
6254 source: crate::ingest::RowSource,
6255 graph_options: GraphSearchOptions,
6256 graph_bound_mode: GraphBoundMode,
6257 request: crate::ingest::SearchRequest<'_>,
6258 k: usize,
6259 vector_preparation: &mut prepared::PreparedVectorQuery<'_>,
6260 accounting: &Arc<stats::Accounting>,
6261 snapshot: &Arc<PublishedSnapshot>,
6262 generation: u64,
6263 control: &QueryControl,
6264 competitive_distance: Option<f32>,
6265 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
6266 &crate::scan::vector_fault::VectorFaultController,
6267 >,
6268 _tier: SearchTier,
6269) -> Result<SegmentGraphResult<'a>, QueryError> {
6270 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
6271 let cancellation = QueryCancellation::new(control, &lease);
6272 cancellation.check_graph().map_err(map_scan_error)?;
6273 let (graph, graph_validated, prepared_entry_seed_discovered, norm_range) =
6274 match graph_bound_mode {
6275 GraphBoundMode::Shared => {
6276 let prepared = segment
6277 .graph_search_cache
6278 .prepare_shared(segment, &cancellation)
6279 .map_err(map_graph_cache_error)?;
6280 (
6281 prepared.graph,
6282 prepared.graph_validated,
6283 prepared.entry_seed_discovered,
6284 Some(prepared.norm_range),
6285 )
6286 }
6287 #[cfg(test)]
6288 GraphBoundMode::Independent => {
6289 let (graph, graph_validated) = segment
6290 .graph_search_cache
6291 .bind_graph(segment)
6292 .map_err(map_graph_cache_error)?;
6293 (graph, graph_validated, false, None)
6294 }
6295 };
6296 let rescore = query_rescore_rows_for_search(
6297 segment,
6298 #[cfg(any(test, feature = "test-support"))]
6299 vector_fault_controller,
6300 #[cfg(any(test, feature = "test-support"))]
6301 vector_fault_tier(_tier),
6302 )?;
6303 let node_count = graph.node_count() as usize;
6304 let segment_k = k.min(node_count);
6305 let has_tombstones = alive.tombstone_count() != 0;
6306 let target_live_k = if has_tombstones {
6307 usize::try_from(alive.live_count())
6308 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?
6309 .min(k)
6310 } else {
6311 segment_k
6312 };
6313 let maximum_candidate_k = if has_tombstones {
6314 let tombstone_count = usize::try_from(alive.tombstone_count())
6315 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6316 target_live_k
6317 .checked_add(tombstone_count)
6318 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?
6319 .min(node_count)
6320 .max(segment_k)
6321 } else {
6322 segment_k
6323 };
6324 if let Some(ef) = graph_options.ef()
6325 && ef < segment_k
6326 {
6327 return Err(QueryError::Graph(
6328 crate::graph::search::GraphSearchError::AdaptiveEf(
6329 crate::graph::search::AdaptiveEfError::ExplicitBelowK { k: segment_k, ef },
6330 ),
6331 ));
6332 }
6333 let request_for = |candidate_k| {
6334 let graph_request = crate::graph::search::GraphSearchRequest::new(
6335 request.vector(),
6336 candidate_k,
6337 graph_options.seed(),
6338 )
6339 .with_profile(graph_options.profile());
6340 graph_options.ef().map_or(graph_request, |ef| {
6345 graph_request.with_ef(ef.max(candidate_k))
6346 })
6347 };
6348 let mut candidate_k = segment_k;
6349 let mut graph_request = request_for(candidate_k);
6350 let effective_ef = graph_request
6351 .effective_ef(graph.node_count() as usize)
6352 .map_err(crate::graph::search::GraphSearchError::AdaptiveEf)
6353 .map_err(QueryError::Graph)?;
6354 let requestable_max_k = maximum_candidate_k;
6355 let scratch_ef = if has_tombstones && requestable_max_k > candidate_k {
6356 request_for(requestable_max_k)
6357 .effective_ef(node_count)
6358 .map_err(crate::graph::search::GraphSearchError::AdaptiveEf)
6359 .map_err(QueryError::Graph)?
6360 } else {
6361 effective_ef
6362 };
6363 if let (Some(norm_range), Some(competitive_distance)) = (norm_range, competitive_distance)
6364 && norm_range.squared_l2_upper_bound(request.vector()) <= f64::from(f32::MAX)
6365 && (norm_range.squared_l2_lower_bound(request.vector()) as f32) > competitive_distance
6366 {
6367 cancellation.check_graph().map_err(map_scan_error)?;
6368 let graph_stats = crate::ingest::GraphSearchStats {
6369 graph_validations: usize::from(graph_validated),
6370 entry_seed_discoveries: usize::from(prepared_entry_seed_discovered),
6371 segments_pruned_by_bound: 1,
6372 ..crate::ingest::GraphSearchStats::default()
6373 };
6374 let plan = crate::planner::SegmentPlan::unfiltered_pruned(source, alive.live_count());
6375 return Ok(SegmentGraphResult::Pruned {
6376 plan,
6377 graph_stats,
6378 _lease: lease,
6379 });
6380 }
6381 let mut scratch = segment
6382 .graph_search_cache
6383 .checkout(graph, scratch_ef, accounting, &cancellation)
6384 .map_err(map_graph_cache_error)?;
6385 let entry_seed_discovered = prepared_entry_seed_discovered || scratch.entry_seed_discovered();
6386 let entries = scratch.entries();
6387 let mut searcher = crate::graph::search::GraphSearcher::with_entry_row_ids(
6388 graph,
6389 rescore,
6390 entries,
6391 scratch.scratch_mut().map_err(map_graph_error)?,
6392 )
6393 .map_err(map_graph_error)?
6394 .with_rescore_validator(segment);
6395 #[cfg(any(test, feature = "test-support"))]
6396 if let Some(after_hops) = graph_options.cancel_after_hops() {
6397 let token = cancellation.cancel_token_for_test().ok_or_else(|| {
6398 QueryError::Graph(crate::graph::search::GraphSearchError::Geometry(
6399 "test hop cancellation requires a Cancel query control".to_owned(),
6400 ))
6401 })?;
6402 let _observed_hops = searcher.cancel_after_hops(after_hops, token);
6403 }
6404 #[cfg(any(test, feature = "test-support"))]
6405 if let Some(controller) = vector_fault_controller {
6406 searcher = searcher.with_vector_fault_controller(
6407 controller,
6408 vector_fault_source(source),
6409 crate::scan::vector_fault::VectorSearchTier::Graph,
6410 );
6411 }
6412 let mut traversal_dims_touched = 0_u64;
6413 let mut traversal_bytes_read = 0_u64;
6414 let mut traversal_epoch_clears = 0_usize;
6415 let mut traversal_candidates_scored = 0_usize;
6416 let mut traversal_candidates_rescored = 0_usize;
6417 let result = loop {
6418 let result = searcher.search_reusing(
6419 graph_request,
6420 Some(&cancellation),
6421 vector_preparation,
6422 segment,
6423 source,
6424 )?;
6425 let counters = result.counters();
6426 traversal_dims_touched = traversal_dims_touched
6427 .checked_add(counters.dims_touched())
6428 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6429 traversal_bytes_read = traversal_bytes_read
6430 .checked_add(counters.bytes_read())
6431 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6432 traversal_epoch_clears = traversal_epoch_clears
6433 .checked_add(usize::from(counters.visited_epoch_cleared()))
6434 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6435 traversal_candidates_scored = traversal_candidates_scored
6436 .checked_add(counters.candidates_scored())
6437 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6438 traversal_candidates_rescored = traversal_candidates_rescored
6439 .checked_add(counters.candidates_rescored())
6440 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6441
6442 if !has_tombstones {
6443 break result;
6444 }
6445 let retained_live = result
6446 .candidates()
6447 .iter()
6448 .filter(|candidate| alive.is_alive(candidate.row_id()))
6449 .count();
6450 if retained_live >= target_live_k {
6451 break result;
6452 }
6453 if candidate_k >= requestable_max_k {
6454 return Err(QueryError::Graph(
6455 crate::graph::search::GraphSearchError::Geometry(format!(
6456 "graph traversal exhausted {candidate_k} candidates but retained {retained_live} live rows, expected {target_live_k}"
6457 )),
6458 ));
6459 }
6460 let widened_k = candidate_k
6461 .saturating_mul(2)
6462 .max(candidate_k.saturating_add(1))
6463 .min(requestable_max_k);
6464 if widened_k <= candidate_k {
6465 return Err(QueryError::Graph(
6466 crate::graph::search::GraphSearchError::Geometry(format!(
6467 "graph traversal could not widen beyond {candidate_k} candidates"
6468 )),
6469 ));
6470 }
6471 candidate_k = widened_k;
6472 graph_request = request_for(candidate_k);
6473 };
6474 let graph_stats = crate::ingest::GraphSearchStats {
6475 segments_traversed: 1,
6476 graph_validations: usize::from(graph_validated),
6477 entry_seed_discoveries: usize::from(entry_seed_discovered),
6478 visited_epoch_clears: traversal_epoch_clears,
6479 candidates_scored: traversal_candidates_scored,
6480 candidates_rescored: traversal_candidates_rescored,
6481 segments_pruned_by_bound: 0,
6482 };
6483 let plan = crate::planner::SegmentPlan::unfiltered_graph(
6484 source,
6485 alive.live_count(),
6486 graph_options.ef(),
6487 effective_ef,
6488 graph_options.profile(),
6489 );
6490 let caller_thread = std::thread::current().id();
6491 drop(searcher);
6492 Ok(SegmentGraphResult::Traversed {
6493 result,
6494 plan,
6495 dims_touched: traversal_dims_touched,
6496 bytes_read: traversal_bytes_read,
6497 caller_thread,
6498 graph_stats,
6499 _scratch: scratch,
6500 _lease: lease,
6501 })
6502}
6503
6504fn add_traversal_stats(
6505 stats: &mut crate::ingest::GraphSearchStats,
6506 delta: crate::ingest::GraphSearchStats,
6507) -> Result<(), QueryError> {
6508 stats.segments_traversed = stats
6509 .segments_traversed
6510 .checked_add(delta.segments_traversed)
6511 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6512 stats.graph_validations = stats
6513 .graph_validations
6514 .checked_add(delta.graph_validations)
6515 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6516 stats.entry_seed_discoveries = stats
6517 .entry_seed_discoveries
6518 .checked_add(delta.entry_seed_discoveries)
6519 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6520 stats.visited_epoch_clears = stats
6521 .visited_epoch_clears
6522 .checked_add(delta.visited_epoch_clears)
6523 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6524 stats.candidates_scored = stats
6525 .candidates_scored
6526 .checked_add(delta.candidates_scored)
6527 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6528 stats.candidates_rescored = stats
6529 .candidates_rescored
6530 .checked_add(delta.candidates_rescored)
6531 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6532 stats.segments_pruned_by_bound = stats
6533 .segments_pruned_by_bound
6534 .checked_add(delta.segments_pruned_by_bound)
6535 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6536 Ok(())
6537}
6538
6539#[allow(clippy::too_many_arguments)]
6540fn scan_sealed_segment(
6541 pool: Option<&pool::QueryPool>,
6542 snapshot: &Arc<PublishedSnapshot>,
6543 segment: &crate::segment::reader::SegmentReader,
6544 alive: &crate::meta::AliveSet,
6545 generation: u64,
6546 request: crate::ingest::SearchRequest<'_>,
6547 k: usize,
6548 scan_options: crate::scan::ScanOptions,
6549 control: &QueryControl,
6550 bit4_query: &std::cell::OnceCell<Result<crate::quant::Bit4Query, crate::quant::QuantError>>,
6551 int8_query: &std::cell::OnceCell<Result<crate::quant::Int8Query, crate::quant::QuantError>>,
6552 full_precision: bool,
6553 exact_lane: usize,
6554 accounting: &Arc<stats::Accounting>,
6555 exact_memory: &mut ExactScanMemory,
6556 _source: crate::ingest::RowSource,
6557 _tier: SearchTier,
6558 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
6559 &crate::scan::vector_fault::VectorFaultController,
6560 >,
6561) -> Result<crate::scan::ScanOutcome, QueryError> {
6562 use crate::scan::{ScanQuery, ScanRequest, ScanRows};
6563
6564 if full_precision {
6565 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
6566 let vectors = exact_rescore_rows_for_search(
6567 segment,
6568 #[cfg(any(test, feature = "test-support"))]
6569 vector_fault_controller,
6570 #[cfg(any(test, feature = "test-support"))]
6571 vector_fault_tier(_tier),
6572 )?;
6573 return scan_squared_l2(
6574 vectors,
6575 segment.meta().row_count as usize,
6576 alive,
6577 request.vector(),
6578 k,
6579 ExactScanDispatch {
6580 pool,
6581 accounting,
6582 control,
6583 lease: &lease,
6584 thread_budget: scan_options.thread_budget,
6585 lane: exact_lane,
6586 },
6587 exact_memory,
6588 #[cfg(any(test, feature = "test-support"))]
6589 vector_fault_controller,
6590 #[cfg(any(test, feature = "test-support"))]
6591 vector_fault_source(_source),
6592 #[cfg(any(test, feature = "test-support"))]
6593 vector_fault_tier(_tier),
6594 );
6595 }
6596
6597 let query_pool = pool.ok_or(QueryError::Store(StoreError::Synchronization {
6598 component: "sealed scan-tier query pool",
6599 }))?;
6600 match segment.meta().scheme {
6601 0 => execute_store_scan(
6602 query_pool,
6603 ScanRequest {
6604 query: ScanQuery::F32(request.vector()),
6605 rows: ScanRows::F32BorrowedRowMajor(
6606 segment
6607 .query_f32_codes()
6608 .map_err(StoreError::Segment)
6609 .map_err(QueryError::Store)?,
6610 ),
6611 row_mask: Some(alive.scan_mask()),
6612 },
6613 k,
6614 scan_options,
6615 control.clone(),
6616 SnapshotLease::new_at(Arc::clone(snapshot), generation),
6617 #[cfg(any(test, feature = "test-support"))]
6618 vector_fault_controller,
6619 #[cfg(any(test, feature = "test-support"))]
6620 vector_fault_source(_source),
6621 #[cfg(any(test, feature = "test-support"))]
6622 vector_fault_tier(_tier),
6623 ),
6624 4 => execute_store_scan(
6625 query_pool,
6626 ScanRequest {
6627 query: ScanQuery::Bit4(
6628 bit4_query
6629 .get_or_init(|| crate::quant::prepare_bit4_query(request.vector(), 0))
6630 .as_ref()
6631 .map_err(|error| crate::scan::ScanError::Quant(error.clone()))
6632 .map_err(QueryError::Scan)?,
6633 ),
6634 rows: ScanRows::Bit4RowMajor {
6635 codes: segment
6636 .query_bit4_codes()
6637 .map_err(StoreError::Segment)
6638 .map_err(QueryError::Store)?,
6639 factors: segment
6640 .query_bit4_factors()
6641 .map_err(StoreError::Segment)
6642 .map_err(QueryError::Store)?,
6643 },
6644 row_mask: Some(alive.scan_mask()),
6645 },
6646 k,
6647 scan_options,
6648 control.clone(),
6649 SnapshotLease::new_at(Arc::clone(snapshot), generation),
6650 #[cfg(any(test, feature = "test-support"))]
6651 vector_fault_controller,
6652 #[cfg(any(test, feature = "test-support"))]
6653 vector_fault_source(_source),
6654 #[cfg(any(test, feature = "test-support"))]
6655 vector_fault_tier(_tier),
6656 ),
6657 2 => {
6658 let factors = segment.query_int8_factors().map_err(QueryError::Store)?;
6659 execute_store_scan(
6660 query_pool,
6661 ScanRequest {
6662 query: ScanQuery::Int8(
6663 int8_query
6664 .get_or_init(|| crate::quant::prepare_int8_query(request.vector()))
6665 .as_ref()
6666 .map_err(|error| crate::scan::ScanError::Quant(error.clone()))
6667 .map_err(QueryError::Scan)?,
6668 ),
6669 rows: ScanRows::Int8RowMajor {
6670 codes: segment
6671 .query_int8_codes()
6672 .map_err(StoreError::Segment)
6673 .map_err(QueryError::Store)?,
6674 factors: factors.as_slice(),
6675 },
6676 row_mask: Some(alive.scan_mask()),
6677 },
6678 k,
6679 scan_options,
6680 control.clone(),
6681 SnapshotLease::new_at(Arc::clone(snapshot), generation),
6682 #[cfg(any(test, feature = "test-support"))]
6683 vector_fault_controller,
6684 #[cfg(any(test, feature = "test-support"))]
6685 vector_fault_source(_source),
6686 #[cfg(any(test, feature = "test-support"))]
6687 vector_fault_tier(_tier),
6688 )
6689 }
6690 scheme => Err(QueryError::Store(StoreError::Segment(
6691 crate::segment::SegmentError::Geometry(format!(
6692 "store search does not support sealed scheme {scheme}"
6693 )),
6694 ))),
6695 }
6696}
6697
6698#[allow(clippy::too_many_arguments)]
6699fn execute_store_scan(
6700 pool: &pool::QueryPool,
6701 request: crate::scan::ScanRequest<'_>,
6702 k: usize,
6703 options: crate::scan::ScanOptions,
6704 control: QueryControl,
6705 lease: SnapshotLease,
6706 #[cfg(any(test, feature = "test-support"))] controller: Option<
6707 &crate::scan::vector_fault::VectorFaultController,
6708 >,
6709 #[cfg(any(test, feature = "test-support"))] source: crate::scan::vector_fault::VectorRowSource,
6710 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
6711) -> Result<crate::scan::ScanOutcome, QueryError> {
6712 #[cfg(any(test, feature = "test-support"))]
6713 if let Some(controller) = controller {
6714 let geometry = crate::scan::scan_geometry(request).map_err(QueryError::Scan)?;
6715 let cancellation = QueryCancellation::new(&control, &lease);
6716 let partition = crate::scan::scan_partition_with_vector_faults(
6717 request,
6718 k,
6719 0..geometry.row_count,
6720 Some(&cancellation),
6721 controller,
6722 source,
6723 tier,
6724 )
6725 .map_err(map_scan_error)?;
6726 return Ok(crate::scan::ScanOutcome {
6727 candidates: partition.candidates,
6728 worst_score: None,
6729 stats: crate::scan::ScanStats {
6730 dims_touched: partition.dims_touched,
6731 bytes_read: partition.bytes_read,
6732 threads_used: 1,
6733 worker_thread_ids: vec![partition.worker_thread_id],
6734 },
6735 });
6736 }
6737 pool.execute(request, k, options, control, lease)
6738}
6739
6740#[cfg(any(test, feature = "test-support"))]
6741fn vector_fault_tier(tier: SearchTier) -> crate::scan::vector_fault::VectorSearchTier {
6742 match tier {
6743 SearchTier::Exact => crate::scan::vector_fault::VectorSearchTier::Exact,
6744 SearchTier::Scan => crate::scan::vector_fault::VectorSearchTier::Scan,
6745 SearchTier::Graph(_) => crate::scan::vector_fault::VectorSearchTier::Graph,
6746 SearchTier::Auto => crate::scan::vector_fault::VectorSearchTier::Auto,
6747 }
6748}
6749
6750#[cfg(any(test, feature = "test-support"))]
6751fn vector_fault_source(
6752 source: crate::ingest::RowSource,
6753) -> crate::scan::vector_fault::VectorRowSource {
6754 match source {
6755 crate::ingest::RowSource::Active => crate::scan::vector_fault::VectorRowSource::Active,
6756 crate::ingest::RowSource::Sealed(segment) => {
6757 crate::scan::vector_fault::VectorRowSource::Sealed(*segment.as_bytes())
6758 }
6759 }
6760}
6761
6762pub(crate) fn auto_uses_full_precision(
6763 snapshot: &PublishedSnapshot,
6764 active: &crate::ingest::ActiveSegment,
6765) -> bool {
6766 let mut has_exact = false;
6767 let mut has_estimated = !active.is_empty();
6768 for segment in snapshot.segments() {
6769 if segment
6770 .directory()
6771 .iter()
6772 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
6773 {
6774 return true;
6775 }
6776 if segment.meta().row_count == 0 {
6777 continue;
6778 }
6779 match segment.meta().scheme {
6780 0 => has_exact = true,
6781 2 | 4 => has_estimated = true,
6782 _ => {}
6783 }
6784 }
6785 has_exact && has_estimated
6786}
6787
6788pub(crate) fn exact_rescore_rows(
6798 segment: &crate::segment::reader::SegmentReader,
6799) -> Result<&[f32], QueryError> {
6800 segment
6801 .query_rescore_f32()
6802 .map_err(|source| QueryError::Store(StoreError::Segment(source)))
6803}
6804
6805fn exact_rescore_rows_for_search<'a>(
6806 segment: &'a crate::segment::reader::SegmentReader,
6807 #[cfg(any(test, feature = "test-support"))] controller: Option<
6808 &crate::scan::vector_fault::VectorFaultController,
6809 >,
6810 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
6811) -> Result<&'a [f32], QueryError> {
6812 #[cfg(any(test, feature = "test-support"))]
6813 if let Some(available_rows) = controller.and_then(|controller| {
6814 controller.missing_rescore_rows(
6815 crate::scan::vector_fault::VectorRowSource::Sealed(*segment.meta().id.as_bytes()),
6816 tier,
6817 crate::scan::vector_fault::MissingRescoreSite::ExactRescoreRows,
6818 segment.meta().row_count,
6819 )
6820 }) {
6821 return Err(missing_rescore_rows_error(segment, available_rows));
6822 }
6823 exact_rescore_rows(segment)
6824}
6825
6826pub(crate) fn query_rescore_rows(
6827 segment: &crate::segment::reader::SegmentReader,
6828) -> Result<&[f32], QueryError> {
6829 segment
6830 .query_rescore_f32()
6831 .map_err(|source| QueryError::Store(StoreError::Segment(source)))
6832}
6833
6834fn query_rescore_rows_for_search<'a>(
6835 segment: &'a crate::segment::reader::SegmentReader,
6836 #[cfg(any(test, feature = "test-support"))] controller: Option<
6837 &crate::scan::vector_fault::VectorFaultController,
6838 >,
6839 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
6840) -> Result<&'a [f32], QueryError> {
6841 #[cfg(any(test, feature = "test-support"))]
6842 if let Some(available_rows) = controller.and_then(|controller| {
6843 controller.missing_rescore_rows(
6844 crate::scan::vector_fault::VectorRowSource::Sealed(*segment.meta().id.as_bytes()),
6845 tier,
6846 crate::scan::vector_fault::MissingRescoreSite::QueryRescoreRows,
6847 segment.meta().row_count,
6848 )
6849 }) {
6850 return Err(missing_rescore_rows_error(segment, available_rows));
6851 }
6852 query_rescore_rows(segment)
6853}
6854
6855#[cfg(any(test, feature = "test-support"))]
6856fn missing_rescore_rows_error(
6857 segment: &crate::segment::reader::SegmentReader,
6858 available_rows: u32,
6859) -> QueryError {
6860 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
6861 format!(
6862 "exact scores unavailable for segment {}: expected {} rows, got {available_rows}",
6863 segment.meta().id,
6864 segment.meta().row_count,
6865 ),
6866 )))
6867}
6868
6869fn retain_global_top_k(candidates: &mut Vec<crate::ingest::SearchCandidate>, k: usize) {
6870 candidates.sort_unstable_by(crate::ingest::compare_search_candidates);
6871 candidates.truncate(k);
6872}
6873
6874fn global_competitive_distance(
6875 candidates: &[crate::ingest::SearchCandidate],
6876 k: usize,
6877) -> Option<f32> {
6878 if k == 0 || candidates.len() < k {
6879 return None;
6880 }
6881 candidates
6882 .get(k.saturating_sub(1))
6883 .map(|candidate| -candidate.score())
6884 .filter(|distance| distance.is_finite() && *distance >= 0.0)
6885}
6886
6887fn map_graph_cache_error(error: graph_cache::GraphCacheError) -> QueryError {
6888 match error {
6889 graph_cache::GraphCacheError::Store(error) => QueryError::Store(error),
6890 graph_cache::GraphCacheError::Search(error) => map_graph_error(error),
6891 }
6892}
6893
6894pub(crate) fn map_graph_error(error: crate::graph::search::GraphSearchError) -> QueryError {
6895 match error {
6896 crate::graph::search::GraphSearchError::ExactRescoreUnavailable(detail) => {
6897 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
6898 detail,
6899 )))
6900 }
6901 crate::graph::search::GraphSearchError::Cancelled { partial } => {
6902 QueryError::Cancelled { partial }
6903 }
6904 crate::graph::search::GraphSearchError::Timeout { partial } => {
6905 QueryError::Timeout { partial }
6906 }
6907 crate::graph::search::GraphSearchError::ReadCancelled { partial } => {
6908 QueryError::ReadCancelled { partial }
6909 }
6910 crate::graph::search::GraphSearchError::Scan(error) => QueryError::Scan(error),
6911 error => QueryError::Graph(error),
6912 }
6913}
6914
6915#[allow(
6916 clippy::too_many_arguments,
6917 reason = "scoped dispatch and test source identity stay explicit"
6918)]
6919fn scan_active_squared_l2(
6920 active: &crate::ingest::ActiveSegment,
6921 alive: &crate::meta::AliveSet,
6922 query: &[f32],
6923 k: usize,
6924 dispatch: ExactScanDispatch<'_>,
6925 memory: &mut ExactScanMemory,
6926 #[cfg(any(test, feature = "test-support"))] controller: Option<
6927 &crate::scan::vector_fault::VectorFaultController,
6928 >,
6929 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
6930) -> Result<crate::scan::ScanOutcome, QueryError> {
6931 scan_squared_l2(
6932 active.vectors(),
6933 active.row_count(),
6934 alive,
6935 query,
6936 k,
6937 dispatch,
6938 memory,
6939 #[cfg(any(test, feature = "test-support"))]
6940 controller,
6941 #[cfg(any(test, feature = "test-support"))]
6942 crate::scan::vector_fault::VectorRowSource::Active,
6943 #[cfg(any(test, feature = "test-support"))]
6944 tier,
6945 )
6946}
6947
6948#[derive(Debug)]
6949enum ExactScanError {
6950 Query(QueryError),
6951 Narrowing { row_id: usize, score: f64 },
6952}
6953
6954impl From<QueryError> for ExactScanError {
6955 fn from(error: QueryError) -> Self {
6956 Self::Query(error)
6957 }
6958}
6959
6960impl ExactScanError {
6961 fn into_query(self) -> QueryError {
6962 match self {
6963 Self::Query(error) => error,
6964 Self::Narrowing { row_id, .. } => {
6965 QueryError::Scan(crate::scan::ScanError::NonFiniteScore { row_id })
6966 }
6967 }
6968 }
6969}
6970
6971pub(crate) struct ExactScanMemory {
6972 candidates: stats::AccountedCounter,
6973 worker_ids: stats::AccountedCounter,
6974}
6975
6976impl ExactScanMemory {
6977 pub(crate) fn new(accounting: &Arc<stats::Accounting>) -> Result<Self, QueryError> {
6978 Ok(Self {
6979 candidates: stats::AccountedCounter::new(
6980 accounting,
6981 stats::AllocationComponent::Temporary,
6982 )
6983 .map_err(QueryError::Store)?,
6984 worker_ids: stats::AccountedCounter::new(
6985 accounting,
6986 stats::AllocationComponent::Temporary,
6987 )
6988 .map_err(QueryError::Store)?,
6989 })
6990 }
6991}
6992
6993struct ExactScanDispatch<'a> {
6994 pool: Option<&'a pool::QueryPool>,
6995 accounting: &'a Arc<stats::Accounting>,
6996 control: &'a QueryControl,
6997 lease: &'a SnapshotLease,
6998 thread_budget: usize,
6999 lane: usize,
7000}
7001
7002#[allow(
7003 clippy::too_many_arguments,
7004 reason = "scoped dispatch and test source identity stay explicit"
7005)]
7006fn scan_squared_l2(
7007 vectors: &[f32],
7008 row_count: usize,
7009 alive: &crate::meta::AliveSet,
7010 query: &[f32],
7011 k: usize,
7012 dispatch: ExactScanDispatch<'_>,
7013 memory: &mut ExactScanMemory,
7014 #[cfg(any(test, feature = "test-support"))] controller: Option<
7015 &crate::scan::vector_fault::VectorFaultController,
7016 >,
7017 #[cfg(any(test, feature = "test-support"))] source: crate::scan::vector_fault::VectorRowSource,
7018 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7019) -> Result<crate::scan::ScanOutcome, QueryError> {
7020 if let Some(pool) = dispatch.pool
7021 && pool.exact_workers(row_count, query.len(), dispatch.thread_budget) > 1
7022 {
7023 return pool.execute_exact(
7024 vectors,
7025 row_count,
7026 alive,
7027 query,
7028 k,
7029 dispatch.thread_budget,
7030 dispatch.lane,
7031 dispatch.control,
7032 dispatch.lease,
7033 dispatch.accounting,
7034 memory,
7035 #[cfg(any(test, feature = "test-support"))]
7036 controller,
7037 #[cfg(any(test, feature = "test-support"))]
7038 source,
7039 #[cfg(any(test, feature = "test-support"))]
7040 tier,
7041 );
7042 }
7043 let cancellation = QueryCancellation::new(dispatch.control, dispatch.lease);
7044 #[cfg(any(test, feature = "test-support"))]
7045 if controller.is_some_and(|controller| controller.has_exact_row_hooks()) {
7046 return scan_exact_partition::<false>(
7047 vectors,
7048 row_count,
7049 0..row_count,
7050 alive,
7051 query,
7052 k,
7053 &cancellation,
7054 memory,
7055 #[cfg(any(test, feature = "test-support"))]
7056 controller,
7057 #[cfg(any(test, feature = "test-support"))]
7058 source,
7059 #[cfg(any(test, feature = "test-support"))]
7060 tier,
7061 )
7062 .map_err(ExactScanError::into_query);
7063 }
7064 scan_serial_squared_l2(
7065 vectors,
7066 row_count,
7067 alive,
7068 query,
7069 k,
7070 &cancellation,
7071 &mut memory.candidates,
7072 &mut memory.worker_ids,
7073 #[cfg(any(test, feature = "test-support"))]
7074 controller,
7075 #[cfg(any(test, feature = "test-support"))]
7076 source,
7077 #[cfg(any(test, feature = "test-support"))]
7078 tier,
7079 )
7080}
7081
7082#[allow(clippy::too_many_arguments)]
7083fn scan_exact_partition<const OFFSET_ROWS: bool>(
7084 vectors: &[f32],
7085 row_count: usize,
7086 range: std::ops::Range<usize>,
7087 alive: &crate::meta::AliveSet,
7088 query: &[f32],
7089 k: usize,
7090 cancellation: &QueryCancellation<'_>,
7091 memory: &mut ExactScanMemory,
7092 #[cfg(any(test, feature = "test-support"))] controller: Option<
7093 &crate::scan::vector_fault::VectorFaultController,
7094 >,
7095 #[cfg(any(test, feature = "test-support"))] source: crate::scan::vector_fault::VectorRowSource,
7096 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7097) -> Result<crate::scan::ScanOutcome, ExactScanError> {
7098 if query.is_empty() {
7099 return Err(QueryError::Scan(crate::scan::ScanError::ZeroDimension).into());
7100 }
7101 let expected = row_count
7102 .checked_mul(query.len())
7103 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7104 if vectors.len() != expected {
7105 return Err(QueryError::Scan(crate::scan::ScanError::RowDataLength {
7106 dimension: query.len(),
7107 actual: vectors.len(),
7108 })
7109 .into());
7110 }
7111 if range.start > range.end || range.end > row_count || (!OFFSET_ROWS && range.start != 0) {
7114 return Err(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow).into());
7115 }
7116 let mut live_rows = 0;
7117 for row in range.clone() {
7118 if row.is_multiple_of(64) {
7119 cancellation.check_graph().map_err(map_scan_error)?;
7120 }
7121 let local_row = u32::try_from(row)
7122 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7123 live_rows += usize::from(alive.is_alive(local_row));
7124 }
7125 cancellation.check_graph().map_err(map_scan_error)?;
7126 let mut collector = crate::scan::topk::ExactTopK::try_new(k, live_rows, &mut memory.candidates)
7127 .map_err(QueryError::Store)?;
7128 let mut worst_score: Option<f32> = None;
7129 let mut narrowing_error: Option<(usize, f64)> = None;
7130 let start = range
7131 .start
7132 .checked_mul(query.len())
7133 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7134 let end = range
7135 .end
7136 .checked_mul(query.len())
7137 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7138 let partition_vectors = vectors
7139 .get(start..end)
7140 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7141 let rows = partition_vectors
7142 .chunks_exact(query.len())
7143 .enumerate()
7144 .map(|(local, vector)| {
7145 (
7146 if OFFSET_ROWS {
7147 range.start + local
7148 } else {
7149 local
7150 },
7151 vector,
7152 )
7153 })
7154 .filter(|(row, _)| u32::try_from(*row).is_ok_and(|row| alive.is_alive(row)));
7155 #[cfg(any(test, feature = "test-support"))]
7156 let mut trace =
7157 controller.and_then(|controller| controller.trace_rows(source, tier, range.clone()));
7158 let scored_rows = crate::quant::exact_squared_l2_with_sink(
7159 query,
7160 rows,
7161 |_row, is_checkpoint| {
7162 #[cfg(any(test, feature = "test-support"))]
7163 if let Some(controller) = controller {
7164 if let Some(trace) = &mut trace {
7165 trace.checked(_row);
7166 }
7167 if is_checkpoint {
7168 cancellation.check_graph().map_err(map_scan_error)?;
7169 }
7170 if controller.after_eligible_row(source, tier, _row) {
7171 return Err(QueryError::Cancelled { partial: false });
7172 }
7173 return Ok(());
7174 }
7175 if is_checkpoint {
7176 cancellation.check_graph().map_err(map_scan_error)?;
7177 }
7178 Ok(())
7179 },
7180 |row_id, exact| {
7181 let score = exact as f32;
7182 if !score.is_finite() {
7183 if narrowing_error.is_none_or(|(row, previous)| {
7187 exact.total_cmp(&previous).is_gt()
7188 || (exact.total_cmp(&previous).is_eq() && row_id < row)
7189 }) {
7190 narrowing_error = Some((row_id, exact));
7191 }
7192 return Ok(());
7193 }
7194 if worst_score.is_none_or(|worst| score.total_cmp(&worst).is_lt()) {
7195 worst_score = Some(score);
7196 }
7197 collector
7198 .try_push(
7199 crate::scan::ScanCandidate { row_id, score },
7200 &mut memory.candidates,
7201 )
7202 .map_err(QueryError::Store)
7203 },
7204 )
7205 .map_err(|error| match error {
7206 crate::quant::RescoreCheckError::Rescore(error) => map_l2_rescore_error(error),
7207 crate::quant::RescoreCheckError::Check(error) => error,
7208 })?;
7209 cancellation.check_graph().map_err(map_scan_error)?;
7210 if let Some((row_id, score)) = narrowing_error {
7211 return Err(ExactScanError::Narrowing { row_id, score });
7212 }
7213 let (candidates, _collector_capacity) = collector
7214 .try_into_sorted_with_ties(&mut memory.candidates)
7215 .map_err(QueryError::Store)?;
7216 #[cfg(any(test, feature = "test-support"))]
7217 if let Some(controller) = controller {
7218 controller.record_exact_scan(crate::scan::vector_fault::ExactScanWork {
7219 sorted_items: candidates.len(),
7220 collector_capacity: _collector_capacity,
7221 scored_rows,
7222 worst_score,
7223 ..Default::default()
7224 });
7225 }
7226 let dims = u64::try_from(query.len())
7227 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7228 let scored_rows = u64::try_from(scored_rows)
7229 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7230 let dims_touched = scored_rows
7231 .checked_mul(dims)
7232 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7233 let bytes_read = dims_touched
7234 .checked_mul(std::mem::size_of::<f32>() as u64)
7235 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7236 memory
7237 .worker_ids
7238 .set(std::mem::size_of::<std::thread::ThreadId>())
7239 .map_err(QueryError::Store)?;
7240 let mut worker_thread_ids = Vec::new();
7241 worker_thread_ids.try_reserve_exact(1).map_err(|_| {
7242 QueryError::Store(StoreError::AllocationFailed {
7243 needed: memory.worker_ids.bytes(),
7244 component: "exact worker IDs",
7245 })
7246 })?;
7247 worker_thread_ids.push(std::thread::current().id());
7248 Ok(crate::scan::ScanOutcome {
7249 candidates,
7250 worst_score,
7251 stats: crate::scan::ScanStats {
7252 dims_touched,
7253 bytes_read,
7254 threads_used: 1,
7255 worker_thread_ids,
7256 },
7257 })
7258}
7259
7260#[allow(clippy::too_many_arguments)]
7263fn scan_serial_squared_l2(
7264 vectors: &[f32],
7265 row_count: usize,
7266 alive: &crate::meta::AliveSet,
7267 query: &[f32],
7268 k: usize,
7269 cancellation: &QueryCancellation<'_>,
7270 memory: &mut stats::AccountedCounter,
7271 worker_memory: &mut stats::AccountedCounter,
7272 #[cfg(any(test, feature = "test-support"))] controller: Option<
7273 &crate::scan::vector_fault::VectorFaultController,
7274 >,
7275 #[cfg(any(test, feature = "test-support"))] source: crate::scan::vector_fault::VectorRowSource,
7276 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7277) -> Result<crate::scan::ScanOutcome, QueryError> {
7278 if query.is_empty() {
7279 return Err(QueryError::Scan(crate::scan::ScanError::ZeroDimension));
7280 }
7281 let expected = row_count
7282 .checked_mul(query.len())
7283 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7284 if vectors.len() != expected {
7285 return Err(QueryError::Scan(crate::scan::ScanError::RowDataLength {
7286 dimension: query.len(),
7287 actual: vectors.len(),
7288 }));
7289 }
7290 let mut live_rows = 0;
7293 for row in 0..row_count {
7294 if row.is_multiple_of(64) {
7295 cancellation.check_graph().map_err(map_scan_error)?;
7296 }
7297 let local_row = u32::try_from(row)
7298 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7299 live_rows += usize::from(alive.is_alive(local_row));
7300 }
7301 cancellation.check_graph().map_err(map_scan_error)?;
7302 let mut collector =
7303 crate::scan::topk::ExactTopK::try_new(k, live_rows, memory).map_err(QueryError::Store)?;
7304 let mut worst_score: Option<f32> = None;
7305 let mut narrowing_error: Option<(usize, f64)> = None;
7306 let rows = vectors
7307 .chunks_exact(query.len())
7308 .enumerate()
7309 .filter(|(row, _)| u32::try_from(*row).is_ok_and(|row| alive.is_alive(row)));
7310 let scored_rows = crate::quant::exact_squared_l2_with_sink(
7311 query,
7312 rows,
7313 |_row, is_checkpoint| {
7314 if is_checkpoint {
7315 cancellation.check_graph().map_err(map_scan_error)?;
7316 }
7317 #[cfg(any(test, feature = "test-support"))]
7318 if controller
7319 .is_some_and(|controller| controller.after_eligible_row(source, tier, _row))
7320 {
7321 return Err(QueryError::Cancelled { partial: false });
7322 }
7323 Ok(())
7324 },
7325 |row_id, exact| {
7326 let score = exact as f32;
7327 if !score.is_finite() {
7328 if narrowing_error.is_none_or(|(row, previous)| {
7332 exact.total_cmp(&previous).is_gt()
7333 || (exact.total_cmp(&previous).is_eq() && row_id < row)
7334 }) {
7335 narrowing_error = Some((row_id, exact));
7336 }
7337 return Ok(());
7338 }
7339 if worst_score.is_none_or(|worst| score.total_cmp(&worst).is_lt()) {
7340 worst_score = Some(score);
7341 }
7342 collector
7343 .try_push(crate::scan::ScanCandidate { row_id, score }, memory)
7344 .map_err(QueryError::Store)
7345 },
7346 )
7347 .map_err(|error| match error {
7348 crate::quant::RescoreCheckError::Rescore(error) => map_l2_rescore_error(error),
7349 crate::quant::RescoreCheckError::Check(error) => error,
7350 })?;
7351 cancellation.check_graph().map_err(map_scan_error)?;
7352 if let Some((row_id, _)) = narrowing_error {
7353 return Err(QueryError::Scan(crate::scan::ScanError::NonFiniteScore {
7354 row_id,
7355 }));
7356 }
7357 let (candidates, _collector_capacity) = collector
7358 .try_into_sorted_with_ties(memory)
7359 .map_err(QueryError::Store)?;
7360 #[cfg(any(test, feature = "test-support"))]
7361 if let Some(controller) = controller {
7362 controller.record_exact_scan(crate::scan::vector_fault::ExactScanWork {
7363 sorted_items: candidates.len(),
7364 collector_capacity: _collector_capacity,
7365 scored_rows,
7366 worst_score,
7367 ..Default::default()
7368 });
7369 }
7370 let dims = u64::try_from(query.len())
7371 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7372 let scored_rows = u64::try_from(scored_rows)
7373 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7374 let dims_touched = scored_rows
7375 .checked_mul(dims)
7376 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7377 let bytes_read = dims_touched
7378 .checked_mul(std::mem::size_of::<f32>() as u64)
7379 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7380 worker_memory
7381 .set(std::mem::size_of::<std::thread::ThreadId>())
7382 .map_err(QueryError::Store)?;
7383 let mut worker_thread_ids = Vec::new();
7384 worker_thread_ids.try_reserve_exact(1).map_err(|_| {
7385 QueryError::Store(StoreError::AllocationFailed {
7386 needed: worker_memory.bytes(),
7387 component: "exact worker IDs",
7388 })
7389 })?;
7390 worker_thread_ids.push(std::thread::current().id());
7391 Ok(crate::scan::ScanOutcome {
7392 candidates,
7393 worst_score,
7394 stats: crate::scan::ScanStats {
7395 dims_touched,
7396 bytes_read,
7397 threads_used: 1,
7398 worker_thread_ids,
7399 },
7400 })
7401}
7402
7403fn map_l2_rescore_error(error: crate::quant::RescoreError) -> QueryError {
7404 let error = match error {
7405 crate::quant::RescoreError::ZeroDimension => crate::scan::ScanError::ZeroDimension,
7406 crate::quant::RescoreError::QueryDimension { expected, actual }
7407 | crate::quant::RescoreError::CoarseScoreCount { expected, actual }
7408 | crate::quant::RescoreError::CandidateRowCount { expected, actual } => {
7409 crate::scan::ScanError::RowDataLength {
7410 dimension: expected,
7411 actual,
7412 }
7413 }
7414 crate::quant::RescoreError::RowDataLength { dimension, actual } => {
7415 crate::scan::ScanError::RowDataLength { dimension, actual }
7416 }
7417 crate::quant::RescoreError::CandidateRowOutOfRange {
7418 row_index,
7419 row_count,
7420 ..
7421 } => crate::scan::ScanError::RowDataLength {
7422 dimension: row_count,
7423 actual: row_index,
7424 },
7425 crate::quant::RescoreError::NonFiniteCoarseScore { index }
7426 | crate::quant::RescoreError::NonFiniteExactScore { row_index: index } => {
7427 crate::scan::ScanError::NonFiniteScore { row_id: index }
7428 }
7429 crate::quant::RescoreError::ZeroK
7430 | crate::quant::RescoreError::ZeroOversample
7431 | crate::quant::RescoreError::InsufficientCandidates { .. }
7432 | crate::quant::RescoreError::ArithmeticOverflow => {
7433 crate::scan::ScanError::ArithmeticOverflow
7434 }
7435 error @ (crate::quant::RescoreError::ZeroCandidateLimit
7436 | crate::quant::RescoreError::CandidateLimitExceeded { .. }) => {
7437 crate::scan::ScanError::Rescore(error)
7438 }
7439 };
7440 QueryError::Scan(error)
7441}
7442
7443fn map_scan_error(error: crate::scan::ScanError) -> QueryError {
7444 match error {
7445 crate::scan::ScanError::Cancelled { partial } => QueryError::Cancelled { partial },
7446 crate::scan::ScanError::Timeout { partial } => QueryError::Timeout { partial },
7447 crate::scan::ScanError::ReadCancelled { partial } => QueryError::ReadCancelled { partial },
7448 error => QueryError::Scan(error),
7449 }
7450}
7451
7452#[allow(clippy::too_many_arguments)]
7453fn merge_store_outcome(
7454 outcome: crate::scan::ScanOutcome,
7455 source: crate::ingest::RowSource,
7456 document: impl Fn(usize) -> Result<Option<crate::ingest::DocumentVersion>, QueryError>,
7457 candidates: &mut Vec<crate::ingest::SearchCandidate>,
7458 dims_touched: &mut u64,
7459 bytes_read: &mut u64,
7460 worker_thread_ids: &mut Vec<std::thread::ThreadId>,
7461 exact_score: bool,
7462 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
7463 &crate::scan::vector_fault::VectorFaultController,
7464 >,
7465 #[cfg(any(test, feature = "test-support"))]
7466 vector_fault_tier: crate::scan::vector_fault::VectorSearchTier,
7467) -> Result<(), QueryError> {
7468 *dims_touched = dims_touched
7469 .checked_add(outcome.stats.dims_touched)
7470 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7471 *bytes_read = bytes_read
7472 .checked_add(outcome.stats.bytes_read)
7473 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7474 for worker in outcome.stats.worker_thread_ids {
7475 if !worker_thread_ids.contains(&worker) {
7476 worker_thread_ids.push(worker);
7477 }
7478 }
7479 reserve_global_candidates(
7480 candidates,
7481 outcome.candidates.len(),
7482 #[cfg(any(test, feature = "test-support"))]
7483 vector_fault_controller,
7484 #[cfg(any(test, feature = "test-support"))]
7485 vector_fault_tier,
7486 )?;
7487 for candidate in outcome.candidates {
7488 let local_row = u32::try_from(candidate.row_id)
7489 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7490 candidates.push(crate::ingest::SearchCandidate::new(
7491 crate::ingest::GlobalRowId::new(source, local_row),
7492 document(candidate.row_id)?,
7493 candidate.score,
7494 exact_score,
7495 ));
7496 }
7497 Ok(())
7498}
7499
7500pub(crate) fn reserve_global_candidates(
7501 candidates: &mut Vec<crate::ingest::SearchCandidate>,
7502 additional: usize,
7503 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
7504 &crate::scan::vector_fault::VectorFaultController,
7505 >,
7506 #[cfg(any(test, feature = "test-support"))]
7507 vector_fault_tier: crate::scan::vector_fault::VectorSearchTier,
7508) -> Result<(), QueryError> {
7509 if additional == 0 {
7510 return Ok(());
7511 }
7512 let items = u64::try_from(additional)
7513 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7514 let candidate_bytes = u64::try_from(std::mem::size_of::<crate::ingest::SearchCandidate>())
7515 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7516 let needed = items
7517 .checked_mul(candidate_bytes)
7518 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7519 #[cfg(any(test, feature = "test-support"))]
7520 if vector_fault_controller.is_some_and(|controller| {
7521 controller.deny_global_candidate_allocation(vector_fault_tier, items, needed)
7522 }) {
7523 return Err(QueryError::Store(StoreError::AllocationFailed {
7524 needed,
7525 component: "vector search global candidates",
7526 }));
7527 }
7528 candidates.try_reserve_exact(additional).map_err(|_| {
7529 QueryError::Store(StoreError::AllocationFailed {
7530 needed,
7531 component: "vector search global candidates",
7532 })
7533 })
7534}
7535
7536struct AdmittedLexicalQuery<'a> {
7537 generation: u64,
7538 active: Arc<crate::ingest::ActiveSegment>,
7539 snapshot: snapshot::ReadSnapshot,
7540 active_query: ActiveQuery<'a>,
7541}
7542
7543struct AdmittedVectorSearch<'a> {
7544 execution_options: SearchOptions,
7545 pool: Option<Arc<pool::QueryPool>>,
7546 snapshot: snapshot::ReadSnapshot,
7547 active_segment: Arc<crate::ingest::ActiveSegment>,
7548 generation: u64,
7549 _active_query: ActiveQuery<'a>,
7550}
7551
7552struct ActiveQuery<'a> {
7553 count: &'a AtomicU64,
7554}
7555
7556impl Drop for ActiveQuery<'_> {
7557 fn drop(&mut self) {
7558 self.count.fetch_sub(1, Ordering::Relaxed);
7559 }
7560}
7561
7562impl Drop for Store {
7563 fn drop(&mut self) {
7564 self.close_best_effort();
7565 }
7566}
7567
7568#[cfg(test)]
7569mod lexical_assembly_tests;
7570
7571#[cfg(test)]
7572#[allow(clippy::expect_used, clippy::panic)]
7573mod tests {
7574 use std::error::Error;
7575 use std::sync::Arc;
7576
7577 use tempfile::tempdir;
7578
7579 use crate::lifecycle::durability::DurabilityPolicy;
7580 use crate::manifest::io::commit_manifest;
7581 use crate::manifest::{Manifest, ManifestError};
7582 use crate::meta::Schema;
7583 use crate::vfs::StdVfs;
7584 use crate::vfs::crash::MemoryVfs;
7585
7586 use super::durability::{CommitTier, DurabilityMode, DurabilityPolicyError};
7587 use super::{
7588 CancelToken, GraphSearchOptions, ManualMonotonicClock, OpenOptions, QueryControl,
7589 SearchOptions, SearchTier, Store, StoreError, StoreTestDependencies,
7590 resolve_hybrid_leg_results,
7591 };
7592
7593 #[test]
7594 fn astra_12_vocabulary_charge_survives_eviction_and_close() {
7595 use crate::fts::{index::DEFAULT_FIELD, query::LexicalQuery};
7596 use crate::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
7597 let directory = tempdir().expect("directory");
7598 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
7599 store
7600 .ingest(IngestBatch::new(vec![
7601 IngestDocument::new(
7602 DocumentVersion::new(DocId::new(1), Revision::new(1)),
7603 vec![1.0, 0.0],
7604 )
7605 .with_text("alpha alphabet beta"),
7606 ]))
7607 .expect("ingest");
7608 store
7609 .search_lexical_structured(
7610 &LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD),
7611 1,
7612 64,
7613 QueryControl::Cancel(CancelToken::new()),
7614 )
7615 .expect("build cache");
7616 let dictionary = {
7617 let entry = store.lexical_index_cache.entry.lock().expect("assembly");
7618 let vocabulary = entry
7619 .as_ref()
7620 .expect("assembly")
7621 .assembly
7622 .vocabulary
7623 .lock()
7624 .expect("vocabulary");
7625 Arc::clone(vocabulary.as_ref().expect("cached"))
7626 };
7627 let bytes = dictionary._memory.bytes();
7628 assert!(bytes > 17);
7629 let second_owner = Arc::clone(&dictionary);
7630 drop(
7631 store
7632 .lexical_index_cache
7633 .entry
7634 .lock()
7635 .expect("evict")
7636 .take(),
7637 );
7638 assert_eq!(store.stats().expect("held bytes").cache_bytes, bytes);
7639 store
7640 .search_lexical_structured(
7641 &LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD),
7642 1,
7643 64,
7644 QueryControl::Cancel(CancelToken::new()),
7645 )
7646 .expect("replacement cache");
7647 assert!(store.stats().expect("both dictionaries").cache_bytes > bytes);
7648 store.close().expect("close");
7649 assert_eq!(
7650 store
7651 .accounting
7652 .audit()
7653 .expect("closed accounting")
7654 .cache_bytes,
7655 bytes
7656 );
7657 drop(dictionary);
7658 assert_eq!(
7659 store.accounting.audit().expect("second owner").cache_bytes,
7660 bytes
7661 );
7662 assert_eq!(
7663 second_owner.view.iter().collect::<Vec<_>>(),
7664 vec![b"alpha".as_slice(), b"alphabet", b"beta"]
7665 );
7666 drop(second_owner);
7667 assert_eq!(store.accounting.audit().expect("last owner").cache_bytes, 0);
7668 assert_eq!(
7669 store.accounting.audit().expect("scratch").temporary_bytes,
7670 0
7671 );
7672 }
7673
7674 #[test]
7675 fn astra_12_refused_vocabulary_is_not_published() {
7676 use crate::fts::{index::DEFAULT_FIELD, query::LexicalQuery};
7677 use crate::ingest::{
7678 DocId, DocumentVersion, IngestBatch, IngestDocument, Revision, StoreLexicalError,
7679 };
7680 let directory = tempdir().expect("directory");
7681 let writer = Store::open(directory.path(), OpenOptions::default()).expect("open");
7682 writer
7683 .ingest(IngestBatch::new(vec![
7684 IngestDocument::new(
7685 DocumentVersion::new(DocId::new(1), Revision::new(1)),
7686 vec![1.0, 0.0],
7687 )
7688 .with_text("alpha alphabet beta"),
7689 ]))
7690 .expect("ingest");
7691 writer.seal().expect("seal");
7692 writer.close().expect("close writer");
7693 let store = Store::open(
7694 directory.path(),
7695 OpenOptions::default().with_max_temp_bytes(1),
7696 )
7697 .expect("small budget");
7698 for _ in 0..2 {
7699 let result = store.search_lexical_structured(
7700 &LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD),
7701 1,
7702 64,
7703 QueryControl::Cancel(CancelToken::new()),
7704 );
7705 assert!(
7706 matches!(
7707 result,
7708 Err(StoreLexicalError::Query(super::QueryError::Store(
7709 StoreError::BudgetExceeded {
7710 component: "temporary",
7711 ..
7712 }
7713 )))
7714 ),
7715 "{result:?}"
7716 );
7717 assert_eq!(store.stats().expect("released").temporary_bytes, 0);
7718 let entry = store.lexical_index_cache.entry.lock().expect("assembly");
7719 assert!(
7720 entry
7721 .as_ref()
7722 .expect("assembled")
7723 .assembly
7724 .vocabulary
7725 .lock()
7726 .expect("dictionary")
7727 .is_none()
7728 );
7729 }
7730 store.close().expect("close");
7731 assert_eq!(
7732 store
7733 .accounting
7734 .audit()
7735 .expect("released cache")
7736 .cache_bytes,
7737 0
7738 );
7739 }
7740
7741 #[test]
7742 fn astra_03_lexical_assembly_charge_survives_cache_eviction() {
7743 use crate::fts::index::DEFAULT_FIELD;
7744 use crate::fts::search::TermQuery;
7745 use crate::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
7746 let directory = tempdir().expect("directory");
7747 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
7748 store
7749 .ingest(IngestBatch::new(vec![
7750 IngestDocument::new(
7751 DocumentVersion::new(DocId::new(1), Revision::new(1)),
7752 vec![1.0, 0.0],
7753 )
7754 .with_text("alpha beta"),
7755 ]))
7756 .expect("ingest");
7757 store
7758 .search_lexical(
7759 &TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]),
7760 1,
7761 QueryControl::Cancel(CancelToken::new()),
7762 )
7763 .expect("build lexical cache");
7764 let (first, second) = {
7765 let entry = store.lexical_index_cache.entry.lock().expect("cache lock");
7766 let assembly = &entry.as_ref().expect("cached assembly").assembly;
7767 (Arc::clone(assembly), Arc::clone(assembly))
7768 };
7769 let charged = store.stats().expect("charged stats").cache_bytes;
7770 let evicted = store
7771 .lexical_index_cache
7772 .entry
7773 .lock()
7774 .expect("cache lock")
7775 .take();
7776 drop(evicted);
7777 assert_eq!(
7778 first.index.document_count(),
7779 1,
7780 "eviction left the query data alive"
7781 );
7782 assert_eq!(
7783 store.stats().expect("held stats").cache_bytes,
7784 charged,
7785 "cache eviction must not release a query-owned assembly's charge"
7786 );
7787 drop(first);
7788 assert_eq!(
7789 store.stats().expect("second owner stats").cache_bytes,
7790 charged
7791 );
7792 drop(second);
7793 let released = store.stats().expect("released stats").cache_bytes;
7794 assert_eq!(
7795 released, 0,
7796 "the final owner must release the full assembly charge"
7797 );
7798 println!("assembly cache bytes: retained={charged}, after final owner={released}");
7799 store.close().expect("close");
7800 }
7801
7802 const CEILING_ROWS: [[f32; 4]; 3] = [
7804 [0.5, -0.25, 0.75, 0.125],
7805 [3.0, 1.5, -2.5, 0.25],
7806 [0.0, 0.0, 0.0, 0.0],
7807 ];
7808
7809 fn publish_f32_segment(directory: &std::path::Path) {
7811 let created = Store::open(directory, OpenOptions::default()).expect("create f32 store");
7812 created.close().expect("close f32 store");
7813
7814 let schema = Schema::timestamp_only();
7815 let mut columns = crate::meta::ColumnStoreBuilder::new(schema.clone());
7816 for row in 0..CEILING_ROWS.len() {
7817 columns
7818 .push_row(i64::try_from(row).expect("row timestamp"), &[])
7819 .expect("timestamp row");
7820 }
7821 let columns = columns.finish().expect("columns");
7822 let alive =
7823 crate::meta::AliveSet::new(u32::try_from(CEILING_ROWS.len()).expect("row count"));
7824 let rescore = CEILING_ROWS.iter().flatten().copied().collect::<Vec<f32>>();
7825 let codes = rescore
7826 .iter()
7827 .flat_map(|value| value.to_bits().to_le_bytes())
7828 .collect::<Vec<u8>>();
7829 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
7830 .expect("f32 fixture durability");
7831 let meta = crate::segment::writer::write_segment(
7832 &StdVfs,
7833 directory,
7834 crate::segment::writer::SegmentBuild {
7835 id: crate::segment::SegmentId::new(11, [7; 10]),
7836 scheme: 0,
7837 dims: 4,
7838 codes: &codes,
7839 factors: crate::segment::writer::SegmentFactors::F32,
7840 rescore: &rescore,
7841 columns: &columns,
7842 alive: &alive,
7843 },
7844 policy,
7845 )
7846 .expect("write scheme-0 segment");
7847 commit_manifest(
7848 &StdVfs,
7849 directory,
7850 &Manifest {
7851 generation: 1,
7852 log_seq: 0,
7853 segments: vec![meta],
7854 epochs: Vec::new(),
7855 epoch_alias: None,
7856 schema,
7857 },
7858 policy,
7859 )
7860 .expect("publish scheme-0 segment");
7861 }
7862
7863 fn ingest_ceiling_rows(store: &Store, keep_active: usize, one_segment_per_row: bool) {
7866 use crate::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
7867
7868 let sealed = CEILING_ROWS.len() - keep_active;
7869 for (index, row) in CEILING_ROWS.iter().enumerate() {
7870 let version = DocumentVersion::new(
7871 DocId::new(u128::try_from(index).expect("document id")),
7872 Revision::new(1),
7873 );
7874 store
7875 .ingest(IngestBatch::new(vec![IngestDocument::new(
7876 version,
7877 row.to_vec(),
7878 )]))
7879 .expect("ingest ceiling row");
7880 let seal_here = if one_segment_per_row {
7881 index < sealed
7882 } else {
7883 index + 1 == sealed
7884 };
7885 if seal_here {
7886 store.seal().expect("seal ceiling rows");
7887 }
7888 }
7889 }
7890
7891 fn assert_remembered_ceiling_is_exact(store: &Store, expected_segments: usize) {
7894 let queries: [Vec<f32>; 3] = [
7895 vec![0.0, 0.0, 0.0, 0.0],
7896 vec![1.0, -1.0, 0.5, 2.0],
7897 vec![7.5, 0.25, -3.5, 0.125],
7898 ];
7899 let admitted = store
7900 .admit_vector_search(SearchOptions::default())
7901 .expect("admit ceiling query");
7902 let snapshot = &admitted.snapshot;
7903 let active = &admitted.active_segment;
7904 assert_eq!(snapshot.segments().len(), expected_segments);
7905 for segment in snapshot.segments() {
7906 assert!(
7907 segment.cached_vector_ceiling_norm_range().is_none(),
7908 "a segment remembered an enclosure before any query walked it"
7909 );
7910 }
7911
7912 for query in &queries {
7913 let cold = super::exact_vector_ceiling(snapshot, active, query).expect("cold ceiling");
7914 for segment in snapshot.segments() {
7915 let remembered = segment
7916 .cached_vector_ceiling_norm_range()
7917 .expect("a walked segment must remember its enclosure");
7918 let recomputed = super::compute_segment_vector_ceiling_norm_range(segment)
7919 .expect("recomputed enclosure");
7920 assert_eq!(
7921 remembered.endpoint_bits(),
7922 recomputed.endpoint_bits(),
7923 "remembered enclosure differs from the recomputed one"
7924 );
7925 }
7926 let warm = super::exact_vector_ceiling(snapshot, active, query).expect("warm ceiling");
7927 let recomputed = super::fold_exact_vector_ceiling(
7928 snapshot,
7929 active,
7930 query,
7931 super::compute_segment_vector_ceiling_norm_range,
7932 )
7933 .expect("recomputed ceiling");
7934 assert_eq!(
7935 cold.to_bits(),
7936 warm.to_bits(),
7937 "the remembered ceiling changed between queries"
7938 );
7939 assert_eq!(
7940 cold.to_bits(),
7941 recomputed.to_bits(),
7942 "the remembered ceiling is not the recomputed ceiling"
7943 );
7944 }
7945 }
7946
7947 #[test]
7948 fn astra_00_elapsed_includes_ceiling_work() {
7949 use std::time::{Duration, Instant};
7950 struct CeilingClock {
7951 base: Instant,
7952 snapshot: Arc<super::PublishedSnapshot>,
7953 }
7954 impl super::MonotonicClock for CeilingClock {
7955 fn now(&self) -> Instant {
7956 if self
7959 .snapshot
7960 .segments()
7961 .iter()
7962 .all(|segment| segment.cached_vector_ceiling_norm_range().is_some())
7963 {
7964 self.base + Duration::from_secs(7)
7965 } else {
7966 self.base
7967 }
7968 }
7969 }
7970 let directory = tempdir().expect("ceiling timer store");
7971 let mut store = Store::open(directory.path(), OpenOptions::default()).expect("open");
7972 ingest_ceiling_rows(&store, 0, false);
7973 let admitted = store
7974 .admit_vector_search(SearchOptions::default())
7975 .expect("snapshot");
7976 assert_eq!(admitted.snapshot.segments().len(), 1);
7977 assert!(
7978 admitted
7979 .snapshot
7980 .segments()
7981 .iter()
7982 .all(|segment| { segment.cached_vector_ceiling_norm_range().is_none() })
7983 );
7984 let snapshot = Arc::clone(&admitted.snapshot);
7985 drop(admitted);
7986 store.clock = Arc::new(CeilingClock {
7987 base: Instant::now(),
7988 snapshot,
7989 });
7990 let outcome = store
7991 .search(
7992 crate::ingest::SearchRequest::new(&[1.0, 0.0, 0.0, 0.0]),
7993 1,
7994 SearchOptions::default().with_tier(SearchTier::Exact),
7995 QueryControl::Cancel(CancelToken::new()),
7996 )
7997 .expect("exact query");
7998 assert_eq!(outcome.diagnostics.elapsed, Duration::from_secs(7));
7999 store.clock = Arc::new(super::SystemMonotonicClock);
8001 store.close().expect("close");
8002 }
8003
8004 #[test]
8005 fn the_remembered_vector_ceiling_is_the_recomputed_ceiling_for_every_scheme() {
8006 let bit4_directory = tempdir().expect("Bit4 ceiling store directory");
8007 let bit4 = Store::open(bit4_directory.path(), OpenOptions::default())
8008 .expect("open Bit4 ceiling store");
8009 ingest_ceiling_rows(&bit4, 0, false);
8010 assert_remembered_ceiling_is_exact(&bit4, 1);
8011 bit4.close().expect("close Bit4 ceiling store");
8012
8013 let split_directory = tempdir().expect("multi-segment ceiling store directory");
8014 let split = Store::open(split_directory.path(), OpenOptions::default())
8015 .expect("open multi-segment ceiling store");
8016 ingest_ceiling_rows(&split, 0, true);
8017 assert_remembered_ceiling_is_exact(&split, CEILING_ROWS.len());
8018 split.close().expect("close multi-segment ceiling store");
8019
8020 let int8_directory = tempdir().expect("Int8 ceiling store directory");
8021 let int8 = Store::open_with_test_dependencies(
8022 int8_directory.path(),
8023 OpenOptions::default(),
8024 StoreTestDependencies::new(Arc::new(StdVfs), Arc::new(super::SystemMonotonicClock))
8025 .with_vector_seal_scheme(crate::quant::QuantScheme::Int8),
8026 )
8027 .expect("open Int8 ceiling store");
8028 ingest_ceiling_rows(&int8, 0, false);
8029 assert_remembered_ceiling_is_exact(&int8, 1);
8030 int8.close().expect("close Int8 ceiling store");
8031
8032 let f32_directory = tempdir().expect("F32 ceiling store directory");
8033 publish_f32_segment(f32_directory.path());
8034 let f32_store = Store::open(f32_directory.path(), OpenOptions::default())
8035 .expect("open F32 ceiling store");
8036 assert_remembered_ceiling_is_exact(&f32_store, 1);
8037 f32_store.close().expect("close F32 ceiling store");
8038
8039 let active_directory = tempdir().expect("active ceiling store directory");
8040 let active = Store::open(active_directory.path(), OpenOptions::default())
8041 .expect("open active ceiling store");
8042 ingest_ceiling_rows(&active, 2, false);
8043 assert!(
8044 !active
8045 .admit_vector_search(SearchOptions::default())
8046 .expect("admit active ceiling query")
8047 .active_segment
8048 .is_empty(),
8049 "the active-segment case sealed every row"
8050 );
8051 assert_remembered_ceiling_is_exact(&active, 1);
8052 active.close().expect("close active ceiling store");
8053 }
8054
8055 #[test]
8056 fn exact_and_graph_tiers_preserve_quantizer_validation_errors() {
8057 let directory = tempdir().expect("store directory");
8058 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
8059 let query = vec![0.0_f32; crate::kernels::MAX_DOT_I8_DIMENSION + 1];
8060
8061 for tier in [
8062 SearchTier::Exact,
8063 SearchTier::Graph(GraphSearchOptions::default()),
8064 ] {
8065 let error = store
8066 .search(
8067 crate::ingest::SearchRequest::new(&query),
8068 1,
8069 SearchOptions::default().with_tier(tier),
8070 QueryControl::Cancel(CancelToken::new()),
8071 )
8072 .expect_err("oversized quantized query must retain its error");
8073 assert!(matches!(
8074 error,
8075 super::QueryError::Scan(crate::scan::ScanError::Quant(
8076 crate::quant::QuantError::DimensionTooLarge { actual, maximum }
8077 )) if actual == query.len()
8078 && maximum == crate::kernels::MAX_DOT_I8_DIMENSION
8079 ));
8080 }
8081 }
8082
8083 #[cfg(feature = "allocation-audit")]
8084 #[test]
8085 fn exact_tier_query_does_not_prepare_quantized_forms() {
8086 let directory = tempdir().expect("store directory");
8087 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
8088 let query = [0.25_f32; 128];
8089 let options = SearchOptions::default().with_tier(SearchTier::Exact);
8090 let warm = store
8091 .search(
8092 crate::ingest::SearchRequest::new(&query),
8093 1,
8094 options,
8095 QueryControl::Cancel(CancelToken::new()),
8096 )
8097 .expect("warm exact-tier search");
8098 assert!(warm.candidates.is_empty());
8099 let control = QueryControl::Cancel(CancelToken::new());
8100
8101 let (outcome, report) = crate::allocation_audit::audit_engine_path(|| {
8102 store.search(
8103 crate::ingest::SearchRequest::new(&query),
8104 1,
8105 options,
8106 control,
8107 )
8108 });
8109 let outcome = outcome.expect("audited exact-tier search");
8110 assert!(outcome.candidates.is_empty());
8111 assert_eq!(report.allocations, 0, "Exact tier prepared quantized forms");
8112 assert_eq!(report.attributed_bytes, 0);
8113 assert_eq!(report.unattributed_bytes, 0);
8114 }
8115
8116 #[test]
8117 fn astra_18_scoped_work_is_joined_before_cancel_returns() {
8118 use std::sync::atomic::{AtomicBool, Ordering};
8119 use std::sync::mpsc;
8120 let worker = super::pool::LexicalWorker::start().expect("worker");
8121 let token = CancelToken::new();
8122 let work_token = token.clone();
8123 let finished = AtomicBool::new(false);
8124 let mut borrowed_rows = Vec::new();
8125 let (started_tx, started_rx) = mpsc::sync_channel(0);
8126 let (release_tx, release_rx) = mpsc::sync_channel(0);
8127 let worker_rows = &mut borrowed_rows;
8128 let worker_finished = &finished;
8129 let (caller, lexical) = worker
8130 .run_scoped(
8131 move || {
8132 started_tx.send(()).expect("entered borrowed job");
8133 release_rx.recv().expect("caller has cancelled");
8134 assert!(work_token.is_cancelled());
8135 worker_rows.push(42_u32);
8136 worker_finished.store(true, Ordering::Release);
8137 Err::<(), _>(crate::fusion::FusionError::Cancelled { partial: false })
8138 },
8139 || {
8140 started_rx.recv().expect("worker entered");
8141 token.cancel();
8142 assert!(!finished.load(Ordering::Acquire));
8143 release_tx.send(()).expect("release borrowed job");
8144 Err::<(), _>(crate::fusion::FusionError::Cancelled { partial: false })
8145 },
8146 )
8147 .expect("scoped dispatch");
8148 assert!(finished.load(Ordering::Acquire));
8149 assert_eq!(borrowed_rows, [42]);
8150 assert_eq!(
8151 resolve_hybrid_leg_results(caller, lexical.expect("worker returned")),
8152 Err(crate::fusion::FusionError::Cancelled { partial: false }),
8153 );
8154 worker.stop_and_join().expect("join worker");
8155 }
8156
8157 #[test]
8158 fn astra_18_control_error_precedence_is_completion_order_independent() {
8159 use crate::fusion::{FusionError, FusionLeg, LegFailureKind};
8160 use std::sync::mpsc;
8161 let failures = [
8162 FusionError::Leg {
8163 leg: FusionLeg::Vector,
8164 kind: LegFailureKind::Scan,
8165 detail: "scan".to_owned(),
8166 },
8167 FusionError::Cancelled { partial: false },
8168 FusionError::Timeout { partial: false },
8169 FusionError::ReadCancelled { partial: false },
8170 ];
8171 let worker = super::pool::LexicalWorker::start().expect("worker");
8172 for (a, vector) in failures.iter().enumerate() {
8173 for (b, lexical) in failures.iter().enumerate() {
8174 for lexical_first in [false, true] {
8175 let (lexical_tx, lexical_rx) = mpsc::sync_channel(1);
8176 let (vector_tx, vector_rx) = mpsc::sync_channel(1);
8177 let (vector_result, lexical_result) = worker
8178 .run_scoped(
8179 move || {
8180 if !lexical_first {
8181 vector_rx.recv().expect("vector completed");
8182 }
8183 if lexical_first {
8184 lexical_tx.send(()).expect("lexical completed");
8185 }
8186 Err::<(), _>(lexical.clone())
8187 },
8188 || {
8189 if lexical_first {
8190 lexical_rx.recv().expect("lexical completed");
8191 }
8192 if !lexical_first {
8193 vector_tx.send(()).expect("vector completed");
8194 }
8195 Err::<(), _>(vector.clone())
8196 },
8197 )
8198 .expect("dispatch");
8199 assert_eq!(
8200 resolve_hybrid_leg_results(vector_result, lexical_result.expect("joined")),
8201 Err(failures.get(a.max(b)).expect("known failure index").clone()),
8202 "vector={a} lexical={b} lexical_first={lexical_first}",
8203 );
8204 }
8205 }
8206 }
8207 worker.stop_and_join().expect("join worker");
8208 }
8209
8210 #[test]
8211 fn hybrid_dual_failure_precedence_is_independent_of_leg_completion_order() {
8212 use crate::fusion::{FusionError, FusionLeg, LegFailureKind};
8213
8214 let vector_failure = FusionError::Leg {
8215 leg: FusionLeg::Vector,
8216 kind: LegFailureKind::Scan,
8217 detail: "vector failure".to_owned(),
8218 };
8219 let lexical_failure = FusionError::Leg {
8220 leg: FusionLeg::Lexical,
8221 kind: LegFailureKind::Lexical,
8222 detail: "lexical failure".to_owned(),
8223 };
8224 assert_eq!(
8225 resolve_hybrid_leg_results::<(), ()>(
8226 Err(vector_failure.clone()),
8227 Err(lexical_failure.clone()),
8228 ),
8229 Err(vector_failure.clone()),
8230 );
8231 assert_eq!(
8232 resolve_hybrid_leg_results::<(), ()>(
8233 Err(vector_failure.clone()),
8234 Err(FusionError::Cancelled { partial: false }),
8235 ),
8236 Err(FusionError::Cancelled { partial: false }),
8237 );
8238 assert_eq!(
8239 resolve_hybrid_leg_results::<(), ()>(
8240 Err(FusionError::Timeout { partial: false }),
8241 Err(FusionError::ReadCancelled { partial: false }),
8242 ),
8243 Err(FusionError::ReadCancelled { partial: false }),
8244 );
8245 assert_eq!(
8246 resolve_hybrid_leg_results::<(), ()>(
8247 Err(FusionError::Cancelled { partial: false }),
8248 Err(FusionError::Timeout { partial: false }),
8249 ),
8250 Err(FusionError::Timeout { partial: false }),
8251 );
8252 assert_eq!(
8253 resolve_hybrid_leg_results::<(), ()>(
8254 Err(FusionError::ReadCancelled { partial: false }),
8255 Err(lexical_failure),
8256 ),
8257 Err(FusionError::ReadCancelled { partial: false }),
8258 );
8259 }
8260
8261 #[test]
8262 fn accounting_errors_report_the_typed_context() {
8263 let statistics = StoreError::Statistics {
8264 component: "mapped resident bytes",
8265 source: std::io::Error::other("mincore failed"),
8266 };
8267 assert_eq!(
8268 statistics.to_string(),
8269 "store statistics mapped resident bytes: mincore failed"
8270 );
8271 assert_eq!(
8272 statistics.source().map(ToString::to_string),
8273 Some("mincore failed".to_owned())
8274 );
8275
8276 let budget = StoreError::BudgetExceeded {
8277 needed: 65,
8278 budget: 64,
8279 component: "wal",
8280 };
8281 assert_eq!(
8282 budget.to_string(),
8283 "store wal allocation needs 65 bytes, budget is 64 bytes"
8284 );
8285 assert!(budget.source().is_none());
8286
8287 let allocation = StoreError::AllocationFailed {
8288 needed: 23,
8289 component: "snapshot",
8290 };
8291 assert_eq!(
8292 allocation.to_string(),
8293 "store snapshot allocator rejected 23 bytes"
8294 );
8295 assert!(allocation.source().is_none());
8296 }
8297
8298 #[test]
8299 fn open_options_resolve_durability_before_filesystem_mutation() {
8300 let parent = tempdir().expect("parent directory");
8301 let store_path = parent.path().join("not-created");
8302 let options =
8303 OpenOptions::new().with_durability(DurabilityMode::Attached, CommitTier::Durable);
8304
8305 let error = match Store::open(&store_path, options) {
8306 Ok(store) => {
8307 drop(store);
8308 panic!("attached durability unexpectedly opened")
8309 }
8310 Err(error) => error,
8311 };
8312
8313 assert!(matches!(
8314 error,
8315 StoreError::Durability(DurabilityPolicyError::AttachedNotYetSupported)
8316 ));
8317 assert!(!store_path.exists(), "rejected open created store files");
8318 }
8319
8320 #[test]
8321 fn memory_vfs_store_open_does_not_create_host_directory() {
8322 let parent = tempdir().expect("parent directory");
8323 let store_path = parent.path().join("memory-only-store");
8324 let dependencies = StoreTestDependencies::new(
8325 Arc::new(MemoryVfs::new()),
8326 Arc::new(ManualMonotonicClock::new()),
8327 );
8328
8329 let store =
8330 Store::open_with_test_dependencies(&store_path, OpenOptions::read_only(), dependencies)
8331 .expect("MemoryVfs-backed store open");
8332 store.close().expect("MemoryVfs-backed store close");
8333
8334 assert!(
8335 !store_path.exists(),
8336 "MemoryVfs-backed Store touched the host filesystem"
8337 );
8338 }
8339
8340 #[test]
8341 fn open_refuses_a_manifest_ahead_of_the_checked_wal() {
8342 let directory = tempdir().expect("store directory");
8343 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::Ordered)
8344 .expect("derived policy");
8345 commit_manifest(
8346 &StdVfs,
8347 directory.path(),
8348 &Manifest {
8349 generation: 1,
8350 log_seq: 1,
8351 segments: Vec::new(),
8352 epochs: Vec::new(),
8353 epoch_alias: None,
8354 schema: Schema::new(Vec::new()).expect("schema"),
8355 },
8356 policy,
8357 )
8358 .expect("manifest");
8359
8360 let error = match Store::open(directory.path(), OpenOptions::default()) {
8361 Ok(store) => {
8362 drop(store);
8363 panic!("manifest ahead of missing WAL unexpectedly opened")
8364 }
8365 Err(error) => error,
8366 };
8367
8368 assert!(matches!(
8369 error,
8370 StoreError::Manifest(ManifestError::AheadOfLog {
8371 snapshot: 1,
8372 durable: 0
8373 })
8374 ));
8375 }
8376}