From e5dd8488f13df400a6b066dc4bbf6fa3cec1a4ec Mon Sep 17 00:00:00 2001 From: xzfc <5121426+xzfc@users.noreply.github.com> Date: Tue, 20 Jan 2026 10:07:35 +0000 Subject: [PATCH] GET /optimizations: re-plan on each request (#7945) --- .../src/collection/collection_ops.rs | 10 ++--- .../src/collection_manager/optimizers/mod.rs | 3 -- .../optimizers/segment_optimizer.rs | 41 +++++++++++------- lib/collection/src/shards/local_shard/mod.rs | 26 ++++++++++- lib/collection/src/shards/replica_set/mod.rs | 11 ++++- lib/collection/src/shards/shard.rs | 14 +++++- lib/collection/src/update_handler.rs | 15 ++----- .../src/update_workers/optimization_worker.rs | 43 ++----------------- 8 files changed, 83 insertions(+), 80 deletions(-) diff --git a/lib/collection/src/collection/collection_ops.rs b/lib/collection/src/collection/collection_ops.rs index cbbf8aa14c..407fc9120d 100644 --- a/lib/collection/src/collection/collection_ops.rs +++ b/lib/collection/src/collection/collection_ops.rs @@ -439,16 +439,14 @@ impl Collection { let Some(log) = replica_set.optimizers_log().await else { continue; }; - - let log = log.lock(); - let IndexingProgressViews { ongoing, completed } = log.progress_views(); - pending.merge(&log.pending); - drop(log); - + let IndexingProgressViews { ongoing, completed } = log.lock().progress_views(); all_ongoing.extend(ongoing); if let Some(all_completed) = all_completed.as_mut() { all_completed.extend(completed); } + if let Some(shard_pending) = replica_set.pending_optimizations().await { + pending.merge(&shard_pending); + } } // Sort - see `OptimizationsResponse` doc all_ongoing.sort_by_key(|v| Reverse(v.started_at())); diff --git a/lib/collection/src/collection_manager/optimizers/mod.rs b/lib/collection/src/collection_manager/optimizers/mod.rs index 84be1079e3..565c55d7a8 100644 --- a/lib/collection/src/collection_manager/optimizers/mod.rs +++ b/lib/collection/src/collection_manager/optimizers/mod.rs @@ -9,8 +9,6 @@ use segment::common::anonymize::Anonymize; use serde::{Deserialize, Serialize}; use super::holders::segment_holder::SegmentId; -use crate::operations::types::PendingOptimizations; - pub mod config_mismatch_optimizer; pub mod indexing_optimizer; pub mod merge_optimizer; @@ -26,7 +24,6 @@ const KEEP_LAST_TRACKERS: usize = 16; #[derive(Default, Clone, Debug)] pub struct TrackerLog { descriptions: VecDeque, - pub pending: PendingOptimizations, } #[derive(Clone, Debug, Default)] diff --git a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs index e98747781c..a2304d1a56 100644 --- a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs @@ -979,6 +979,7 @@ pub struct OptimizationPlanner<'a> { running: usize, /// This goes into [`Self::scheduled`]. + /// Should be set before calling [`Self::plan`]. optimizer: Option>, } @@ -999,21 +1000,7 @@ impl<'a> OptimizationPlanner<'a> { &self.remaining } - /// Should set called before calling [`Self::plan`] if you want the - /// optimizer to be attached to [`Self::scheduled`]. - pub fn set_optimizer(&mut self, optimizer: Arc) { - self.optimizer = Some(optimizer); - } - - pub fn scheduled(&self) -> &Vec<(Option>, Vec)> { - &self.scheduled - } - - pub fn into_scheduled(self) -> Vec<(Option>, Vec)> { - self.scheduled - } - - /// Like [`Self::into_scheduled`], but drops the optimizers. + /// Returns [`Self::scheduled`], but without `Option>` part. #[cfg(test)] pub fn into_scheduled_for_test(self) -> Vec> { self.scheduled @@ -1038,3 +1025,27 @@ impl<'a> OptimizationPlanner<'a> { self.scheduled.push((self.optimizer.clone(), segments)); } } + +/// Plans optimizations for the given segments and optimizers. +/// +/// Returns a list of scheduled optimizations, each containing the +/// corresponding optimizer and a batch of segment IDs to be optimized. +pub fn plan_optimizations( + segments: &SegmentHolder, + optimizers: &[Arc], +) -> Vec<(Arc, Vec)> { + let mut planner = OptimizationPlanner::new( + segments.running_optimizations.count(), + segments.iter_original(), + ); + for optimizer in optimizers { + planner.optimizer = Some(Arc::clone(optimizer)); + optimizer.plan_optimizations(&mut planner); + } + planner + .scheduled + .into_iter() + .inspect(|(optimizer, _segments)| debug_assert!(optimizer.is_some())) + .filter_map(|(optimizer, segments)| Some((optimizer?, segments))) + .collect() +} diff --git a/lib/collection/src/shards/local_shard/mod.rs b/lib/collection/src/shards/local_shard/mod.rs index afc189a93b..626274361d 100644 --- a/lib/collection/src/shards/local_shard/mod.rs +++ b/lib/collection/src/shards/local_shard/mod.rs @@ -68,14 +68,15 @@ use crate::collection_manager::holders::segment_holder::{ LockedSegment, LockedSegmentHolder, SegmentHolder, }; use crate::collection_manager::optimizers::TrackerLog; +use crate::collection_manager::optimizers::segment_optimizer::plan_optimizations; use crate::collection_manager::segments_searcher::SegmentsSearcher; use crate::common::file_utils::{move_dir, move_file}; use crate::config::CollectionConfigInternal; use crate::operations::OperationWithClockTag; use crate::operations::shared_storage_config::SharedStorageConfig; use crate::operations::types::{ - CollectionError, CollectionResult, OptimizersStatus, ShardInfoInternal, ShardStatus, - check_sparse_compatible_with_segment_config, + CollectionError, CollectionResult, OptimizersStatus, PendingOptimizations, ShardInfoInternal, + ShardStatus, check_sparse_compatible_with_segment_config, }; use crate::optimizers_builder::{OptimizersConfig, build_optimizers, clear_temp_segments}; use crate::shards::CollectionId; @@ -987,6 +988,27 @@ impl LocalShard { Arc::clone(&self.optimizers_log) } + /// Call [`plan_optimizations`] and return summary. + pub fn pending_optimizations(&self) -> PendingOptimizations { + let segments = self.segments.read(); + let scheduled = plan_optimizations(&segments, &self.optimizers); + let mut pending_segments = 0; + let mut points = 0; + for (_, segment_ids) in scheduled.iter() { + pending_segments += segment_ids.len(); + for &segment_id in segment_ids { + if let Some(LockedSegment::Original(segment)) = segments.get(segment_id) { + points += segment.read().available_point_count(); + } + } + } + PendingOptimizations { + optimizations: scheduled.len(), + segments: pending_segments, + points, + } + } + /// Get the recovery point for the current shard /// /// This is sourced from the last seen clocks from other nodes that we know about. diff --git a/lib/collection/src/shards/replica_set/mod.rs b/lib/collection/src/shards/replica_set/mod.rs index 4e0bd63ebb..bc7d4b1420 100644 --- a/lib/collection/src/shards/replica_set/mod.rs +++ b/lib/collection/src/shards/replica_set/mod.rs @@ -39,7 +39,9 @@ use crate::common::collection_size_stats::CollectionSizeStats; use crate::common::snapshots_manager::SnapshotStorageManager; use crate::config::CollectionConfigInternal; use crate::operations::shared_storage_config::SharedStorageConfig; -use crate::operations::types::{CollectionError, CollectionResult, UpdateResult, UpdateStatus}; +use crate::operations::types::{ + CollectionError, CollectionResult, PendingOptimizations, UpdateResult, UpdateStatus, +}; use crate::operations::{CollectionUpdateOperations, point_ops}; use crate::optimizers_builder::OptimizersConfig; use crate::shards::channel_service::ChannelService; @@ -1345,6 +1347,13 @@ impl ShardReplicaSet { let local = self.local.read().await; local.as_ref().and_then(|shard| shard.optimizers_log()) } + + pub async fn pending_optimizations(&self) -> Option { + let local = self.local.read().await; + local + .as_ref() + .and_then(|shard| shard.pending_optimizations()) + } } /// Represents a change in replica set, due to scaling of `replication_factor` diff --git a/lib/collection/src/shards/shard.rs b/lib/collection/src/shards/shard.rs index ba3e79ef33..5c36e42165 100644 --- a/lib/collection/src/shards/shard.rs +++ b/lib/collection/src/shards/shard.rs @@ -16,7 +16,9 @@ use super::local_shard::clock_map::RecoveryPoint; use super::update_tracker::UpdateTracker; use crate::collection_manager::optimizers::TrackerLog; use crate::operations::operation_effect::{EstimateOperationEffectArea, OperationEffectArea}; -use crate::operations::types::{CollectionError, CollectionResult, OptimizersStatus}; +use crate::operations::types::{ + CollectionError, CollectionResult, OptimizersStatus, PendingOptimizations, +}; use crate::shards::dummy_shard::DummyShard; use crate::shards::forward_proxy_shard::ForwardProxyShard; use crate::shards::local_shard::LocalShard; @@ -248,6 +250,16 @@ impl Shard { Some(optimizers_log) } + pub fn pending_optimizations(&self) -> Option { + Some(match self { + Self::Local(local_shard) => local_shard.pending_optimizations(), + Self::Proxy(proxy_shard) => proxy_shard.wrapped_shard.pending_optimizations(), + Self::ForwardProxy(proxy_shard) => proxy_shard.wrapped_shard.pending_optimizations(), + Self::QueueProxy(proxy_shard) => proxy_shard.wrapped_shard()?.pending_optimizations(), + Self::Dummy(_) => return None, + }) + } + pub async fn shard_recovery_point(&self) -> CollectionResult { match self { Self::Local(local_shard) => Ok(local_shard.recovery_point().await), diff --git a/lib/collection/src/update_handler.rs b/lib/collection/src/update_handler.rs index 0f634fd3e3..23f05e3187 100644 --- a/lib/collection/src/update_handler.rs +++ b/lib/collection/src/update_handler.rs @@ -16,7 +16,7 @@ use crate::collection::payload_index_schema::PayloadIndexSchema; use crate::collection_manager::holders::segment_holder::LockedSegmentHolder; use crate::collection_manager::optimizers::TrackerLog; use crate::collection_manager::optimizers::segment_optimizer::{ - OptimizationPlanner, SegmentOptimizer, + SegmentOptimizer, plan_optimizations, }; use crate::common::stoppable_task::StoppableTaskHandle; use crate::operations::CollectionUpdateOperations; @@ -310,17 +310,8 @@ impl UpdateHandler { pub(crate) fn check_optimizer_conditions(&self) -> (bool, bool) { // Check if Qdrant triggered any optimizations since starting at all let has_triggered_any_optimizers = self.has_triggered_optimizers.load(Ordering::Relaxed); - - let segments = self.segments.read(); - let mut planner = OptimizationPlanner::new( - segments.running_optimizations.count(), - segments.iter_original(), - ); - let has_suboptimal_optimizers = self.optimizers.iter().any(|optimizer| { - optimizer.plan_optimizations(&mut planner); - !planner.scheduled().is_empty() - }); - + let has_suboptimal_optimizers = + !plan_optimizations(&self.segments.read(), &self.optimizers).is_empty(); (has_triggered_any_optimizers, has_suboptimal_optimizers) } diff --git a/lib/collection/src/update_workers/optimization_worker.rs b/lib/collection/src/update_workers/optimization_worker.rs index 2641a61b89..f3afe2376b 100644 --- a/lib/collection/src/update_workers/optimization_worker.rs +++ b/lib/collection/src/update_workers/optimization_worker.rs @@ -10,10 +10,8 @@ use common::panic; use common::save_on_disk::SaveOnDisk; use parking_lot::Mutex; use segment::common::operation_error::{OperationError, OperationResult}; -use segment::entry::SegmentEntry; use segment::index::hnsw_index::num_rayon_threads; use segment::types::QuantizationConfig; -use shard::locked_segment::LockedSegment; use shard::payload_index_schema::PayloadIndexSchema; use shard::segment_holder::LockedSegmentHolder; use tokio::sync::mpsc::{Receiver, Sender}; @@ -25,12 +23,12 @@ use tokio::time::timeout; use crate::collection_manager::collection_updater::CollectionUpdater; use crate::collection_manager::optimizers::segment_optimizer::{ - OptimizationPlanner, OptimizerThresholds, + OptimizerThresholds, plan_optimizations, }; use crate::collection_manager::optimizers::{Tracker, TrackerLog, TrackerStatus}; use crate::common::stoppable_task::{StoppableTaskHandle, spawn_stoppable}; use crate::config::CollectionParams; -use crate::operations::types::{CollectionError, CollectionResult, PendingOptimizations}; +use crate::operations::types::{CollectionError, CollectionResult}; use crate::shards::update_tracker::UpdateTracker; use crate::update_handler::{Optimizer, OptimizerSignal}; use crate::update_workers::UpdateWorkers; @@ -275,43 +273,8 @@ impl UpdateWorkers { let mut handles = vec![]; let is_optimization_failed = Arc::new(AtomicBool::new(false)); - let scheduled; - let mut pending_segments = 0; - let mut pending_points = 0; - { - let segments = segments.read(); - let mut planner = OptimizationPlanner::new( - segments.running_optimizations.count(), - segments.iter_original(), - ); - for optimizer in optimizers.iter() { - planner.set_optimizer(Arc::clone(optimizer)); - optimizer.plan_optimizations(&mut planner); - } - scheduled = planner.into_scheduled(); - - for (_, segment_ids) in &scheduled { - pending_segments += segment_ids.len(); - for &segment_id in segment_ids { - if let Some(LockedSegment::Original(segment)) = segments.get(segment_id) { - pending_points += segment.read().available_point_count(); - } - } - } - } - - optimizers_log.lock().pending = PendingOptimizations { - optimizations: scheduled.len(), - segments: pending_segments, - points: pending_points, - }; - + let scheduled = plan_optimizations(&segments.read(), &optimizers); for (optimizer, segments_to_merge) in scheduled { - let Some(optimizer) = optimizer else { - debug_assert!(false); - continue; - }; - // Return early if we reached the optimization job limit if limit.map(|extra| handles.len() >= extra).unwrap_or(false) { log::trace!("Reached optimization job limit, postponing other optimizations");