use super::*;
use crate::query::DocBitset;
impl<D: DirectoryWriter + 'static> SegmentManager<D> {
pub(crate) async fn commit_with_deletes(
self: &Arc<Self>,
new_segments: &[(String, u32)],
keys: Vec<String>,
staged_deletions: Vec<(String, Arc<crate::index::staged_row::StagedSegment>)>,
) -> Result<()> {
if keys.is_empty() && staged_deletions.is_empty() {
return self.commit(new_segments).await;
}
let field = self
.schema
.primary_field()
.ok_or_else(|| Error::Schema("row deletion requires a primary key".into()))?;
for (id, count) in new_segments {
self.validate_completed_segment(id, *count).await?;
}
let mut st = Arc::clone(&self.state).lock_owned().await;
let manager = Arc::clone(self);
let new_segments = new_segments.to_vec();
self.run_lifecycle_transaction(async move {
let keys: Arc<HashSet<String>> = Arc::new(keys.into_iter().collect());
let mut next = st.metadata.clone();
let mut retired = Vec::new();
let mut outputs = Vec::new();
for (segment_id, info) in st
.metadata
.segment_metas
.iter()
.filter(|_| !keys.is_empty())
{
let sid = SegmentId::from_hex(segment_id)
.ok_or_else(|| Error::Corruption("invalid deletion source".into()))?;
let mut fields = crate::segment::reader::loader::load_fast_fields_file(
manager.directory.as_ref(),
&SegmentFiles::new(sid.0),
&manager.schema,
)
.await?;
let ff = fields
.remove(&field.0)
.ok_or_else(|| Error::Corruption("primary-key column is missing".into()))?;
let cancellation = manager.active_operations.cancellation_flag();
let lookup_cancellation = Arc::clone(&cancellation);
let keys = Arc::clone(&keys);
let num_docs = info.num_docs;
let (ff, ordinals) = tokio::task::spawn_blocking(move || {
let ordinals = crate::segment::deletion::target_ordinals(
&ff,
num_docs,
keys.iter().map(String::as_str),
|| {
if lookup_cancellation.load(Ordering::Acquire) {
Err(Error::IndexClosed)
} else {
Ok(())
}
},
)?;
Ok::<_, Error>((ff, ordinals))
})
.await
.map_err(|error| {
Error::Internal(format!("deletion key lookup failed: {error}"))
})??;
if ordinals.is_empty() {
continue;
}
let mut alive = match &info.deletions {
Some(meta) => {
(*meta.load(manager.directory.as_ref(), info.num_docs).await?).clone()
}
None => DocBitset::all(info.num_docs),
};
let (alive, changed) = tokio::task::spawn_blocking(move || {
let changed = crate::segment::deletion::clear_target_rows(
&ff,
&ordinals,
&mut alive,
|| {
if cancellation.load(Ordering::Acquire) {
Err(Error::IndexClosed)
} else {
Ok(())
}
},
)?;
Ok::<_, Error>((alive, changed))
})
.await
.map_err(|error| Error::Internal(format!("deletion worker failed: {error}")))??;
if !changed {
continue;
}
let id = SegmentId::new();
let claim = manager.protect_new_segment(id.to_hex())?;
let cleanup = manager.output_cleanup_guard(id);
outputs.push((claim, cleanup));
let deletion = crate::segment::deletion::write(
manager.directory.as_ref(),
id,
info.num_docs,
&alive,
)
.await?;
if let Some(old) = &info.deletions {
retired.push(old.id.clone());
}
next.segment_metas.get_mut(segment_id).unwrap().deletions = Some(deletion);
}
for (id, count) in &new_segments {
if !next.has_segment(id) {
next.add_segment(id.clone(), *count);
}
}
for (segment_id, staged) in &staged_deletions {
if st.metadata.has_segment(segment_id) {
return Err(Error::Corruption(
"staged visibility must refer to a new segment".into(),
));
}
let info = next
.segment_metas
.get_mut(segment_id)
.ok_or_else(|| Error::Corruption("staged segment is missing".into()))?;
let alive = staged.live_rows(info.num_docs).ok_or_else(|| {
Error::Internal("staged visibility has no cancelled rows".into())
})?;
let id = SegmentId::new();
let claim = manager.protect_new_segment(id.to_hex())?;
let cleanup = manager.output_cleanup_guard(id);
outputs.push((claim, cleanup));
info.deletions = Some(
crate::segment::deletion::write(
manager.directory.as_ref(),
id,
info.num_docs,
&alive,
)
.await?,
);
}
next.publication_generation = next
.publication_generation
.checked_add(1)
.ok_or_else(|| Error::Corruption("publication generation overflow".into()))?;
next.save(manager.directory.as_ref()).await?;
for id in next.owned_ids() {
manager.tracker.register(&id);
}
let previous = manager.published_generation();
manager
.published_generation
.store(Arc::new(PublishedIndexGeneration {
publication_id: next.publication_generation,
schema: Arc::clone(&previous.schema),
trained_vectors: previous.trained_vectors.clone(),
}));
retired.retain(|id| !next.owns_id(id));
st.metadata = next;
for (_, cleanup) in &mut outputs {
cleanup.disarm();
}
let ready = manager.tracker.mark_for_deletion(&retired);
drop(st);
(manager.delete_fn)(ready);
Ok(())
})
.await
}
}
impl<D: DirectoryWriter + 'static> SegmentManager<D> {
pub async fn compact_segment(
self: &Arc<Self>,
segment_id: &str,
memory_budget: usize,
) -> Result<bool> {
self.compact_segment_admitted(segment_id, memory_budget, None)
.await
}
pub async fn compaction_candidates(&self, threshold: f64) -> Result<Vec<(String, u32, u32)>> {
validate_threshold(threshold)?;
if threshold == 0.0 || self.force_merge_active.load(Ordering::Acquire) > 0 {
return Ok(Vec::new());
}
let active = self.active_operations.snapshot();
let paused = self.paused_reorder_segments();
let quarantined = self.quarantined_segments.lock().clone();
let st = self.state.lock().await;
let mut candidates: Vec<_> = st
.metadata
.segment_metas
.iter()
.filter(|(id, meta)| {
meta.num_deleted_docs() > 0
&& meta.deleted_ratio() >= threshold
&& !active.contains(*id)
&& !paused.contains(*id)
&& !quarantined.contains(*id)
})
.map(|(id, meta)| (id.clone(), meta.num_docs, meta.num_deleted_docs()))
.collect();
candidates.sort_unstable_by(|a, b| {
(f64::from(b.2) / f64::from(b.1))
.total_cmp(&(f64::from(a.2) / f64::from(a.1)))
.then_with(|| a.1.cmp(&b.1))
.then_with(|| a.0.cmp(&b.0))
});
Ok(candidates)
}
pub async fn compact_segment_if_eligible(
self: &Arc<Self>,
segment_id: &str,
threshold: f64,
memory_budget: usize,
) -> Result<bool> {
validate_threshold(threshold)?;
if threshold == 0.0 {
return Ok(false);
}
let result = self
.compact_segment_admitted(segment_id, memory_budget, Some(threshold))
.await;
match &result {
Ok(true) => self.clear_reorder_retry(segment_id),
Err(error) if !matches!(error, Error::IndexClosed) => {
self.pause_reorder_retries(segment_id, error).await
}
_ => {}
}
result
}
async fn compact_segment_admitted(
self: &Arc<Self>,
segment_id: &str,
memory_budget: usize,
threshold: Option<f64>,
) -> Result<bool> {
if memory_budget < 1024 * 1024 {
return Err(Error::Schema(
"compaction memory budget must be at least 1 MiB".into(),
));
}
let sid = SegmentId::from_hex(segment_id)
.ok_or_else(|| Error::Document("invalid compaction segment ID".into()))?;
if !self.active_operations.is_accepting() {
return Err(Error::IndexClosed);
}
if self.quarantined_segments.lock().contains(segment_id) {
return Err(Error::Corruption("compaction source is quarantined".into()));
}
if threshold.is_some()
&& (self.force_merge_active.load(Ordering::Acquire) > 0
|| self.paused_reorder_segments().contains(segment_id))
{
return Ok(false);
}
let (global_permit, permit, reorder_permit) = if threshold.is_some() {
let Ok(global) = Arc::clone(&self.global_merge_permits).try_acquire_owned() else {
return Ok(false);
};
let Ok(local) = Arc::clone(&self.merge_permits).try_acquire_owned() else {
return Ok(false);
};
let Ok(reorder) = self.reorder_permits.try_acquire_optimizer() else {
return Ok(false);
};
(global, local, reorder)
} else {
let (global, local) = self.acquire_maintenance_capacity().await?;
let reorder = tokio::select! {
biased;
() = self.active_operations.wait_for_shutdown() => return Err(Error::IndexClosed),
permit = self.reorder_permits.acquire(ReorderPriority::Optimizer) => permit.map_err(|_| Error::IndexClosed)?,
};
(global, local, reorder)
};
let manager = Arc::clone(self);
let source_id = segment_id.to_owned();
self.run_lifecycle_transaction(async move {
let _permit = permit;
let _global_permit = global_permit;
let _reorder_permit = reorder_permit;
let start = std::time::Instant::now();
let output = SegmentId::new();
let (claim, snapshot, visibility, schema) = {
let st = manager.state.lock().await;
let Some(info) = st.metadata.segment_metas.get(&source_id) else {
return Ok(false);
};
if let Some(threshold) = threshold
&& (manager.force_merge_active.load(Ordering::Acquire) > 0 || info.deleted_ratio() < threshold) {
return Ok(false);
}
let Some(visibility) = info.deletions.clone() else {
return Ok(false);
};
let Some(claim) = manager.active_operations.try_register(vec![source_id.clone(), output.to_hex()]) else {
return if threshold.is_some() { Ok(false) } else {
Err(Error::Internal("compaction source is busy; retry after maintenance finishes".into()))
};
};
let acquired = manager
.tracker
.acquire(&[source_id.clone(), visibility.id.clone()]);
let snapshot = SegmentSnapshot::with_delete_fn(
Arc::clone(&manager.tracker),
acquired,
Arc::clone(&manager.delete_fn),
);
(
claim,
snapshot,
visibility,
manager.published_generation().schema.clone(),
)
};
let _claim = claim;
let _snapshot = snapshot;
let mut cleanup = manager.output_cleanup_guard(output);
let mut reader = SegmentReader::open_with_term_cache_budget(
manager.directory.as_ref(),
sid,
schema.clone(),
manager.term_cache_blocks,
manager.term_cache_budget_bytes,
)
.await?;
reader
.load_deletions(manager.directory.as_ref(), visibility.clone())
.await?;
let physical = reader.num_docs();
let meta = manager
.build_compacted(reader, output, memory_budget)
.await?;
let removed = visibility.num_deleted;
if meta.num_docs.checked_add(removed) != Some(physical) {
return Err(Error::Corruption("compaction row count mismatch".into()));
}
let expected = HashMap::from([(source_id.clone(), Some(visibility))]);
manager
.replace_segments(
std::slice::from_ref(&source_id),
output.to_hex(),
meta.num_docs,
ReplacementLayout::Compacted,
Some(&expected),
)
.await?;
cleanup.disarm();
log::info!("[compaction] index={} source={} physical_rows={} removed_rows={} removed_ratio={:.6} live_rows={} elapsed_secs={:.3}",
manager.schema.index_label(), source_id, physical, removed,
f64::from(removed) / f64::from(physical), meta.num_docs, start.elapsed().as_secs_f64());
Ok(true)
})
.await
}
pub(super) async fn build_compacted(
&self,
reader: SegmentReader,
output: SegmentId,
memory_budget: usize,
) -> Result<crate::segment::SegmentMeta> {
let directory = Arc::clone(&self.directory);
let merger = crate::segment::SegmentMerger::new(self.published_generation().schema.clone())
.with_posting_config(self.optimization, self.posting_codec)
.with_term_dict_block_size(self.term_dict_block_size)
.with_cancellation(self.active_operations.cancellation_flag())
.with_background_pool(Some(self.background_cpu_pool()));
let runtime = tokio::runtime::Handle::current();
let (meta, _) = tokio::task::spawn_blocking(move || {
runtime.block_on(merger.compact(directory.as_ref(), &reader, output, memory_budget))
})
.await
.map_err(|error| Error::Internal(format!("compaction worker failed: {error}")))??;
Ok(meta)
}
}
fn validate_threshold(threshold: f64) -> Result<()> {
if !threshold.is_finite() || !(0.0..=1.0).contains(&threshold) {
return Err(Error::Schema("compaction deleted ratio must be finite and in 0..=1 (0 disables automatic compaction)".into()));
}
Ok(())
}