From f683003e8ee7d285e86c855c0ff5d5e3807929a7 Mon Sep 17 00:00:00 2001
From: Loakesh Indiran <66478092+Loaki07@users.noreply.github.com>
Date: Thu, 22 May 2025 15:54:42 +0530
Subject: [PATCH] fix: add lock to download file (#6875) (#6880)

fix #6875
---
 src/job/file_downloader.rs | 54 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/src/job/file_downloader.rs b/src/job/file_downloader.rs
index 08cb8a3584..bd3f730145 100644
--- a/src/job/file_downloader.rs
+++ b/src/job/file_downloader.rs
@@ -26,6 +26,28 @@ use tokio::sync::{
 /// (trace_id, account, file, size, cache_type)
 type FileInfo = (String, String, String, usize, file_data::CacheType);
 
+mod processing_files {
+    use hashbrown::HashSet;
+    use parking_lot::RwLock;
+
+    use super::*;
+
+    static PROCESSING_FILES: Lazy<RwLock<HashSet<String>>> =
+        Lazy::new(|| RwLock::new(HashSet::new()));
+
+    pub fn is_processing(file_name: &str) -> bool {
+        PROCESSING_FILES.read().contains(file_name)
+    }
+
+    pub fn add(file_name: &str) {
+        PROCESSING_FILES.write().insert(file_name.to_string());
+    }
+
+    pub fn remove(file_name: &str) {
+        PROCESSING_FILES.write().remove(file_name);
+    }
+}
+
 struct DownloadQueue {
     sender: Sender<FileInfo>,
     receiver: Arc<Mutex<Receiver<FileInfo>>>,
@@ -87,6 +109,19 @@ pub async fn run() -> Result<(), anyhow::Error> {
                         break;
                     }
                     Some((trace_id, account, file, file_size, cache)) => {
+                        // check if the file is already being downloaded
+                        if processing_files::is_processing(&file) {
+                            log::warn!(
+                                "[trace_id {trace_id}] [thread {thread}] search->storage: file {} is already being downloaded, will skip it",
+                                file
+                            );
+                            continue;
+                        }
+
+                        // add the file to processing set
+                        processing_files::add(&file);
+
+                        // download the file
                         match download_file(thread, &trace_id, &account, &file, file_size, cache)
                             .await
                         {
@@ -110,6 +145,9 @@ pub async fn run() -> Result<(), anyhow::Error> {
                             }
                         }
 
+                        // remove the file from processing set
+                        processing_files::remove(&file);
+
                         // update metrics
                         metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE
                             .with_label_values(&[])
@@ -162,6 +200,19 @@ pub async fn run() -> Result<(), anyhow::Error> {
                         let file_info = PRIORITY_FILE_DOWNLOAD_CHANNEL.pop().await;
                         match file_info {
                             Some((trace_id, account, file, file_size, cache)) => {
+                                 // check if the file is already being downloaded
+                                if processing_files::is_processing(&file) {
+                                    log::warn!(
+                                        "[trace_id {trace_id}] [thread {thread}] search->storage: file {} is already being downloaded, will skip it",
+                                        file
+                                    );
+                                    return;
+                                }
+
+                                // add the file to processing set
+                                processing_files::add(&file);
+
+                                // download the file
                                 match download_file(thread, &trace_id, &account, &file, file_size, cache).await {
                                     Ok(data_len) => {
                                         if data_len > 0 && data_len != file_size {
@@ -183,6 +234,9 @@ pub async fn run() -> Result<(), anyhow::Error> {
                                     }
                                 }
 
+                                // remove the file from processing set
+                                processing_files::remove(&file);
+
                                 // update metrics
                                 metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE
                                     .with_label_values(&[])