chore(cubestore): Bulk in memory chunks drop (#6785)

Author: waralexrom

rust/cubestore/cubestore/src/cluster/message.rs

@@ -43,7 +43,7 @@ pub enum NetworkMessage {
     },
     FreeMemoryChunkResult(Result<(), CubeError>),
 
-    FreeDeletedMemoryChunks,
+    FreeDeletedMemoryChunks(Vec<u64>),
     FreeDeletedMemoryChunksResult(Result<(), CubeError>),
 
     MetaStoreCall(MetaStoreRpcMethodCall),

rust/cubestore/cubestore/src/cluster/mod.rs

@@ -132,7 +132,11 @@ pub trait Cluster: DIService + Send + Sync {
     ) -> Result<(), CubeError>;
 
     async fn free_memory_chunk(&self, node_name: &str, chunk_id: u64) -> Result<(), CubeError>;
-    async fn free_deleted_memory_chunks(&self, node_name: &str) -> Result<(), CubeError>;
+    async fn free_deleted_memory_chunks(
+        &self,
+        node_name: &str,
+        chunk_ids: Vec<u64>,
+    ) -> Result<(), CubeError>;
 
     fn job_result_listener(&self) -> JobResultListener;
 
@@ -446,9 +450,16 @@ impl Cluster for ClusterImpl {
         }
     }
 
-    async fn free_deleted_memory_chunks(&self, node_name: &str) -> Result<(), CubeError> {
+    async fn free_deleted_memory_chunks(
+        &self,
+        node_name: &str,
+        chunk_ids: Vec<u64>,
+    ) -> Result<(), CubeError> {
         let response = self
-            .send_or_process_locally(node_name, NetworkMessage::FreeDeletedMemoryChunks)
+            .send_or_process_locally(
+                node_name,
+                NetworkMessage::FreeDeletedMemoryChunks(chunk_ids),
+            )
             .await?;
         match response {
             NetworkMessage::FreeDeletedMemoryChunksResult(r) => r,
@@ -594,14 +605,14 @@ impl Cluster for ClusterImpl {
                 let res = chunk_store.free_memory_chunk(chunk_id).await;
                 NetworkMessage::FreeMemoryChunkResult(res)
             }
-            NetworkMessage::FreeDeletedMemoryChunks => {
+            NetworkMessage::FreeDeletedMemoryChunks(ids) => {
                 let chunk_store = self
                     .injector
                     .upgrade()
                     .unwrap()
                     .get_service_typed::<dyn ChunkDataStore>()
                     .await;
-                let res = chunk_store.free_deleted_memory_chunks().await;
+                let res = chunk_store.free_deleted_memory_chunks(ids).await;
                 NetworkMessage::FreeDeletedMemoryChunksResult(res)
             }
             NetworkMessage::FreeMemoryChunkResult(_) => {

rust/cubestore/cubestore/src/scheduler/mod.rs

@@ -44,6 +44,7 @@ pub struct SchedulerImpl {
     reconcile_loop: WorkerLoop,
     chunk_processing_loop: WorkerLoop,
     chunk_events_queue: Mutex<Vec<(SystemTime, u64)>>,
+    in_memory_chunks_to_delete: Mutex<Vec<(String, u64)>>, //(node, chunk_is)
     node_last_actions: Mutex<HashMap<String, LastNodeActionTimes>>,
 }
 
@@ -97,6 +98,7 @@ impl SchedulerImpl {
             gc_queue,
             reconcile_loop: WorkerLoop::new("Reconcile"),
             chunk_events_queue: Mutex::new(Vec::with_capacity(1000)),
+            in_memory_chunks_to_delete: Mutex::new(Vec::with_capacity(1000)),
             node_last_actions: Mutex::new(workers),
             chunk_processing_loop: WorkerLoop::new("ChunkProcessing"),
         }
@@ -756,9 +758,8 @@ impl SchedulerImpl {
                     .get_partition(chunk.get_row().get_partition_id())
                     .await?;
                 let node_name = self.cluster.node_name_by_partition(&partition);
-                self.cluster
-                    .free_memory_chunk(&node_name, chunk.get_id())
-                    .await?;
+                let mut to_delete = self.in_memory_chunks_to_delete.lock().await;
+                to_delete.push((node_name, chunk.get_id()))
             } else if chunk.get_row().uploaded() {
                 let file_name =
                     ChunkStore::chunk_remote_path(chunk.get_id(), chunk.get_row().suffix());
@@ -899,6 +900,39 @@ impl SchedulerImpl {
             self.process_active_chunks(active_chunks).await?;
             self.process_inactive_chunks(inactive_chunks).await?;
         }
+        {
+            let chunks_to_delete = {
+                let mut chunks = self.in_memory_chunks_to_delete.lock().await;
+                if chunks.is_empty() {
+                    Vec::new()
+                } else {
+                    let mut result = Vec::new();
+                    std::mem::swap(&mut *chunks, &mut result);
+                    result
+                }
+            };
+            if !chunks_to_delete.is_empty() {
+                let chunks_to_delete = chunks_to_delete.into_iter().into_group_map();
+                for (node, ids) in chunks_to_delete {
+                    if !ids.is_empty() {
+                        if let Err(e) = self
+                            .cluster
+                            .free_deleted_memory_chunks(&node, ids.clone())
+                            .await
+                        {
+                            log::error!(
+                                "Error while trying release in memory chunks in node {}: {}",
+                                node,
+                                e
+                            );
+
+                            let mut chunks = self.in_memory_chunks_to_delete.lock().await;
+                            chunks.extend(ids.iter().map(|v| (node.clone(), *v)));
+                        }
+                    }
+                }
+            }
+        }
 
         Ok(())
     }

rust/cubestore/cubestore/src/store/mod.rs

@@ -252,7 +252,7 @@ pub trait ChunkDataStore: DIService + Send + Sync {
     ) -> Result<(Vec<ArrayRef>, Vec<u64>), CubeError>;
     async fn add_memory_chunk(&self, chunk_id: u64, batch: RecordBatch) -> Result<(), CubeError>;
     async fn free_memory_chunk(&self, chunk_id: u64) -> Result<(), CubeError>;
-    async fn free_deleted_memory_chunks(&self) -> Result<(), CubeError>;
+    async fn free_deleted_memory_chunks(&self, chunk_ids: Vec<u64>) -> Result<(), CubeError>;
     async fn add_persistent_chunk(
         &self,
         index: IdRow<Index>,
@@ -705,19 +705,13 @@ impl ChunkDataStore for ChunkStore {
         memory_chunks.remove(&chunk_id);
         Ok(())
     }
-    #[tracing::instrument(level = "trace", skip(self))]
-    async fn free_deleted_memory_chunks(&self) -> Result<(), CubeError> {
-        let existing_chunk_ids = self
-            .meta_store
-            .get_all_node_in_memory_chunks(self.cluster.server_name().to_string())
-            .await?
-            .into_iter()
-            .map(|c| c.get_id())
-            .collect::<HashSet<_>>();
 
+    #[tracing::instrument(level = "trace", skip(self))]
+    async fn free_deleted_memory_chunks(&self, chunk_ids: Vec<u64>) -> Result<(), CubeError> {
+        let ids_set = chunk_ids.into_iter().collect::<HashSet<_>>();
         {
             let mut memory_chunks = self.memory_chunks.write().await;
-            memory_chunks.retain(|id, _| existing_chunk_ids.contains(id));
+            memory_chunks.retain(|id, _| !ids_set.contains(id));
         }
 
         self.report_in_memory_metrics().await?;