moved from query handler to scheduled job

Author: nikhilsinhaparseable

src/handlers/http/modal/ingest_server.rs

@@ -129,6 +129,9 @@ impl ParseableServer for IngestServer {
         let (cancel_tx, cancel_rx) = oneshot::channel();
         thread::spawn(|| sync::handler(cancel_rx));
 
+        // Initialize memory release scheduler
+        crate::memory::init_memory_release_scheduler()?;
+
         tokio::spawn(airplane::server());
 
         // Ingestors shouldn't have to deal with OpenId auth flow

src/handlers/http/modal/query_server.rs

@@ -129,6 +129,9 @@ impl ParseableServer for QueryServer {
             analytics::init_analytics_scheduler()?;
         }
 
+        // Initialize memory release scheduler
+        crate::memory::init_memory_release_scheduler()?;
+
         if init_cluster_metrics_schedular().is_ok() {
             info!("Cluster metrics scheduler started successfully");
         }

src/handlers/http/modal/server.rs

@@ -154,6 +154,9 @@ impl ParseableServer for Server {
             analytics::init_analytics_scheduler()?;
         }
 
+        // Initialize memory release scheduler
+        crate::memory::init_memory_release_scheduler()?;
+
         tokio::spawn(handlers::livetail::server());
         tokio::spawn(handlers::airplane::server());
 

src/handlers/http/query.rs

@@ -51,7 +51,7 @@ use crate::query::error::ExecuteError;
 use crate::query::{CountsRequest, Query as LogicalQuery, execute};
 use crate::query::{QUERY_SESSION, resolve_stream_names};
 use crate::rbac::Users;
-use crate::response::{QueryResponse, force_memory_release};
+use crate::response::QueryResponse;
 use crate::storage::ObjectStorageError;
 use crate::utils::actix::extract_session_key_from_req;
 use crate::utils::time::{TimeParseError, TimeRange};
@@ -245,8 +245,8 @@ async fn handle_non_streaming_query(
         .insert_header((TIME_ELAPSED_HEADER, total_time.as_str()))
         .json(response);
 
-    // Force memory release after HTTP response is fully created
-    force_memory_release();
+    // // Force memory release after HTTP response is fully created
+    // force_memory_release();
 
     Ok(http_response)
 }
@@ -346,8 +346,6 @@ fn create_batch_processor(
             let bytes_result = Bytes::from(format!("{response}\n"));
             drop(response); // Explicit cleanup
 
-            force_memory_release();
-
             Ok(bytes_result)
         }
         Err(e) => Err(actix_web::error::ErrorInternalServerError(e)),

src/lib.rs

@@ -30,6 +30,7 @@ pub mod event;
 pub mod handlers;
 pub mod hottier;
 mod livetail;
+pub mod memory;
 mod metadata;
 pub mod metastore;
 pub mod metrics;

src/main.rs

@@ -32,6 +32,7 @@ use tracing_subscriber::util::SubscriberInitExt;
 use tracing_subscriber::{EnvFilter, Registry, fmt};
 
 // Use jemalloc as the global allocator
+#[cfg(not(target_env = "msvc"))]
 #[global_allocator]
 static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 

src/memory.rs

@@ -0,0 +1,75 @@
+/*
+ * Parseable Server (C) 2022 - 2024 Parseable, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+use std::ffi::CString;
+use std::time::Duration;
+
+use clokwerk::AsyncScheduler;
+use tracing::{info, warn};
+
+/// Force memory release using jemalloc
+pub fn force_memory_release() {
+    // Advance epoch to refresh statistics and trigger potential cleanup
+    if let Err(e) = tikv_jemalloc_ctl::epoch::mib().and_then(|mib| mib.advance()) {
+        warn!("Failed to advance jemalloc epoch: {:?}", e);
+    }
+
+    // Purge each initialized arena
+    if let Ok(n) = tikv_jemalloc_ctl::arenas::narenas::read() {
+        for i in 0..n {
+            if let Ok(name) = CString::new(format!("arena.{i}.purge")) {
+                unsafe {
+                    let ret = tikv_jemalloc_sys::mallctl(
+                        name.as_ptr(),
+                        std::ptr::null_mut(),
+                        std::ptr::null_mut(),
+                        std::ptr::null_mut(),
+                        0,
+                    );
+                    if ret != 0 {
+                        warn!("Arena purge failed for index {i} with code: {ret}");
+                    }
+                }
+            }
+        }
+    } else {
+        warn!("Failed to read jemalloc arenas.narenas");
+    }
+}
+
+/// Initialize memory management scheduler
+pub fn init_memory_release_scheduler() -> anyhow::Result<()> {
+    info!("Setting up scheduler for memory release");
+
+    let mut scheduler = AsyncScheduler::new();
+    scheduler
+        .every(clokwerk::Interval::Hours(1))
+        .run(move || async {
+            info!("Running scheduled memory release");
+            force_memory_release();
+        });
+
+    tokio::spawn(async move {
+        loop {
+            scheduler.run_pending().await;
+            tokio::time::sleep(Duration::from_secs(60)).await; // Check every minute
+        }
+    });
+
+    Ok(())
+}

src/response.rs

@@ -16,14 +16,10 @@
  *
  */
 
-use std::ffi::CString;
-use std::sync::Mutex;
-use std::time::{Duration, Instant};
-
 use crate::{handlers::http::query::QueryError, utils::arrow::record_batches_to_json};
 use datafusion::arrow::record_batch::RecordBatch;
 use serde_json::{Value, json};
-use tracing::{debug, info, warn};
+use tracing::info;
 
 pub struct QueryResponse {
     pub records: Vec<RecordBatch>,
@@ -70,49 +66,3 @@ impl QueryResponse {
         Ok(response)
     }
 }
-
-impl Drop for QueryResponse {
-    fn drop(&mut self) {
-        force_memory_release();
-    }
-}
-
-// Rate-limited memory release with proper error handling
-static LAST_PURGE: Mutex<Option<Instant>> = Mutex::new(None);
-const PURGE_INTERVAL: Duration = Duration::from_secs(3600); // 1 hour
-pub fn force_memory_release() {
-    {
-        let mut last_purge = LAST_PURGE.lock().unwrap();
-        if let Some(last) = *last_purge {
-            if last.elapsed() < PURGE_INTERVAL {
-                return;
-            }
-        }
-        *last_purge = Some(Instant::now());
-    }
-
-    // Advance epoch to refresh statistics and trigger potential cleanup
-    if let Err(e) = tikv_jemalloc_ctl::epoch::mib().and_then(|mib| mib.advance()) {
-        warn!("Failed to advance jemalloc epoch: {:?}", e);
-    }
-
-    // Purge all arenas using MALLCTL_ARENAS_ALL
-    if let Ok(arena_purge) = CString::new("arena.4096.purge") {
-        unsafe {
-            let ret = tikv_jemalloc_sys::mallctl(
-                arena_purge.as_ptr(),
-                std::ptr::null_mut(), // oldp (not reading)
-                std::ptr::null_mut(), // oldlenp (not reading)
-                std::ptr::null_mut(), // newp (void operation)
-                0,                    // newlen (void operation)
-            );
-            if ret != 0 {
-                warn!("Arena purge failed with code: {}", ret);
-            } else {
-                debug!("Successfully purged all jemalloc arenas");
-            }
-        }
-    } else {
-        warn!("Failed to create CString for arena purge");
-    }
-}

src/utils/arrow/mod.rs

@@ -65,7 +65,7 @@ pub fn record_batches_to_json(records: &[RecordBatch]) -> Result<Vec<Map<String,
     // Use a cursor to avoid extra allocations during parsing
     let json_rows: Vec<Map<String, Value>> = {
         let cursor = std::io::Cursor::new(buf);
-        serde_json::from_reader(cursor).unwrap_or_else(|_| Vec::with_capacity(0))
+        serde_json::from_reader(cursor)?
     };
 
     Ok(json_rows)