Update queue dont keep ops in ram (#7951)

* update queue dont keep ops in ram showcase

* always load operation from WAL

* revert operations buffering

cound only pending in update worker operations

fix typo

use channel size instead of wal index

* remove result expect

* decrease buffering const

* review remarks

* are you happy codespell

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<usize> = Some(1);
 pub const DEFAULT_SNAPSHOTS_PATH: &str = "./snapshots";
 

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(),
             }));

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<Box<CollectionUpdateOperations>>,
     /// Callback notification channel
     pub sender: Option<oneshot::Sender<CollectionResult<usize>>>,
     /// Hardware measurement for the operation

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,

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<Option<OperationWithClockTag>> {
+        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.

lib/shard/src/wal.rs

@@ -79,6 +79,21 @@ impl<R: DeserializeOwned + Serialize> SerdeWal<R> {
         }
     }
 
+    pub fn read_single_record(&self, idx: u64) -> Result<Option<R>> {
+        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<Item = (u64, R)> + '_ {
         // 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"),