Fix LB timeout in QueryServer

server/src/handlers/http/modal/server.rs

@@ -36,16 +36,18 @@ use crate::storage;
 use crate::sync;
 use crate::users::dashboards::DASHBOARDS;
 use crate::users::filters::FILTERS;
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
 
-use actix_web::web::resource;
+use actix_web::web::{resource, Data};
 use actix_web::Resource;
 use actix_web::Scope;
 use actix_web::{web, App, HttpServer};
 use actix_web_prometheus::PrometheusMetrics;
 use actix_web_static_files::ResourceFiles;
 use async_trait::async_trait;
+use tokio::sync::Mutex;
 
 use crate::{
     handlers::http::{
@@ -84,7 +86,10 @@ impl ParseableServer for Server {
         };
 
         let create_app_fn = move || {
+            let query_map: query::QueryMap = Arc::new(Mutex::new(HashMap::new()));
+
             App::new()
+                .app_data(Data::new(query_map))
                 .wrap(prometheus.clone())
                 .configure(|cfg| Server::configure_routes(cfg, oidc_client.clone()))
                 .wrap(actix_web::middleware::Logger::default())

server/src/handlers/http/query.rs

@@ -18,7 +18,7 @@
 
 use actix_web::http::header::ContentType;
 use actix_web::web::{self, Json};
-use actix_web::{FromRequest, HttpRequest, Responder};
+use actix_web::{FromRequest, HttpRequest, HttpResponse};
 use chrono::{DateTime, Utc};
 use datafusion::common::tree_node::TreeNode;
 use datafusion::error::DataFusionError;
@@ -28,10 +28,13 @@ use http::StatusCode;
 use std::collections::HashMap;
 use std::pin::Pin;
 use std::sync::Arc;
-use std::time::Instant;
+use std::time::{Duration, Instant};
+use tokio::sync::Mutex;
+use tokio::time::sleep;
 
 use crate::event::error::EventError;
 use crate::handlers::http::fetch_schema;
+use crate::rbac::map::SessionKey;
 use arrow_array::RecordBatch;
 
 use crate::event::commit_schema;
@@ -41,7 +44,7 @@ use crate::option::{Mode, CONFIG};
 use crate::query::error::ExecuteError;
 use crate::query::Query as LogicalQuery;
 use crate::query::{TableScanVisitor, QUERY_SESSION};
-use crate::querycache::{CacheMetadata, QueryCacheManager};
+use crate::querycache::{generate_hash, CacheMetadata, QueryCacheManager};
 use crate::rbac::role::{Action, Permission};
 use crate::rbac::Users;
 use crate::response::QueryResponse;
@@ -50,7 +53,7 @@ use crate::storage::ObjectStorageError;
 use crate::utils::actix::extract_session_key_from_req;
 
 /// Query Request through http endpoint.
-#[derive(Debug, serde::Deserialize, serde::Serialize)]
+#[derive(Debug, serde::Deserialize, serde::Serialize, Clone)]
 #[serde(rename_all = "camelCase")]
 pub struct Query {
     pub query: String,
@@ -64,81 +67,142 @@ pub struct Query {
     pub filter_tags: Option<Vec<String>>,
 }
 
+pub type QueryMap = Arc<Mutex<HashMap<u64, QueryStatus>>>;
 enum QueryStatus {
     Processing,
-    Result(Query),
+    Result(QueryResponse),
 }
 
-pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Responder, QueryError> {
+pub async fn query(
+    req: HttpRequest,
+    query_request: Query,
+    query_map: web::Data<QueryMap>,
+) -> Result<HttpResponse, QueryError> {
     let session_state = QUERY_SESSION.state();
 
-    // get the logical plan and extract the table name
+    // Generate hash for the query based on start, end, and query string
+    let hash = generate_hash(
+        &query_request.start_time,
+        &query_request.end_time,
+        &query_request.query,
+    );
+
+    // Check if the query is already being processed or completed
+    {
+        let query_map = query_map.lock().await;
+        if let Some(status) = query_map.get(&hash) {
+            match status {
+                QueryStatus::Processing => {
+                    // Wait for 55 seconds and check again
+                    sleep(Duration::from_secs(55)).await;
+                    if let Some(QueryStatus::Result(response)) = query_map.get(&hash) {
+                        return Ok(response.clone().to_http()?);
+                    } else {
+                        return Ok(HttpResponse::Accepted().finish());
+                    }
+                }
+                QueryStatus::Result(response) => {
+                    return Ok(response.clone().to_http()?);
+                }
+            }
+        }
+    }
+
+    // Insert the query into the map as Processing
+    {
+        let mut query_map = query_map.lock().await;
+        query_map.insert(hash.clone(), QueryStatus::Processing);
+    }
+
+    // Clone necessary data for the spawned task
+    let query_request_clone = query_request.clone();
+    let hash_clone = hash.clone();
+    let session_state_clone = session_state.clone();
+    let creds = extract_session_key_from_req(&req).unwrap().to_owned();
+
+    // Spawn a separate task to process the query and cache the results
+    tokio::spawn(async move {
+        let mut query_map = query_map.lock().await;
+
+        let result = process_query(
+            query_request_clone,
+            Arc::new(session_state_clone),
+            creds,
+        ).await;
+
+        // Update the query status in the map
+        match result {
+            Ok(response) => {
+                query_map.insert(hash_clone, QueryStatus::Result(response));
+            }
+            Err(err) => {
+                log::error!("Error processing query: {:?}", err);
+                query_map.remove(&hash_clone);
+            }
+        }
+    });
+
+    // Wait for 55 seconds and respond with HTTP 202
+    sleep(Duration::from_secs(55)).await;
+    Ok(HttpResponse::Accepted().finish())
+}
+
+async fn process_query(
+    query_request: Query,
+    session_state: Arc<SessionState>,
+    creds: SessionKey,
+) -> Result<QueryResponse, QueryError> {
     let raw_logical_plan = session_state
         .create_logical_plan(&query_request.query)
-        .await?;
+        .await.unwrap();
 
-    // create a visitor to extract the table name
     let mut visitor = TableScanVisitor::default();
     let _ = raw_logical_plan.visit(&mut visitor);
-    let _stream = visitor
+
+    let visitor_clone = visitor.clone();
+    let stream = visitor
         .top()
         .ok_or_else(|| QueryError::MalformedQuery("Table Name not found in SQL"))?;
 
-    let query_cache_manager = QueryCacheManager::global(CONFIG.parseable.query_cache_size).await?;
-
-    // deal with cached data
-    if let Ok(results) = get_results_from_cache(
-        query_cache_manager,
-        &query_request.start_time,
-        &query_request.end_time,
-        &query_request.query,
-        query_request.send_null,
-        query_request.fields,
-    )
-    .await
-    {
-        return results.to_http();
-    };
-
-    let tables = visitor.into_inner();
+    // Process the query
+    let tables = visitor_clone.into_inner();
     update_schema_when_distributed(tables).await?;
+
     let mut query: LogicalQuery = into_query(&query_request, &session_state).await?;
 
-    let creds = extract_session_key_from_req(&req)?;
     let permissions = Users.get_permissions(&creds);
 
-    let table_name = query
-        .first_table_name()
+    let table_name = query.first_table_name()
         .ok_or_else(|| QueryError::MalformedQuery("No table name found in query"))?;
 
     authorize_and_set_filter_tags(&mut query, permissions, &table_name)?;
 
     let time = Instant::now();
     let (records, fields) = query.execute(table_name.clone()).await?;
-    // deal with cache saving
-    if let Err(err) = put_results_in_cache(
-        query_cache_manager,
-        &table_name,
+
+    let query_cache_manager = QueryCacheManager::global(CONFIG.parseable.query_cache_size)
+        .await?;
+
+    // Cache the results
+    put_results_in_cache(
+        &query_cache_manager,
+        stream,
         &records,
         query.start.to_rfc3339(),
         query.end.to_rfc3339(),
-        query_request.query,
+        query_request.query.clone(),
     )
-    .await
-    {
-        log::error!("{}", err);
-    };
+    .await?;
 
+    // Create the response
     let response = QueryResponse {
         records,
         fields,
         fill_null: query_request.send_null,
         with_fields: query_request.fields,
-    }
-    .to_http()?;
+    };
 
     let time = time.elapsed().as_secs_f64();
-
     QUERY_EXECUTE_TIME
         .with_label_values(&[&table_name])
         .observe(time);

server/src/query.rs

@@ -176,7 +176,7 @@ impl Query {
     }
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Clone)]
 pub(crate) struct TableScanVisitor {
     tables: Vec<String>,
 }

server/src/querycache.rs

@@ -390,7 +390,7 @@ impl QueryCacheManager {
     }
 }
 
-fn generate_hash(start: &str, end: &str, query: &str) -> u64 {
+pub fn generate_hash(start: &str, end: &str, query: &str) -> u64 {
     let mut hasher = DefaultHasher::new();
     start.hash(&mut hasher);
     end.hash(&mut hasher);

server/src/response.rs

@@ -23,12 +23,13 @@ use crate::{
         record_batches_to_json,
     },
 };
-use actix_web::{web, Responder};
+use actix_web::HttpResponse;
 use datafusion::arrow::record_batch::RecordBatch;
 use itertools::Itertools;
 use serde_json::{json, Value};
 use tonic::{Response, Status};
 
+#[derive(Clone)]
 pub struct QueryResponse {
     pub records: Vec<RecordBatch>,
     pub fields: Vec<String>,
@@ -37,7 +38,7 @@ pub struct QueryResponse {
 }
 
 impl QueryResponse {
-    pub fn to_http(&self) -> Result<impl Responder, QueryError> {
+    pub fn to_http(&self) -> Result<HttpResponse, QueryError> {
         log::info!("{}", "Returning query results");
         let records: Vec<&RecordBatch> = self.records.iter().collect();
         let mut json_records = record_batches_to_json(&records)?;
@@ -62,7 +63,7 @@ impl QueryResponse {
             Value::Array(values)
         };
 
-        Ok(web::Json(response))
+        Ok(HttpResponse::Ok().json(response))
     }
 
     pub fn into_flight(self) -> Result<Response<DoGetStream>, Status> {