feat: add timeout to fusion search (#174)

* feat: add timeout to fusion search * feat: add timeout to fusion search #174
infinilabs · Feb 23, 2025 · 65d1e05 · 65d1e05
1 parent a25bfd7
commit 65d1e05
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 3 deletions.
diff --git a/docs/content.en/docs/release-notes/_index.md b/docs/content.en/docs/release-notes/_index.md
@@ -10,6 +10,8 @@ Information about release notes of Coco Server is provided here.
 ## Latest (In development)
 
 ### Features
+- Add timeout to fusion search #174
+
 ### Breaking changes
 ### Bug fix
 - Fix to access deeplink for linux #148

diff --git a/src-tauri/src/search/mod.rs b/src-tauri/src/search/mod.rs
@@ -6,7 +6,10 @@ use crate::common::traits::SearchError;
 use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use tauri::{AppHandle, Manager, Runtime};
+use tokio::time::{timeout, Duration};
+
 #[tauri::command]
 pub async fn query_coco_fusion<R: Runtime>(
     app_handle: AppHandle<R>,
@@ -22,6 +25,10 @@ pub async fn query_coco_fusion<R: Runtime>(
 
     let sources_list = sources_future.await;
 
+    // Time limit for each query
+    let timeout_duration = Duration::from_millis(100);
+
+    // Push all queries into futures
     for query_source in sources_list {
         let query_source_type = query_source.get_type().clone();
         sources.insert(query_source_type.id.clone(), query_source_type);
@@ -30,7 +37,11 @@ pub async fn query_coco_fusion<R: Runtime>(
         let query_source_clone = query_source.clone(); // Clone Arc to avoid ownership issues
 
         futures.push(tokio::spawn(async move {
-            query_source_clone.search(query).await
+            // Timeout each query execution
+            timeout(timeout_duration, async {
+                query_source_clone.search(query).await
+            })
+                .await
         }));
     }
 
@@ -41,7 +52,7 @@ pub async fn query_coco_fusion<R: Runtime>(
 
     while let Some(result) = futures.next().await {
         match result {
-            Ok(Ok(response)) => {
+            Ok(Ok(Ok(response))) => {
                 total_hits += response.total_hits;
                 let source_id = response.source.id.clone();
 
@@ -72,6 +83,19 @@ pub async fn query_coco_fusion<R: Runtime>(
                     reason: None,
                 });
             }
+            // Timeout reached, skip this request
+            Ok(_) => {
+                failed_requests.push(FailedRequest {
+                    source: QuerySource {
+                        r#type: "N/A".into(),
+                        name: "N/A".into(),
+                        id: "N/A".into(),
+                    },
+                    status: 0,
+                    error: Some("Query source timed out".to_string()),
+                    reason: None,
+                });
+            }
             Err(_) => {
                 failed_requests.push(FailedRequest {
                     source: QuerySource {
@@ -100,7 +124,7 @@ pub async fn query_coco_fusion<R: Runtime>(
     };
 
     let mut final_hits = Vec::new();
-    let mut seen_docs = std::collections::HashSet::new(); // To track documents we've already added
+    let mut seen_docs = HashSet::new(); // To track documents we've already added
 
     // Distribute hits fairly across sources
     for (_source_id, hits) in &mut hits_per_source {