1use super::*;
15use crate::catalog_cmds::CatalogCommand;
16use crate::jobs::{run_build_publish, BuildPublishJob, JobContext, JobError, JobKind, JobTarget};
17use crate::retention::{PinGuard, PinSource};
18use crate::schema::IndexDef;
19use crate::wal::DdlOp;
20
21#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
23enum IndexBuildKind {
24 Create {
25 definition: IndexDef,
26 },
27 Replace {
28 expected_old_name: String,
29 new_definition: IndexDef,
30 },
31}
32
33impl IndexBuildKind {
34 fn definition(&self) -> &IndexDef {
35 match self {
36 Self::Create { definition } => definition,
37 Self::Replace { new_definition, .. } => new_definition,
38 }
39 }
40}
41
42const INDEX_BUILD_SPEC_VERSION: u16 = 1;
43
44#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
47#[serde(deny_unknown_fields)]
48struct IndexBuildSpec {
49 version: u16,
50 table: String,
51 expected_schema_sequence: u64,
52 kind: IndexBuildKind,
53}
54
55impl IndexBuildSpec {
56 fn encode(&self) -> Result<Vec<u8>> {
57 serde_json::to_vec(self)
58 .map_err(|error| MongrelError::Other(format!("serialize index-build spec: {error}")))
59 }
60
61 fn decode(bytes: &[u8]) -> Result<Self> {
62 let spec: Self = serde_json::from_slice(bytes)
63 .map_err(|error| MongrelError::Other(format!("decode index-build spec: {error}")))?;
64 if spec.version != INDEX_BUILD_SPEC_VERSION {
65 return Err(MongrelError::Other(format!(
66 "unsupported index-build spec version {}",
67 spec.version
68 )));
69 }
70 Ok(spec)
71 }
72}
73
74struct IndexBuildJob<'a> {
76 db: &'a Database,
77 table: String,
78 kind: IndexBuildKind,
79 expected_schema_sequence: u64,
80 snapshot_epoch: Option<Epoch>,
82 pin: Option<PinGuard>,
84 snapshot_row_count: u64,
86 snapshot_rows: std::collections::BTreeMap<RowId, crate::memtable::Row>,
88 artifact: Option<crate::engine::SecondaryIndexArtifact>,
90 memory_reservation: Option<crate::memory::Reservation>,
92 built_through: Option<Epoch>,
94 built_data_generation: Option<u64>,
96 published: bool,
98}
99
100impl IndexBuildJob<'_> {
101 fn catalog_command(&self) -> CatalogCommand {
102 match &self.kind {
103 IndexBuildKind::Create { definition } => CatalogCommand::AddIndex {
104 table: self.table.clone(),
105 index: definition.clone(),
106 },
107 IndexBuildKind::Replace {
108 expected_old_name,
109 new_definition,
110 } => CatalogCommand::ReplaceIndex {
111 table: self.table.clone(),
112 expected_schema_sequence: self.expected_schema_sequence,
113 expected_old_name: expected_old_name.clone(),
114 new_definition: new_definition.clone(),
115 },
116 }
117 }
118
119 fn release_pin(&mut self) {
120 self.pin = None;
121 }
122
123 fn reserve_hidden_memory(
124 &mut self,
125 schema: &Schema,
126 rows: &[crate::memtable::Row],
127 ) -> std::result::Result<(), JobError> {
128 let row_bytes = rows.iter().fold(0u64, |total, row| {
129 total.saturating_add(row.estimated_bytes().saturating_add(64))
130 });
131 let mut estimate = row_bytes.saturating_mul(2);
132 if self.definition().kind == crate::schema::IndexKind::Ann {
133 let dim = schema
134 .columns
135 .iter()
136 .find(|column| column.id == self.definition().column_id)
137 .and_then(|column| match column.ty {
138 crate::schema::TypeId::Embedding { dim } => Some(dim as u64),
139 _ => None,
140 })
141 .ok_or_else(|| {
142 JobError::Phase(format!(
143 "ANN index {} has no embedding column",
144 self.definition().name
145 ))
146 })?;
147 let options = self.definition().options.ann.clone().unwrap_or_default();
148 let vector_bytes = match options.quantization {
149 crate::schema::AnnQuantization::BinarySign => dim.div_ceil(8),
150 crate::schema::AnnQuantization::Dense => dim.saturating_mul(4),
151 };
152 let levels = (crate::index::hnsw::MAX_HNSW_LEVEL as u64).saturating_add(1);
153 let adjacency_slots = (options.m as u64)
154 .saturating_mul(2)
155 .saturating_add((options.m as u64).saturating_mul(levels.saturating_sub(1)));
156 let per_node = vector_bytes
157 .saturating_add(std::mem::size_of::<RowId>() as u64)
158 .saturating_add(levels.saturating_mul(std::mem::size_of::<Vec<usize>>() as u64))
159 .saturating_add(
160 adjacency_slots.saturating_mul(std::mem::size_of::<usize>() as u64),
161 );
162 estimate = estimate
164 .saturating_add(per_node.saturating_mul(rows.len() as u64).saturating_mul(2));
165 estimate = estimate.saturating_add(
166 (options.ef_construction as u64)
167 .saturating_mul(dim.saturating_add(32))
168 .saturating_mul(8),
169 );
170 }
171
172 let requested = usize::try_from(estimate).unwrap_or(usize::MAX);
173 let governor = self.db.memory_governor();
174 let result = match self.memory_reservation.as_mut() {
175 Some(reservation) => reservation.resize(estimate),
176 None => governor
177 .try_reserve(estimate, crate::memory::MemoryClass::Compaction)
178 .map(|reservation| self.memory_reservation = Some(reservation)),
179 };
180 result.map_err(|error| {
181 let limit = match error {
182 crate::memory::MemoryError::Exhausted { available, .. } => {
183 usize::try_from(available).unwrap_or(usize::MAX)
184 }
185 crate::memory::MemoryError::LowPriorityRejected { .. } => 0,
186 crate::memory::MemoryError::InvalidConfig(_) => {
187 usize::try_from(governor.class_budget(crate::memory::MemoryClass::Compaction))
188 .unwrap_or(usize::MAX)
189 }
190 };
191 JobError::ResourceLimitExceeded {
192 resource: "index build memory",
193 requested,
194 limit,
195 }
196 })
197 }
198
199 fn definition(&self) -> &IndexDef {
200 self.kind.definition()
201 }
202
203 fn publication_already_applied(&self, schema: &Schema) -> bool {
204 schema.schema_id == self.expected_schema_sequence.saturating_add(1)
205 && schema
206 .indexes
207 .iter()
208 .any(|index| index == self.definition())
209 && match &self.kind {
210 IndexBuildKind::Create { .. } => true,
211 IndexBuildKind::Replace {
212 expected_old_name,
213 new_definition,
214 } => {
215 expected_old_name == &new_definition.name
216 || !schema
217 .indexes
218 .iter()
219 .any(|index| index.name == *expected_old_name)
220 }
221 }
222 }
223
224 fn refresh_hidden(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
229 context.check_cancelled()?;
230 let handle = self.db.table(&self.table).map_err(job_phase_err)?;
231 let (generation, snapshot) = handle
232 .read_generation_with_context(None)
233 .map_err(job_phase_err)?;
234 if self.pin.is_none() {
235 self.pin = Some(
236 Arc::clone(generation.pin_registry())
237 .pin(PinSource::OnlineIndexBuild, snapshot.epoch),
238 );
239 self.snapshot_epoch.get_or_insert(snapshot.epoch);
240 }
241
242 let rows = generation.visible_rows(snapshot).map_err(job_phase_err)?;
243 self.reserve_hidden_memory(generation.schema(), &rows)?;
244 let next: std::collections::BTreeMap<_, _> =
245 rows.into_iter().map(|row| (row.row_id, row)).collect();
246 self.snapshot_rows
247 .retain(|row_id, _| next.contains_key(row_id));
248 for (row_id, row) in next {
249 self.snapshot_rows.insert(row_id, row);
250 }
251
252 let rows: Vec<_> = self.snapshot_rows.values().cloned().collect();
253 let total = rows.len();
254 let definition = self.definition().clone();
255 let mut last_reported = None;
256 let artifact = generation
257 .build_secondary_index_artifact(&definition, &rows, |done, total| {
258 context.check_cancelled().map_err(MongrelError::from)?;
259 if total > 0
260 && last_reported != Some(done)
261 && (done.is_multiple_of(256) || done == total)
262 {
263 context
264 .report_progress(done as u64, total as u64)
265 .map_err(MongrelError::from)?;
266 last_reported = Some(done);
267 }
268 Ok(())
269 })
270 .map_err(job_phase_err)?;
271 self.snapshot_row_count = total as u64;
272 self.built_through = Some(snapshot.epoch);
273 self.built_data_generation = Some(generation.data_generation());
274 self.artifact = Some(artifact);
275 Ok(())
276 }
277}
278
279impl BuildPublishJob for IndexBuildJob<'_> {
280 fn checkpoint_state(&self) -> Vec<u8> {
281 let epoch = self.snapshot_epoch.map(|e| e.0).unwrap_or(0);
283 let built_through = self.built_through.map(|epoch| epoch.0).unwrap_or(0);
284 let built_data_generation = self.built_data_generation.unwrap_or(0);
285 let published = u8::from(self.published);
286 let mut out = Vec::with_capacity(33);
287 out.extend_from_slice(&epoch.to_le_bytes());
288 out.extend_from_slice(&self.snapshot_row_count.to_le_bytes());
289 out.extend_from_slice(&built_through.to_le_bytes());
290 out.extend_from_slice(&built_data_generation.to_le_bytes());
291 out.push(published);
292 out
293 }
294
295 fn restore_checkpoint(&mut self, state: &[u8]) -> std::result::Result<(), JobError> {
296 if state.len() < 17 {
297 return Err(JobError::Phase(
298 "index-build checkpoint is truncated".into(),
299 ));
300 }
301 let mut epoch_bytes = [0u8; 8];
302 epoch_bytes.copy_from_slice(&state[0..8]);
303 let epoch = u64::from_le_bytes(epoch_bytes);
304 if epoch != 0 {
305 self.snapshot_epoch = Some(Epoch(epoch));
306 }
307 let mut count_bytes = [0u8; 8];
308 count_bytes.copy_from_slice(&state[8..16]);
309 self.snapshot_row_count = u64::from_le_bytes(count_bytes);
310 if state.len() >= 33 {
311 let mut built_bytes = [0u8; 8];
312 built_bytes.copy_from_slice(&state[16..24]);
313 let built = u64::from_le_bytes(built_bytes);
314 if built != 0 {
315 self.built_through = Some(Epoch(built));
316 }
317 let mut generation_bytes = [0u8; 8];
318 generation_bytes.copy_from_slice(&state[24..32]);
319 self.built_data_generation = Some(u64::from_le_bytes(generation_bytes));
320 self.published = state[32] != 0;
321 } else if state.len() >= 25 {
322 let mut built_bytes = [0u8; 8];
323 built_bytes.copy_from_slice(&state[16..24]);
324 let built = u64::from_le_bytes(built_bytes);
325 if built != 0 {
326 self.built_through = Some(Epoch(built));
327 }
328 self.published = state[24] != 0;
329 } else {
330 self.published = state[16] != 0;
331 }
332 Ok(())
333 }
334
335 fn record_pending(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
336 context.check_cancelled()?;
337 let catalog = self.db.catalog.read();
340 let entry = catalog
341 .live(&self.table)
342 .ok_or_else(|| JobError::Phase(format!("table {:?} not found", self.table)))?;
343 if entry.schema.schema_id != self.expected_schema_sequence {
344 return Err(JobError::Phase(format!(
345 "schema sequence changed before index build started on {:?}",
346 self.table
347 )));
348 }
349 let command = self.catalog_command();
350 crate::catalog_cmds::apply(&catalog, &command).map_err(job_phase_err)?;
351 Ok(())
352 }
353
354 fn pin_snapshot(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
355 context.check_cancelled()?;
356 if self.pin.is_some() {
357 return Ok(());
358 }
359 let handle = self.db.table(&self.table).map_err(job_phase_err)?;
360 let table = handle.lock();
361 let epoch = table.current_epoch();
362 let pin = Arc::clone(table.pin_registry()).pin(PinSource::OnlineIndexBuild, epoch);
363 self.snapshot_epoch = Some(epoch);
364 self.pin = Some(pin);
365 Ok(())
366 }
367
368 fn build_hidden(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
369 self.refresh_hidden(context)
370 }
371
372 fn catch_up(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
373 self.refresh_hidden(context)
374 }
375
376 fn validate(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
377 context.check_cancelled()?;
378 let catalog = self.db.catalog.read();
379 let entry = catalog
380 .live(&self.table)
381 .ok_or_else(|| JobError::Phase(format!("table {:?} not found", self.table)))?;
382 if entry.schema.schema_id != self.expected_schema_sequence {
383 return Err(JobError::Phase(format!(
384 "schema sequence conflict during index build on {:?}: expected {}, found {}",
385 self.table, self.expected_schema_sequence, entry.schema.schema_id
386 )));
387 }
388 crate::catalog_cmds::apply(&catalog, &self.catalog_command()).map_err(job_phase_err)?;
389 Ok(())
390 }
391
392 fn publish(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
393 context.check_cancelled()?;
394 if self.published {
395 return Ok(());
396 }
397 match self.db.publish_index_build(self, context) {
398 Ok(()) => {
399 self.published = true;
400 Ok(())
401 }
402 Err(error @ MongrelError::DurableCommit { .. }) => {
403 let applied = self
404 .db
405 .catalog
406 .read()
407 .live(&self.table)
408 .is_some_and(|entry| self.publication_already_applied(&entry.schema));
409 if applied {
410 self.published = true;
411 Ok(())
412 } else {
413 Err(job_phase_err(error))
414 }
415 }
416 Err(error) => Err(job_phase_err(error)),
417 }
418 }
419
420 fn release_old(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
421 context.check_cancelled()?;
422 self.release_pin();
425 self.artifact = None;
426 self.memory_reservation = None;
427 self.snapshot_rows.clear();
428 Ok(())
429 }
430
431 fn rollback(&mut self) -> std::result::Result<(), JobError> {
432 self.release_pin();
435 self.artifact = None;
436 self.memory_reservation = None;
437 self.snapshot_rows.clear();
438 Ok(())
439 }
440}
441
442fn job_phase_err(error: MongrelError) -> JobError {
443 match error {
444 MongrelError::Cancelled => JobError::Cancelled,
445 MongrelError::Conflict(message) => JobError::Phase(message),
446 other => JobError::Phase(other.to_string()),
447 }
448}
449
450impl Database {
451 pub fn index_build_target_definition(&self, job_id: u64) -> Result<IndexDef> {
458 let record = self
459 .job_registry
460 .get(job_id)
461 .ok_or_else(|| MongrelError::NotFound(format!("job {job_id} not found")))?;
462 if record.kind != JobKind::IndexBuild {
463 return Err(MongrelError::InvalidArgument(format!(
464 "job {job_id} is not an index build"
465 )));
466 }
467 let definition = record.definition.as_deref().ok_or_else(|| {
468 MongrelError::Other(format!(
469 "index-build job {job_id} has no durable definition"
470 ))
471 })?;
472 Ok(IndexBuildSpec::decode(definition)?
473 .kind
474 .definition()
475 .clone())
476 }
477
478 pub fn create_index(&self, table: &str, definition: IndexDef) -> Result<u64> {
483 self.run_index_build(
484 table,
485 IndexBuildKind::Create {
486 definition: definition.clone(),
487 },
488 Some(definition),
489 None,
490 )
491 }
492
493 pub fn submit_create_index(&self, table: &str, definition: IndexDef) -> Result<u64> {
499 let (job_id, _) = self.submit_index_build(
500 table,
501 IndexBuildKind::Create {
502 definition: definition.clone(),
503 },
504 Some(definition),
505 None,
506 )?;
507 Ok(job_id)
508 }
509
510 pub fn start_create_index(self: &Arc<Self>, table: &str, definition: IndexDef) -> Result<u64> {
512 let (job_id, spec) = self.submit_index_build(
513 table,
514 IndexBuildKind::Create {
515 definition: definition.clone(),
516 },
517 Some(definition),
518 None,
519 )?;
520 self.spawn_index_build(job_id, spec)?;
521 Ok(job_id)
522 }
523
524 pub fn replace_index(
530 &self,
531 table: &str,
532 expected_old_name: &str,
533 new_definition: IndexDef,
534 ) -> Result<u64> {
535 self.run_index_build(
536 table,
537 IndexBuildKind::Replace {
538 expected_old_name: expected_old_name.to_string(),
539 new_definition: new_definition.clone(),
540 },
541 Some(new_definition),
542 Some(expected_old_name),
543 )
544 }
545
546 pub fn submit_replace_index(
548 &self,
549 table: &str,
550 expected_old_name: &str,
551 new_definition: IndexDef,
552 ) -> Result<u64> {
553 let (job_id, _) = self.submit_index_build(
554 table,
555 IndexBuildKind::Replace {
556 expected_old_name: expected_old_name.to_string(),
557 new_definition: new_definition.clone(),
558 },
559 Some(new_definition),
560 Some(expected_old_name),
561 )?;
562 Ok(job_id)
563 }
564
565 pub fn start_replace_index(
567 self: &Arc<Self>,
568 table: &str,
569 expected_old_name: &str,
570 new_definition: IndexDef,
571 ) -> Result<u64> {
572 let (job_id, spec) = self.submit_index_build(
573 table,
574 IndexBuildKind::Replace {
575 expected_old_name: expected_old_name.to_string(),
576 new_definition: new_definition.clone(),
577 },
578 Some(new_definition),
579 Some(expected_old_name),
580 )?;
581 self.spawn_index_build(job_id, spec)?;
582 Ok(job_id)
583 }
584
585 pub fn resume_index_build(self: &Arc<Self>, job_id: u64) -> Result<()> {
588 let record = self
589 .job_registry
590 .get(job_id)
591 .ok_or_else(|| MongrelError::NotFound(format!("job {job_id}")))?;
592 if record.kind != JobKind::IndexBuild {
593 return Err(MongrelError::InvalidArgument(format!(
594 "job {job_id} is not an index build"
595 )));
596 }
597 let definition = record.definition.as_deref().ok_or_else(|| {
598 MongrelError::Other(format!(
599 "index build job {job_id} has no reconstructible definition"
600 ))
601 })?;
602 let spec = IndexBuildSpec::decode(definition)?;
603 match record.state {
604 crate::jobs::JobState::Paused => self
605 .job_registry
606 .resume(job_id)
607 .map_err(MongrelError::from)?,
608 crate::jobs::JobState::Pending => {}
609 state => {
610 return Err(MongrelError::Conflict(format!(
611 "index build job {job_id} cannot resume from {state:?}"
612 )));
613 }
614 }
615 self.spawn_index_build(job_id, spec)
616 }
617
618 pub fn drop_index(&self, table: &str, name: &str) -> Result<()> {
621 use std::sync::atomic::Ordering;
622
623 self.require(&crate::auth::Permission::Ddl)?;
624 if self.poisoned.load(Ordering::Relaxed) {
625 return Err(MongrelError::Other(
626 "database poisoned by fsync error".into(),
627 ));
628 }
629 let _operation = self.admit_operation()?;
630 let _schema_barrier = self.acquire_schema_barrier_exclusive()?;
631 let _ddl = self.ddl_lock.lock();
632 let _security_write = self.security_write()?;
633 self.require(&crate::auth::Permission::Ddl)?;
634
635 let table_id = {
636 let catalog = self.catalog.read();
637 catalog
638 .live(table)
639 .ok_or_else(|| MongrelError::NotFound(format!("table {table:?} not found")))?
640 .table_id
641 };
642 let handle = self
643 .tables
644 .read()
645 .get(&table_id)
646 .cloned()
647 .ok_or_else(|| MongrelError::NotFound(format!("table {table:?} not mounted")))?;
648
649 let command = CatalogCommand::RemoveIndex {
650 table: table.to_string(),
651 name: name.to_string(),
652 };
653 let prepared_schema = match crate::catalog_cmds::apply(&self.catalog.read(), &command)? {
655 crate::catalog_cmds::CatalogDelta::SchemaReplaced { schema, .. } => schema,
656 crate::catalog_cmds::CatalogDelta::NoOp => return Ok(()),
657 other => {
658 return Err(MongrelError::Other(format!(
659 "unexpected catalog delta for drop_index: {other:?}"
660 )));
661 }
662 };
663
664 crate::catalog::inject_hook("index.publish.before")?;
665
666 let durable_epoch = std::cell::Cell::new(None);
667 let result: Result<()> = (|| {
668 let mut table_guard = handle.lock();
669 let commit_lock = Arc::clone(&self.commit_lock);
670 let _commit = commit_lock.lock();
671 let epoch = self.epoch.bump_assigned();
672 let mut epoch_guard = EpochGuard::new(self.epoch.as_ref(), epoch);
673 let txn_id = self.alloc_txn_id()?;
674 let mut next_catalog = self.catalog.read().clone();
675 let catalog_entry_index = next_catalog
676 .tables
677 .iter()
678 .position(|entry| entry.table_id == table_id)
679 .ok_or_else(|| MongrelError::NotFound(format!("table {table:?} not found")))?;
680 self.apply_catalog_command_to(&mut next_catalog, command)?;
681 next_catalog.tables[catalog_entry_index].schema = prepared_schema.clone();
682 next_catalog.db_epoch = next_catalog.db_epoch.max(epoch.0);
683
684 let commit_seq = {
685 let mut wal = self.shared_wal.lock();
686 let append: Result<u64> = (|| {
687 append_catalog_snapshot(&mut wal, txn_id, &next_catalog)?;
688 wal.append_commit(txn_id, epoch, &[])
689 })();
690 append.map_err(|error| self.commit_outcome_unknown(epoch, error))?
691 };
692 let receipt = self.await_durable_commit(txn_id, commit_seq, epoch)?;
693 durable_epoch.set(Some(epoch));
694
695 table_guard.publish_index_drop(prepared_schema)?;
696 let schema = table_guard.schema().clone();
697 drop(table_guard);
698 next_catalog.tables[catalog_entry_index].schema = schema;
699 let catalog_result =
700 catalog::write_atomic(&self.root, &next_catalog, self.meta_dek.as_ref());
701 *self.catalog.write() = next_catalog;
702 self.publish_committed(&receipt, epoch)?;
703 epoch_guard.disarm();
704 if let Err(error) = catalog_result {
705 self.poisoned.store(true, Ordering::Relaxed);
706 self.lifecycle.poison();
707 return Err(MongrelError::DurableCommit {
708 epoch: epoch.0,
709 message: error.to_string(),
710 });
711 }
712 crate::catalog::inject_hook("index.publish.after")?;
713 Ok(())
714 })();
715 result.map_err(|error| match (durable_epoch.get(), error) {
716 (_, error @ MongrelError::DurableCommit { .. }) => error,
717 (Some(epoch), error) => MongrelError::DurableCommit {
718 epoch: epoch.0,
719 message: error.to_string(),
720 },
721 (None, error) => error,
722 })
723 }
724
725 fn run_index_build(
726 &self,
727 table: &str,
728 kind: IndexBuildKind,
729 definition: Option<IndexDef>,
730 expected_old_name: Option<&str>,
731 ) -> Result<u64> {
732 let (job_id, spec) = self.submit_index_build(table, kind, definition, expected_old_name)?;
733 self.drive_index_build(job_id, spec)
734 }
735
736 fn submit_index_build(
737 &self,
738 table: &str,
739 kind: IndexBuildKind,
740 definition: Option<IndexDef>,
741 expected_old_name: Option<&str>,
742 ) -> Result<(u64, IndexBuildSpec)> {
743 self.require(&crate::auth::Permission::Ddl)?;
744 if self.poisoned.load(Ordering::Relaxed) {
745 return Err(MongrelError::Other(
746 "database poisoned by fsync error".into(),
747 ));
748 }
749 if table.is_empty() {
750 return Err(MongrelError::InvalidArgument(
751 "index DDL requires a non-empty table name".into(),
752 ));
753 }
754
755 let (expected_schema_sequence, index_name) = {
757 let catalog = self.catalog.read();
758 let entry = catalog
759 .live(table)
760 .ok_or_else(|| MongrelError::NotFound(format!("table {table:?} not found")))?;
761 if let Some(old) = expected_old_name {
762 if !entry.schema.indexes.iter().any(|index| index.name == old) {
763 return Err(MongrelError::NotFound(format!(
764 "index {old} does not exist on {table}"
765 )));
766 }
767 }
768 if let Some(def) = &definition {
769 def.validate_options()?;
770 }
771 let command = match &kind {
774 IndexBuildKind::Create { definition } => CatalogCommand::AddIndex {
775 table: table.to_string(),
776 index: definition.clone(),
777 },
778 IndexBuildKind::Replace {
779 expected_old_name,
780 new_definition,
781 } => CatalogCommand::ReplaceIndex {
782 table: table.to_string(),
783 expected_schema_sequence: entry.schema.schema_id,
784 expected_old_name: expected_old_name.clone(),
785 new_definition: new_definition.clone(),
786 },
787 };
788 crate::catalog_cmds::apply(&catalog, &command)?;
789 let name = match &kind {
790 IndexBuildKind::Create { definition } => definition.name.clone(),
791 IndexBuildKind::Replace { new_definition, .. } => new_definition.name.clone(),
792 };
793 (entry.schema.schema_id, name)
794 };
795
796 let _ = self.table(table)?;
798
799 let spec = IndexBuildSpec {
800 version: INDEX_BUILD_SPEC_VERSION,
801 table: table.to_string(),
802 expected_schema_sequence,
803 kind,
804 };
805 let job_id = self.job_registry.submit_with_definition(
806 JobKind::IndexBuild,
807 JobTarget {
808 table: table.to_string(),
809 index: Some(index_name),
810 },
811 Some(spec.encode()?),
812 )?;
813 Ok((job_id, spec))
814 }
815
816 fn spawn_index_build(self: &Arc<Self>, job_id: u64, spec: IndexBuildSpec) -> Result<()> {
817 let db = Arc::clone(self);
818 std::thread::Builder::new()
819 .name(format!("mongreldb-index-{job_id}"))
820 .spawn(move || {
821 let _ = db.drive_index_build(job_id, spec);
823 })
824 .map_err(|error| {
825 MongrelError::Io(std::io::Error::new(
826 error.kind(),
827 format!("spawn index build job {job_id}: {error}"),
828 ))
829 })?;
830 Ok(())
831 }
832
833 fn drive_index_build(&self, job_id: u64, spec: IndexBuildSpec) -> Result<u64> {
834 let mut job = IndexBuildJob {
835 db: self,
836 table: spec.table,
837 kind: spec.kind,
838 expected_schema_sequence: spec.expected_schema_sequence,
839 snapshot_epoch: None,
840 pin: None,
841 snapshot_row_count: 0,
842 snapshot_rows: std::collections::BTreeMap::new(),
843 artifact: None,
844 memory_reservation: None,
845 built_through: None,
846 built_data_generation: None,
847 published: false,
848 };
849
850 match run_build_publish(self.job_registry.as_ref(), job_id, &mut job) {
851 Ok(()) => {
852 let record = self.job_registry.get(job_id);
853 match record.map(|r| r.state) {
854 Some(crate::jobs::JobState::Succeeded) => Ok(job_id),
855 Some(crate::jobs::JobState::Paused) => Err(MongrelError::Other(format!(
856 "index build job {job_id} parked mid-run; resume and redrive"
857 ))),
858 Some(state) => Err(MongrelError::Other(format!(
859 "index build job {job_id} ended in unexpected state {state:?}"
860 ))),
861 None => Err(MongrelError::NotFound(format!("job {job_id}"))),
862 }
863 }
864 Err(JobError::Cancelled) => {
865 if job.published
869 || self
870 .job_registry
871 .get(job_id)
872 .is_some_and(|r| r.state == crate::jobs::JobState::Succeeded)
873 {
874 Ok(job_id)
875 } else {
876 Err(MongrelError::Cancelled)
877 }
878 }
879 Err(error) => {
880 if job.published {
881 Ok(job_id)
883 } else {
884 Err(error.into())
885 }
886 }
887 }
888 }
889
890 fn publish_index_build(&self, job: &mut IndexBuildJob<'_>, context: &JobContext) -> Result<()> {
895 use std::sync::atomic::Ordering;
896
897 if self.poisoned.load(Ordering::Relaxed) {
898 return Err(MongrelError::Other(
899 "database poisoned by fsync error".into(),
900 ));
901 }
902 self.require(&crate::auth::Permission::Ddl)?;
904
905 let table_name = job.table.clone();
906 let table_id = {
907 let catalog = self.catalog.read();
908 let entry = catalog
909 .live(&table_name)
910 .ok_or_else(|| MongrelError::NotFound(format!("table {table_name:?} not found")))?;
911 if entry.schema.schema_id != job.expected_schema_sequence {
912 if job.publication_already_applied(&entry.schema) {
913 job.published = true;
914 return Ok(());
915 }
916 return Err(MongrelError::Conflict(format!(
917 "index publish on {table_name}: expected schema sequence {}, found {}",
918 job.expected_schema_sequence, entry.schema.schema_id
919 )));
920 }
921 entry.table_id
922 };
923 let handle =
924 self.tables.read().get(&table_id).cloned().ok_or_else(|| {
925 MongrelError::NotFound(format!("table {table_name:?} not mounted"))
926 })?;
927
928 loop {
929 context.check_cancelled().map_err(MongrelError::from)?;
930 if job.artifact.is_none() || job.built_data_generation.is_none() {
931 job.refresh_hidden(context).map_err(MongrelError::from)?;
932 }
933 let built_data_generation = job.built_data_generation.ok_or_else(|| {
934 MongrelError::Other("index build has no content watermark".into())
935 })?;
936 let command = job.catalog_command();
937 let prepared_schema = match crate::catalog_cmds::apply(&self.catalog.read(), &command)?
938 {
939 crate::catalog_cmds::CatalogDelta::SchemaReplaced { schema, .. } => schema,
940 other => {
941 return Err(MongrelError::Other(format!(
942 "unexpected catalog delta for index publish: {other:?}"
943 )));
944 }
945 };
946 let durable_epoch = std::cell::Cell::new(None);
947
948 let _schema_barrier = self.acquire_schema_barrier_exclusive()?;
951 let _ddl = self.ddl_lock.lock();
952 let _security_write = self.security_write()?;
953 self.require(&crate::auth::Permission::Ddl)?;
954
955 {
957 let catalog = self.catalog.read();
958 let entry = catalog.live(&table_name).ok_or_else(|| {
959 MongrelError::NotFound(format!("table {table_name:?} not found"))
960 })?;
961 if entry.schema.schema_id != job.expected_schema_sequence {
962 if job.publication_already_applied(&entry.schema) {
963 job.published = true;
964 return Ok(());
965 }
966 return Err(MongrelError::Conflict(format!(
967 "index publish on {table_name}: expected schema sequence {}, found {}",
968 job.expected_schema_sequence, entry.schema.schema_id
969 )));
970 }
971 }
972
973 let commit_lock = Arc::clone(&self.commit_lock);
975 let _commit = commit_lock.lock();
976 let mut table_guard = handle.lock();
977 if table_guard.data_generation() != built_data_generation {
978 drop(table_guard);
979 drop(_commit);
980 drop(_security_write);
981 drop(_ddl);
982 drop(_schema_barrier);
983 job.refresh_hidden(context).map_err(MongrelError::from)?;
984 continue;
985 }
986
987 let result: Result<()> = (|| {
988 crate::catalog::inject_hook("index.publish.before")?;
991
992 let epoch = self.epoch.bump_assigned();
993 let mut epoch_guard = EpochGuard::new(self.epoch.as_ref(), epoch);
994 let txn_id = self.alloc_txn_id()?;
995 let mut next_catalog = self.catalog.read().clone();
996 let catalog_entry_index = next_catalog
997 .tables
998 .iter()
999 .position(|entry| entry.table_id == table_id)
1000 .ok_or_else(|| {
1001 MongrelError::NotFound(format!("table {table_name:?} not found"))
1002 })?;
1003 self.apply_catalog_command_to(&mut next_catalog, command)?;
1004 next_catalog.tables[catalog_entry_index].schema = prepared_schema.clone();
1005 next_catalog.db_epoch = next_catalog.db_epoch.max(epoch.0);
1006
1007 let commit_seq = {
1008 let mut wal = self.shared_wal.lock();
1009 let append: Result<u64> = (|| {
1010 let _ = DdlOp::encode_schema(&prepared_schema)?;
1012 append_catalog_snapshot(&mut wal, txn_id, &next_catalog)?;
1013 wal.append_commit(txn_id, epoch, &[])
1014 })();
1015 append.map_err(|error| self.commit_outcome_unknown(epoch, error))?
1016 };
1017 let receipt = self.await_durable_commit(txn_id, commit_seq, epoch)?;
1018 durable_epoch.set(Some(epoch));
1019
1020 let artifact = job.artifact.take().ok_or_else(|| {
1021 MongrelError::Other("index build lost its hidden artifact".into())
1022 })?;
1023 table_guard.publish_index_schema_change(prepared_schema, artifact)?;
1024 let schema = table_guard.schema().clone();
1025 drop(table_guard);
1026
1027 next_catalog.tables[catalog_entry_index].schema = schema;
1028 let catalog_result =
1029 catalog::write_atomic(&self.root, &next_catalog, self.meta_dek.as_ref());
1030 *self.catalog.write() = next_catalog;
1031 self.publish_committed(&receipt, epoch)?;
1032 epoch_guard.disarm();
1033 if let Err(error) = catalog_result {
1034 self.poisoned.store(true, Ordering::Relaxed);
1035 self.lifecycle.poison();
1036 return Err(MongrelError::DurableCommit {
1037 epoch: epoch.0,
1038 message: error.to_string(),
1039 });
1040 }
1041 crate::catalog::inject_hook("index.publish.after")?;
1042 Ok(())
1043 })();
1044 return result.map_err(|error| match (durable_epoch.get(), error) {
1045 (_, error @ MongrelError::DurableCommit { .. }) => error,
1046 (Some(epoch), error) => MongrelError::DurableCommit {
1047 epoch: epoch.0,
1048 message: error.to_string(),
1049 },
1050 (None, error) => error,
1051 });
1052 }
1053 }
1054}
1055
1056#[cfg(test)]
1057mod tests {
1058 use super::*;
1059 use crate::query::{Retriever, RetrieverScore};
1060 use crate::schema::{AnnOptions, AnnQuantization, IndexKind, IndexOptions, TypeId};
1061 use std::sync::Mutex;
1062
1063 static INDEX_DDL_TEST_LOCK: Mutex<()> = Mutex::new(());
1066
1067 fn lock_tests() -> std::sync::MutexGuard<'static, ()> {
1068 INDEX_DDL_TEST_LOCK
1069 .lock()
1070 .unwrap_or_else(|poisoned| poisoned.into_inner())
1071 }
1072
1073 fn embedding_schema(quantization: AnnQuantization) -> Schema {
1074 Schema {
1075 columns: vec![
1076 ColumnDef {
1077 id: 1,
1078 name: "id".into(),
1079 ty: TypeId::Int64,
1080 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
1081 default_value: None,
1082 embedding_source: None,
1083 },
1084 ColumnDef {
1085 id: 2,
1086 name: "embedding".into(),
1087 ty: TypeId::Embedding { dim: 4 },
1088 flags: ColumnFlags::empty(),
1089 default_value: None,
1090 embedding_source: None,
1091 },
1092 ],
1093 indexes: vec![IndexDef {
1094 name: "idx_embed_ann".into(),
1095 column_id: 2,
1096 kind: IndexKind::Ann,
1097 predicate: None,
1098 options: IndexOptions {
1099 ann: Some(AnnOptions {
1100 m: 8,
1101 ef_construction: 32,
1102 ef_search: 16,
1103 quantization,
1104 }),
1105 ..IndexOptions::default()
1106 },
1107 }],
1108 ..Schema::default()
1109 }
1110 }
1111
1112 fn plain_embedding_schema() -> Schema {
1113 Schema {
1114 columns: vec![
1115 ColumnDef {
1116 id: 1,
1117 name: "id".into(),
1118 ty: TypeId::Int64,
1119 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
1120 default_value: None,
1121 embedding_source: None,
1122 },
1123 ColumnDef {
1124 id: 2,
1125 name: "embedding".into(),
1126 ty: TypeId::Embedding { dim: 4 },
1127 flags: ColumnFlags::empty(),
1128 default_value: None,
1129 embedding_source: None,
1130 },
1131 ],
1132 ..Schema::default()
1133 }
1134 }
1135
1136 fn put_row(db: &Database, table: &str, id: i64, embedding: Vec<f32>) {
1137 let mut txn = db.begin();
1138 txn.put(
1139 table,
1140 vec![(1, Value::Int64(id)), (2, Value::Embedding(embedding))],
1141 )
1142 .unwrap();
1143 txn.commit().unwrap();
1144 }
1145
1146 fn row_id_for_pk(db: &Database, table: &str, id: i64) -> RowId {
1147 let handle = db.table(table).unwrap();
1148 let table = handle.lock();
1149 table
1150 .visible_rows(Snapshot::at(table.current_epoch()))
1151 .unwrap()
1152 .into_iter()
1153 .find(|row| row.columns.get(&1) == Some(&Value::Int64(id)))
1154 .map(|row| row.row_id)
1155 .unwrap()
1156 }
1157
1158 fn dense_definition() -> IndexDef {
1159 IndexDef {
1160 name: "idx_embed_ann".into(),
1161 column_id: 2,
1162 kind: IndexKind::Ann,
1163 predicate: None,
1164 options: IndexOptions {
1165 ann: Some(AnnOptions {
1166 quantization: AnnQuantization::Dense,
1167 ..AnnOptions::default()
1168 }),
1169 ..IndexOptions::default()
1170 },
1171 }
1172 }
1173
1174 #[test]
1175 fn create_index_indexes_snapshot_rows() {
1176 let _lock = lock_tests();
1177 let dir = tempfile::tempdir().unwrap();
1178 let db = Database::create(dir.path()).unwrap();
1179 db.create_table("docs", plain_embedding_schema()).unwrap();
1180 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1181 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1182
1183 let definition = IndexDef {
1184 name: "idx_embed_ann".into(),
1185 column_id: 2,
1186 kind: IndexKind::Ann,
1187 predicate: None,
1188 options: IndexOptions {
1189 ann: Some(AnnOptions {
1190 m: 8,
1191 ef_construction: 32,
1192 ef_search: 16,
1193 quantization: AnnQuantization::Dense,
1194 }),
1195 ..IndexOptions::default()
1196 },
1197 };
1198 let job_id = db.create_index("docs", definition).unwrap();
1199 let record = db.job_registry().get(job_id).unwrap();
1200 assert_eq!(record.state, crate::jobs::JobState::Succeeded);
1201
1202 let handle = db.table("docs").unwrap();
1203 let mut table = handle.lock();
1204 let hits = table
1205 .retrieve(&Retriever::Ann {
1206 column_id: 2,
1207 query: vec![1.0, 0.0, 0.0, 0.0],
1208 k: 2,
1209 })
1210 .unwrap();
1211 assert_eq!(hits.len(), 2);
1212 assert!(matches!(
1213 hits[0].score,
1214 RetrieverScore::AnnCosineDistance(d) if d.abs() < 1e-5
1215 ));
1216 assert_eq!(table.schema().schema_id, table.table_id().saturating_add(1));
1219 }
1220
1221 #[test]
1222 fn replace_binary_sign_with_dense_search_works() {
1223 let _lock = lock_tests();
1224 let dir = tempfile::tempdir().unwrap();
1225 let db = Database::create(dir.path()).unwrap();
1226 db.create_table("docs", embedding_schema(AnnQuantization::BinarySign))
1227 .unwrap();
1228 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1229 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1230
1231 let new_def = IndexDef {
1232 name: "idx_embed_ann".into(),
1233 column_id: 2,
1234 kind: IndexKind::Ann,
1235 predicate: None,
1236 options: IndexOptions {
1237 ann: Some(AnnOptions {
1238 m: 8,
1239 ef_construction: 32,
1240 ef_search: 16,
1241 quantization: AnnQuantization::Dense,
1242 }),
1243 ..IndexOptions::default()
1244 },
1245 };
1246 db.replace_index("docs", "idx_embed_ann", new_def).unwrap();
1247
1248 let handle = db.table("docs").unwrap();
1249 let mut table = handle.lock();
1250 let hits = table
1251 .retrieve(&Retriever::Ann {
1252 column_id: 2,
1253 query: vec![1.0, 0.0, 0.0, 0.0],
1254 k: 1,
1255 })
1256 .unwrap();
1257 assert_eq!(hits.len(), 1);
1258 assert!(matches!(
1259 hits[0].score,
1260 RetrieverScore::AnnCosineDistance(_)
1261 ));
1262 assert!(!hits
1263 .iter()
1264 .any(|hit| matches!(hit.score, RetrieverScore::AnnHammingDistance(_))));
1265 }
1266
1267 #[test]
1268 fn drop_index_removes_schema_entry_without_rewrite() {
1269 let _lock = lock_tests();
1270 let dir = tempfile::tempdir().unwrap();
1271 let db = Database::create(dir.path()).unwrap();
1272 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1273 .unwrap();
1274 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1275 db.drop_index("docs", "idx_embed_ann").unwrap();
1276 let handle = db.table("docs").unwrap();
1277 let table = handle.lock();
1278 assert!(table.schema().indexes.is_empty());
1279 let rows = table
1281 .visible_rows(Snapshot::at(table.current_epoch()))
1282 .unwrap();
1283 assert_eq!(rows.len(), 1);
1284 }
1285
1286 #[test]
1287 fn concurrent_writes_progress_during_create() {
1288 use std::sync::Barrier;
1289 let _lock = lock_tests();
1290 let dir = tempfile::tempdir().unwrap();
1291 let db = Arc::new(Database::create(dir.path()).unwrap());
1292 db.create_table("docs", plain_embedding_schema()).unwrap();
1293 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1294
1295 let start = Arc::new(Barrier::new(2));
1296 let writer = {
1297 let db = Arc::clone(&db);
1298 let start = Arc::clone(&start);
1299 std::thread::spawn(move || {
1300 start.wait();
1301 for id in 100..120 {
1302 put_row(&db, "docs", id, vec![0.0, 1.0, 0.0, 0.0]);
1303 }
1304 })
1305 };
1306 start.wait();
1307 let definition = IndexDef {
1308 name: "idx_embed_ann".into(),
1309 column_id: 2,
1310 kind: IndexKind::Ann,
1311 predicate: None,
1312 options: IndexOptions {
1313 ann: Some(AnnOptions {
1314 quantization: AnnQuantization::Dense,
1315 ..AnnOptions::default()
1316 }),
1317 ..IndexOptions::default()
1318 },
1319 };
1320 db.create_index("docs", definition).unwrap();
1321 writer.join().unwrap();
1322 let handle = db.table("docs").unwrap();
1324 let table = handle.lock();
1325 assert!(table
1326 .schema()
1327 .indexes
1328 .iter()
1329 .any(|index| index.name == "idx_embed_ann"));
1330 }
1331
1332 #[test]
1333 fn asynchronous_build_is_observable_and_replays_all_row_deltas() {
1334 use std::sync::Barrier;
1335 use std::time::Duration;
1336
1337 let _lock = lock_tests();
1338 let dir = tempfile::tempdir().unwrap();
1339 let db = Arc::new(Database::create(dir.path()).unwrap());
1340 db.create_table("docs", plain_embedding_schema()).unwrap();
1341 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1342 put_row(&db, "docs", 3, vec![0.0, 0.0, 1.0, 0.0]);
1343 let deleted_row = row_id_for_pk(&db, "docs", 3);
1344 let baseline = db
1345 .memory_governor()
1346 .usage(crate::memory::MemoryClass::Compaction);
1347
1348 let entered = Arc::new(Barrier::new(2));
1349 let release = Arc::new(Barrier::new(2));
1350 let callback_entered = Arc::clone(&entered);
1351 let callback_release = Arc::clone(&release);
1352 let _guard = mongreldb_fault::ScopedGuard::new(
1353 "job.build_hidden.after",
1354 mongreldb_fault::Action::Callback(Arc::new(move |_| {
1355 callback_entered.wait();
1356 callback_release.wait();
1357 })),
1358 );
1359
1360 let job_id = db.start_create_index("docs", dense_definition()).unwrap();
1361 entered.wait();
1362
1363 let running = db.job_registry().get(job_id).unwrap();
1364 assert_eq!(running.state, crate::jobs::JobState::Running);
1365 let persisted = IndexBuildSpec::decode(running.definition.as_deref().unwrap()).unwrap();
1366 assert_eq!(persisted.table, "docs");
1367 assert_eq!(persisted.kind.definition(), &dense_definition());
1368 assert_eq!(
1369 db.index_build_target_definition(job_id).unwrap(),
1370 dense_definition()
1371 );
1372 assert!(
1373 db.memory_governor()
1374 .usage(crate::memory::MemoryClass::Compaction)
1375 > baseline,
1376 "hidden generation must hold a governor reservation"
1377 );
1378
1379 put_row(&db, "docs", 1, vec![0.0, 1.0, 0.0, 0.0]);
1382 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1383 db.transaction(|txn| txn.delete("docs", deleted_row))
1384 .unwrap();
1385
1386 release.wait();
1387 let terminal = db
1388 .job_registry()
1389 .wait_terminal(job_id, Duration::from_secs(30))
1390 .unwrap();
1391 assert_eq!(terminal.state, crate::jobs::JobState::Succeeded);
1392 assert_eq!(
1393 db.memory_governor()
1394 .usage(crate::memory::MemoryClass::Compaction),
1395 baseline,
1396 "successful build must release its reservation"
1397 );
1398
1399 let expected: std::collections::HashSet<_> = [1, 2]
1400 .into_iter()
1401 .map(|id| row_id_for_pk(&db, "docs", id))
1402 .collect();
1403 let handle = db.table("docs").unwrap();
1404 let mut table = handle.lock();
1405 let actual: std::collections::HashSet<_> = table
1406 .retrieve(&Retriever::Ann {
1407 column_id: 2,
1408 query: vec![0.0, 1.0, 0.0, 0.0],
1409 k: 10,
1410 })
1411 .unwrap()
1412 .into_iter()
1413 .map(|hit| hit.row_id)
1414 .collect();
1415 assert_eq!(actual, expected);
1416 }
1417
1418 #[test]
1419 fn cancelling_real_hidden_build_rolls_back_pin_memory_and_schema() {
1420 use std::sync::Barrier;
1421 use std::time::Duration;
1422
1423 let _lock = lock_tests();
1424 let dir = tempfile::tempdir().unwrap();
1425 let db = Arc::new(Database::create(dir.path()).unwrap());
1426 db.create_table("docs", plain_embedding_schema()).unwrap();
1427 for id in 0..32 {
1428 put_row(&db, "docs", id, vec![1.0, id as f32, 0.0, 0.0]);
1429 }
1430 let baseline = db
1431 .memory_governor()
1432 .usage(crate::memory::MemoryClass::Compaction);
1433
1434 let entered = Arc::new(Barrier::new(2));
1435 let release = Arc::new(Barrier::new(2));
1436 let callback_entered = Arc::clone(&entered);
1437 let callback_release = Arc::clone(&release);
1438 let _guard = mongreldb_fault::ScopedGuard::new(
1439 "job.build_hidden.after",
1440 mongreldb_fault::Action::Callback(Arc::new(move |_| {
1441 callback_entered.wait();
1442 callback_release.wait();
1443 })),
1444 );
1445
1446 let job_id = db.start_create_index("docs", dense_definition()).unwrap();
1447 entered.wait();
1448 db.job_registry().cancel(job_id).unwrap();
1449 release.wait();
1450 let terminal = db
1451 .job_registry()
1452 .wait_terminal(job_id, Duration::from_secs(30))
1453 .unwrap();
1454 assert_eq!(terminal.state, crate::jobs::JobState::Failed);
1455 assert!(terminal
1456 .error
1457 .as_deref()
1458 .is_some_and(|error| error.contains("cancel")));
1459 assert_eq!(
1460 db.memory_governor()
1461 .usage(crate::memory::MemoryClass::Compaction),
1462 baseline
1463 );
1464 assert!(
1465 db.table("docs").unwrap().lock().schema().indexes.is_empty(),
1466 "pre-publication cancellation must keep the old schema"
1467 );
1468 }
1469
1470 #[test]
1471 fn pending_index_job_reconstructs_and_resumes_after_reopen() {
1472 use std::time::Duration;
1473
1474 let _lock = lock_tests();
1475 let dir = tempfile::tempdir().unwrap();
1476 let path = dir.path().to_path_buf();
1477 let job_id = {
1478 let db = Database::create(&path).unwrap();
1479 db.create_table("docs", plain_embedding_schema()).unwrap();
1480 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1481 let job_id = db.submit_create_index("docs", dense_definition()).unwrap();
1482 let pending = db.job_registry().get(job_id).unwrap();
1483 assert_eq!(pending.state, crate::jobs::JobState::Pending);
1484 assert!(pending.definition.is_some());
1485 job_id
1486 };
1487
1488 let db = Arc::new(Database::open(&path).unwrap());
1489 db.resume_index_build(job_id).unwrap();
1490 let terminal = db
1491 .job_registry()
1492 .wait_terminal(job_id, Duration::from_secs(30))
1493 .unwrap();
1494 assert_eq!(terminal.state, crate::jobs::JobState::Succeeded);
1495 assert_eq!(
1496 db.table("docs")
1497 .unwrap()
1498 .lock()
1499 .schema()
1500 .indexes
1501 .iter()
1502 .find(|index| index.name == "idx_embed_ann")
1503 .and_then(|index| index.options.ann.as_ref())
1504 .map(|options| options.quantization),
1505 Some(AnnQuantization::Dense)
1506 );
1507 }
1508
1509 #[test]
1510 fn dense_build_is_rejected_by_memory_admission_before_graph_allocation() {
1511 let _lock = lock_tests();
1512 let dir = tempfile::tempdir().unwrap();
1513 let path = dir.path().to_path_buf();
1514 let dim = 4_096u32;
1515 {
1516 let db = Database::create(&path).unwrap();
1517 let mut schema = plain_embedding_schema();
1518 schema.columns[1].ty = TypeId::Embedding { dim };
1519 db.create_table("docs", schema).unwrap();
1520 for id in 0..64 {
1521 let mut vector = vec![0.0; dim as usize];
1522 vector[id as usize] = 1.0;
1523 put_row(&db, "docs", id, vector);
1524 }
1525 }
1526
1527 let db = Database::open_with_options(
1528 &path,
1529 crate::database::OpenOptions::default().with_memory_budget_bytes(2 * 1024 * 1024),
1530 )
1531 .unwrap();
1532 let baseline = db
1533 .memory_governor()
1534 .usage(crate::memory::MemoryClass::Compaction);
1535 let error = db
1536 .create_index("docs", dense_definition())
1537 .expect_err("oversized Dense graph must be denied");
1538 assert!(
1539 matches!(
1540 error,
1541 MongrelError::ResourceLimitExceeded {
1542 resource: "index build memory",
1543 ..
1544 }
1545 ),
1546 "unexpected admission error: {error:?}"
1547 );
1548 assert_eq!(
1549 db.memory_governor()
1550 .usage(crate::memory::MemoryClass::Compaction),
1551 baseline,
1552 "rejected build must not leak a reservation"
1553 );
1554 assert!(db.table("docs").unwrap().lock().schema().indexes.is_empty());
1555 }
1556
1557 #[test]
1558 fn ddl_auth_fail_closed() {
1559 let _lock = lock_tests();
1560 let dir = tempfile::tempdir().unwrap();
1561 let path = dir.path().to_path_buf();
1562 {
1563 let db = Database::create_with_credentials(&path, "admin", "admin-pw").unwrap();
1564 db.create_table("docs", plain_embedding_schema()).unwrap();
1565 db.create_user("alice", "alice-pw").unwrap();
1566 db.create_role("reader").unwrap();
1567 db.grant_permission(
1568 "reader",
1569 crate::auth::Permission::Select {
1570 table: "docs".into(),
1571 },
1572 )
1573 .unwrap();
1574 db.grant_role("alice", "reader").unwrap();
1575 }
1576 let db = Database::open_with_credentials(&path, "alice", "alice-pw").unwrap();
1577 let definition = IndexDef {
1578 name: "idx".into(),
1579 column_id: 2,
1580 kind: IndexKind::Ann,
1581 predicate: None,
1582 options: IndexOptions {
1583 ann: Some(AnnOptions {
1584 quantization: AnnQuantization::Dense,
1585 ..AnnOptions::default()
1586 }),
1587 ..IndexOptions::default()
1588 },
1589 };
1590 let err = db.create_index("docs", definition).unwrap_err();
1591 assert!(
1592 matches!(err, MongrelError::PermissionDenied { .. }),
1593 "unexpected error: {err:?}"
1594 );
1595 }
1596
1597 #[test]
1598 fn cluster_replica_rejects_local_index_ddl() {
1599 let _lock = lock_tests();
1600 use mongreldb_types::ids::{ClusterId, DatabaseId, NodeId};
1601 let dir = tempfile::tempdir().unwrap();
1602 let cluster_id = ClusterId::from_bytes([1; 16]);
1603 let node_id = NodeId::from_bytes([2; 16]);
1604 let database_id = DatabaseId::from_bytes([3; 16]);
1605 let db =
1606 Database::create_cluster_replica(dir.path(), cluster_id, node_id, database_id).unwrap();
1607 assert!(db.is_read_only_replica());
1608 let definition = IndexDef {
1609 name: "idx".into(),
1610 column_id: 2,
1611 kind: IndexKind::Ann,
1612 predicate: None,
1613 options: Default::default(),
1614 };
1615 let err = db.create_index("docs", definition).unwrap_err();
1616 assert!(
1617 matches!(err, MongrelError::ReadOnlyReplica),
1618 "cluster replica must reject local index DDL: {err:?}"
1619 );
1620 }
1621
1622 #[test]
1623 fn fault_before_publish_leaves_old_schema() {
1624 let _lock = lock_tests();
1625 let _guard = mongreldb_fault::ScopedGuard::new(
1628 "index.publish.before",
1629 mongreldb_fault::Action::Fail,
1630 );
1631 let dir = tempfile::tempdir().unwrap();
1632 let db = Database::create(dir.path()).unwrap();
1633 db.create_table("docs", plain_embedding_schema()).unwrap();
1634 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1635
1636 let definition = IndexDef {
1637 name: "idx_embed_ann".into(),
1638 column_id: 2,
1639 kind: IndexKind::Ann,
1640 predicate: None,
1641 options: IndexOptions {
1642 ann: Some(AnnOptions {
1643 quantization: AnnQuantization::Dense,
1644 ..AnnOptions::default()
1645 }),
1646 ..IndexOptions::default()
1647 },
1648 };
1649 let err = db.create_index("docs", definition).unwrap_err();
1650 assert!(
1651 !matches!(err, MongrelError::NotFound(_)),
1652 "should fail at publish, not earlier: {err:?}"
1653 );
1654 let handle = db.table("docs").unwrap();
1655 let table = handle.lock();
1656 assert!(
1657 table.schema().indexes.is_empty(),
1658 "pre-publication fault must leave the old schema active"
1659 );
1660 }
1661
1662 #[test]
1663 fn cancel_pre_publish_leaves_old_index() {
1664 let _lock = lock_tests();
1665 let dir = tempfile::tempdir().unwrap();
1667 let db = Database::create(dir.path()).unwrap();
1668 db.create_table("docs", embedding_schema(AnnQuantization::BinarySign))
1669 .unwrap();
1670 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1671
1672 let job_id = db
1674 .job_registry()
1675 .submit(
1676 JobKind::IndexBuild,
1677 JobTarget {
1678 table: "docs".into(),
1679 index: Some("idx_embed_ann".into()),
1680 },
1681 )
1682 .unwrap();
1683 db.job_registry().cancel(job_id).unwrap();
1684 let record = db.job_registry().get(job_id).unwrap();
1685 assert_eq!(record.state, crate::jobs::JobState::Failed);
1686
1687 let handle = db.table("docs").unwrap();
1689 let table = handle.lock();
1690 let index = table
1691 .schema()
1692 .indexes
1693 .iter()
1694 .find(|index| index.name == "idx_embed_ann")
1695 .unwrap();
1696 assert_eq!(
1697 index
1698 .options
1699 .ann
1700 .as_ref()
1701 .map(|options| options.quantization),
1702 Some(AnnQuantization::BinarySign)
1703 );
1704 }
1705
1706 #[test]
1707 fn replace_dense_with_binary_sign_search_works() {
1708 let _lock = lock_tests();
1709 let dir = tempfile::tempdir().unwrap();
1710 let db = Database::create(dir.path()).unwrap();
1711 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1712 .unwrap();
1713 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1714 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1715
1716 let new_def = IndexDef {
1717 name: "idx_embed_ann".into(),
1718 column_id: 2,
1719 kind: IndexKind::Ann,
1720 predicate: None,
1721 options: IndexOptions {
1722 ann: Some(AnnOptions {
1723 m: 8,
1724 ef_construction: 32,
1725 ef_search: 16,
1726 quantization: AnnQuantization::BinarySign,
1727 }),
1728 ..IndexOptions::default()
1729 },
1730 };
1731 db.replace_index("docs", "idx_embed_ann", new_def).unwrap();
1732
1733 let handle = db.table("docs").unwrap();
1734 let mut table = handle.lock();
1735 let hits = table
1736 .retrieve(&Retriever::Ann {
1737 column_id: 2,
1738 query: vec![1.0, 0.0, 0.0, 0.0],
1739 k: 1,
1740 })
1741 .unwrap();
1742 assert_eq!(hits.len(), 1);
1743 assert!(matches!(
1744 hits[0].score,
1745 RetrieverScore::AnnHammingDistance(_)
1746 ));
1747 assert!(!hits
1748 .iter()
1749 .any(|hit| matches!(hit.score, RetrieverScore::AnnCosineDistance(_))));
1750 }
1751
1752 #[test]
1753 fn fault_after_publish_keeps_new_index() {
1754 let _lock = lock_tests();
1755 let _guard =
1758 mongreldb_fault::ScopedGuard::new("index.publish.after", mongreldb_fault::Action::Fail);
1759 let dir = tempfile::tempdir().unwrap();
1760 let db = Database::create(dir.path()).unwrap();
1761 db.create_table("docs", plain_embedding_schema()).unwrap();
1762 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1763
1764 let definition = IndexDef {
1765 name: "idx_embed_ann".into(),
1766 column_id: 2,
1767 kind: IndexKind::Ann,
1768 predicate: None,
1769 options: IndexOptions {
1770 ann: Some(AnnOptions {
1771 quantization: AnnQuantization::Dense,
1772 ..AnnOptions::default()
1773 }),
1774 ..IndexOptions::default()
1775 },
1776 };
1777 let _ = db.create_index("docs", definition);
1780 let handle = db.table("docs").unwrap();
1781 let table = handle.lock();
1782 let index = table
1783 .schema()
1784 .indexes
1785 .iter()
1786 .find(|index| index.name == "idx_embed_ann");
1787 assert!(
1788 index.is_some(),
1789 "post-publish fault must keep the new index active"
1790 );
1791 assert_eq!(
1792 index
1793 .unwrap()
1794 .options
1795 .ann
1796 .as_ref()
1797 .map(|options| options.quantization),
1798 Some(AnnQuantization::Dense)
1799 );
1800 }
1801
1802 #[test]
1803 fn dense_plaintext_close_reopen_preserves_search() {
1804 let _lock = lock_tests();
1805 let dir = tempfile::tempdir().unwrap();
1806 let path = dir.path().to_path_buf();
1807 {
1808 let db = Database::create(&path).unwrap();
1809 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1810 .unwrap();
1811 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1812 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1813 let handle = db.table("docs").unwrap();
1815 let mut table = handle.lock();
1816 table.flush().unwrap();
1817 let _ = table.publish_read_generation().unwrap();
1818 }
1819 let db = Database::open(&path).unwrap();
1820 let handle = db.table("docs").unwrap();
1821 let mut table = handle.lock();
1822 let hits = table
1823 .retrieve(&Retriever::Ann {
1824 column_id: 2,
1825 query: vec![1.0, 0.0, 0.0, 0.0],
1826 k: 1,
1827 })
1828 .unwrap();
1829 assert_eq!(hits.len(), 1);
1830 assert!(matches!(
1831 hits[0].score,
1832 RetrieverScore::AnnCosineDistance(d) if d.abs() < 1e-4
1833 ));
1834 }
1835}