diff --git a/lib/collection/src/operations/shared_storage_config.rs b/lib/collection/src/operations/shared_storage_config.rs index 600a57f030..8fbb461445 100644 --- a/lib/collection/src/operations/shared_storage_config.rs +++ b/lib/collection/src/operations/shared_storage_config.rs @@ -14,6 +14,10 @@ use crate::shards::transfer::ShardTransferMethod; const DEFAULT_SEARCH_TIMEOUT: Duration = Duration::from_secs(60); const DEFAULT_UPDATE_QUEUE_SIZE: usize = 100; const DEFAULT_UPDATE_QUEUE_SIZE_LISTENER: usize = 10_000; +/// Maximum number of operations which are stored in RAM in update worker queue. +/// If there are more pending operations, operation data +/// will be read from WAL when processing the operation. +pub const DEFAULT_UPDATE_QUEUE_RAM_BUFFER: usize = 500; pub const DEFAULT_IO_SHARD_TRANSFER_LIMIT: Option = Some(1); pub const DEFAULT_SNAPSHOTS_PATH: &str = "./snapshots"; diff --git a/lib/collection/src/shards/local_shard/shard_ops.rs b/lib/collection/src/shards/local_shard/shard_ops.rs index 6155e7c52c..1ff96df3d3 100644 --- a/lib/collection/src/shards/local_shard/shard_ops.rs +++ b/lib/collection/src/shards/local_shard/shard_ops.rs @@ -20,6 +20,7 @@ use tokio::time::error::Elapsed; use crate::collection_manager::segments_searcher::SegmentsSearcher; use crate::operations::OperationWithClockTag; use crate::operations::generalizer::Generalizer; +use crate::operations::shared_storage_config::DEFAULT_UPDATE_QUEUE_RAM_BUFFER; use crate::operations::types::{ CollectionError, CollectionInfo, CollectionResult, CountResult, PointRequestInternal, UpdateResult, UpdateStatus, @@ -70,6 +71,13 @@ impl ShardOperation for LocalShard { let operation_id = { let update_sender = self.update_sender.load(); + // Estimate pending operations count in the channel. + // `Sender::capacity` is returns available slots in the channel regarding tokio docs. + // To calculate pending operations we need to subtract it from the max capacity, + // which is the total capacity defined while creating the channel. + let pending_operations_count = update_sender + .max_capacity() + .saturating_sub(update_sender.capacity()); let channel_permit = update_sender.reserve().await?; // It is *critical* to hold `_wal_lock` while sending operation to the update handler! @@ -90,9 +98,14 @@ impl ShardOperation for LocalShard { Err(err) => return Err(err.into()), }; + // If there are too many pending operations, don't keep operation data in RAM. + // Instead, read operation data from the WAL when processing the operation. + let keep_operation_in_ram = pending_operations_count < DEFAULT_UPDATE_QUEUE_RAM_BUFFER; + let operation = keep_operation_in_ram.then_some(Box::new(operation.operation)); + channel_permit.send(UpdateSignal::Operation(OperationData { op_num: operation_id, - operation: operation.operation, + operation, sender: callback_sender, hw_measurements: hw_measurement_acc.clone(), })); diff --git a/lib/collection/src/update_handler.rs b/lib/collection/src/update_handler.rs index 79723b2ba6..3ed2c8aaea 100644 --- a/lib/collection/src/update_handler.rs +++ b/lib/collection/src/update_handler.rs @@ -7,6 +7,7 @@ use common::counter::hardware_accumulator::HwMeasurementAcc; use common::save_on_disk::SaveOnDisk; use parking_lot::Mutex; use segment::types::SeqNumberType; +use shard::operations::CollectionUpdateOperations; use tokio::runtime::Handle; use tokio::sync::mpsc::{self, Receiver}; use tokio::sync::{Mutex as TokioMutex, oneshot, watch}; @@ -19,7 +20,6 @@ use crate::collection_manager::optimizers::segment_optimizer::{ SegmentOptimizer, plan_optimizations, }; use crate::common::stoppable_task::StoppableTaskHandle; -use crate::operations::CollectionUpdateOperations; use crate::operations::shared_storage_config::SharedStorageConfig; use crate::operations::types::CollectionResult; use crate::shards::CollectionId; @@ -35,8 +35,8 @@ pub type Optimizer = dyn SegmentOptimizer + Sync + Send; pub struct OperationData { /// Sequential number of the operation pub op_num: SeqNumberType, - /// Operation - pub operation: CollectionUpdateOperations, + /// Operation. If None, then the operation data is read from WAL + pub operation: Option>, /// Callback notification channel pub sender: Option>>, /// Hardware measurement for the operation diff --git a/lib/collection/src/update_workers/update_worker.rs b/lib/collection/src/update_workers/update_worker.rs index 08ebddb5da..571cde724f 100644 --- a/lib/collection/src/update_workers/update_worker.rs +++ b/lib/collection/src/update_workers/update_worker.rs @@ -50,6 +50,33 @@ impl UpdateWorkers { let update_operation_lock_clone = update_operation_lock.clone(); let update_tracker_clone = update_tracker.clone(); + let operation = if let Some(operation) = operation { + *operation + } else { + let record = wal.lock().await.read_single_record(op_num); + match record { + Ok(Some(op)) => op.operation, + Ok(None) => { + if let Some(feedback) = sender { + feedback.send(Err(CollectionError::service_error( + format!("Operation {op_num} not found in WAL"), + ))).unwrap_or_else(|_| { + log::debug!("Can't report operation {op_num} result. Assume already not required"); + }); + } + continue; + } + Err(err) => { + if let Some(feedback) = sender { + feedback.send(Err(CollectionError::from(err))).unwrap_or_else(|_| { + log::debug!("Can't report operation {op_num} result. Assume already not required"); + }); + } + continue; + } + } + }; + let operation_result = Self::wait_for_optimization( prevent_unoptimized_threshold_kb, &segments_clone, diff --git a/lib/collection/src/wal_delta.rs b/lib/collection/src/wal_delta.rs index 6a4c135c9d..4300c47b14 100644 --- a/lib/collection/src/wal_delta.rs +++ b/lib/collection/src/wal_delta.rs @@ -69,6 +69,14 @@ impl RecoverableWal { wal_lock.write(operation).map(|op_num| (op_num, wal_lock)) } + /// Read a single record from the WAL by its operation number. + pub async fn read_single_record( + &self, + op_num: u64, + ) -> shard::wal::Result> { + self.wal.lock().await.read_single_record(op_num) + } + /// Take clocks snapshot because we deactivated our replica /// /// Does nothing if a snapshot already existed. Returns `true` if a snapshot was taken. diff --git a/lib/shard/src/wal.rs b/lib/shard/src/wal.rs index 9fd29fdcf5..fd3d63136e 100644 --- a/lib/shard/src/wal.rs +++ b/lib/shard/src/wal.rs @@ -79,6 +79,21 @@ impl SerdeWal { } } + pub fn read_single_record(&self, idx: u64) -> Result> { + if let Some(entry) = self.wal.entry(idx) { + let record: R = serde_cbor::from_slice(&entry) + .or_else(|_err| rmp_serde::from_slice(&entry)) + .map_err(|err| { + WalError::WriteWalError(format!( + "Can't deserialize entry, probably corrupted WAL or version mismatch: {err:?}" + )) + })?; + Ok(Some(record)) + } else { + Ok(None) + } + } + pub fn read(&self, from: u64) -> impl DoubleEndedIterator + '_ { // We have to explicitly do `from..self.first_index() + self.len(false)`, instead of more // concise `from..=self.last_index()`, because if the WAL is empty, `Wal::last_index` @@ -265,7 +280,7 @@ mod tests { use super::*; - #[derive(Debug, Deserialize, Serialize)] + #[derive(Debug, Deserialize, Serialize, PartialEq)] #[serde(rename_all = "snake_case")] #[serde(untagged)] enum TestRecord { @@ -273,13 +288,13 @@ mod tests { Struct2(TestInternalStruct2), } - #[derive(Debug, Deserialize, Serialize)] + #[derive(Debug, Deserialize, Serialize, PartialEq)] #[serde(rename_all = "snake_case")] struct TestInternalStruct1 { data: usize, } - #[derive(Debug, Deserialize, Serialize)] + #[derive(Debug, Deserialize, Serialize, PartialEq)] #[serde(rename_all = "snake_case")] struct TestInternalStruct2 { a: i32, @@ -325,6 +340,16 @@ mod tests { assert_eq!(idx1, 0); assert_eq!(idx2, 1); + assert_eq!( + serde_wal.read_single_record(idx1).unwrap().unwrap(), + record1 + ); + assert_eq!( + serde_wal.read_single_record(idx2).unwrap().unwrap(), + record2 + ); + assert_eq!(serde_wal.read_single_record(100).unwrap(), None); + match record1 { TestRecord::Struct1(x) => assert_eq!(x.data, 10), TestRecord::Struct2(_) => panic!("Wrong structure"),