Optimize worker execution with thread pool

max-lt · max-lt · commit 5820db32c01d · 2025-11-22T20:46:26.000+01:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -11,6 +11,8 @@ actix-web = "4.9.0"
 bytes = "1.6.0"
 log = "0.4.21"
 tokio = "1.43.0"
+tokio-util = { version = "0.7", features = ["rt"] }
+once_cell = "1.19"
 env_logger = "0.11.6"
 http_v02 = { package = "http", version = "0.2.12" }
 sqlx = { version = "0.8.3", features = [ "runtime-tokio", "postgres", "uuid", "bigdecimal", "rust_decimal" ] }
diff --git a/bin/main.rs b/bin/main.rs
@@ -153,11 +153,35 @@ async fn handle_request(data: Data<AppState>, req: HttpRequest, body: Bytes) ->
         request
     };
 
+    // Try to acquire a worker slot from the semaphore with timeout
+    let timeout = openworkers_runner::worker_pool::get_worker_wait_timeout();
+    let permit = match tokio::time::timeout(
+        timeout,
+        openworkers_runner::worker_pool::WORKER_SEMAPHORE
+            .clone()
+            .acquire_owned(),
+    )
+    .await
+    {
+        Ok(Ok(permit)) => permit,
+        Ok(Err(_)) => {
+            error!("semaphore closed unexpectedly");
+            return HttpResponse::InternalServerError()
+                .content_type("text/plain")
+                .body("Internal server error");
+        }
+        Err(_) => {
+            debug!("worker pool saturated after {}ms timeout, returning 503", timeout.as_millis());
+            return HttpResponse::ServiceUnavailable()
+                .content_type("text/plain")
+                .body("Server is overloaded, please try again later");
+        }
+    };
+
     let (res_tx, res_rx) = channel::<http_v02::Response<Bytes>>();
 
-    let handle = openworkers_runner::event_fetch::run_fetch(worker, request, res_tx, data.log_tx.clone());
+    openworkers_runner::event_fetch::run_fetch(worker, request, res_tx, data.log_tx.clone(), permit);
 
-    // TODO: select! on res_rx, timeout and handle.join()
     let response = match res_rx.await {
         Ok(res) => {
             let mut rb = HttpResponse::build(res.status());
@@ -176,8 +200,6 @@ async fn handle_request(data: Data<AppState>, req: HttpRequest, body: Bytes) ->
 
     debug!("handle_request done in {}ms", start.elapsed().as_millis());
 
-    handle.join().unwrap();
-
     response
 }
 
diff --git a/src/event_fetch.rs b/src/event_fetch.rs
@@ -1,5 +1,4 @@
 use std::ops::Deref;
-use std::thread::JoinHandle;
 use std::time::Duration;
 
 use bytes::Bytes;
@@ -8,8 +7,10 @@ use openworkers_runtime::RuntimeLimits;
 use openworkers_runtime::Script;
 use openworkers_runtime::Task;
 use openworkers_runtime::Worker;
+use tokio::sync::OwnedSemaphorePermit;
 
 use crate::store::WorkerData;
+use crate::worker_pool::WORKER_POOL;
 
 type ResTx = tokio::sync::oneshot::Sender<http_v02::Response<Bytes>>;
 
@@ -21,7 +22,8 @@ pub fn run_fetch(
     req: http_v02::Request<Bytes>,
     res_tx: ResTx,
     global_log_tx: std::sync::mpsc::Sender<crate::log::LogMessage>,
-) -> JoinHandle<()> {
+    permit: OwnedSemaphorePermit,
+) {
     let (log_tx, log_handler) = crate::log::create_log_handler(worker.id.clone(), global_log_tx);
 
     let script = Script {
@@ -32,63 +34,56 @@ pub fn run_fetch(
         },
     };
 
-    std::thread::spawn(move || {
-        let local = tokio::task::LocalSet::new();
-
-        let tasks = local.spawn_local(async move {
-            log::debug!("create worker");
-
-            let limits = RuntimeLimits {
-                max_cpu_time_ms: 100,           // 100ms CPU time for fetch tasks
-                max_wall_clock_time_ms: 60_000, // 60s total time for fetch tasks
-                ..Default::default()
-            };
-
-            let mut worker = match Worker::new(script, Some(log_tx), Some(limits)).await {
-                Ok(worker) => worker,
-                Err(err) => {
-                    log::error!("failed to create worker: {err}");
-                    res_tx
-                        .send(
-                            http_v02::Response::builder()
-                                .status(500)
-                                .body(format!("failed to create worker: {err}").into())
-                                .unwrap(),
-                        )
-                        .unwrap();
-
-                    return;
-                }
-            };
-
-            let task = Task::Fetch(Some(FetchInit::new(req, res_tx)));
-
-            log::debug!("exec fetch task with {}ms timeout", FETCH_TIMEOUT_MS);
-
-            // Wrap execution with timeout
-            let timeout_duration = Duration::from_millis(FETCH_TIMEOUT_MS);
-            match tokio::time::timeout(timeout_duration, worker.exec(task)).await {
-                Ok(Ok(())) => log::debug!("exec completed"),
-                Ok(Err(err)) => log::error!("exec did not complete: {err}"),
-                Err(_) => {
-                    log::error!("exec timeout after {}ms", FETCH_TIMEOUT_MS);
-                    // Note: Worker may have already sent a response via FetchInit
-                    // If no response was sent, res_tx will be dropped and client gets an error
-                }
+    // Use the global worker pool instead of spawning a new thread
+    WORKER_POOL.spawn_pinned(move || async move {
+        // Keep the permit alive for the entire worker execution
+        // It will be automatically released when this async block completes
+        let _permit = permit;
+
+        log::debug!("create worker");
+
+        let limits = RuntimeLimits {
+            max_cpu_time_ms: 100,           // 100ms CPU time for fetch tasks
+            max_wall_clock_time_ms: 60_000, // 60s total time for fetch tasks
+            ..Default::default()
+        };
+
+        let mut worker = match Worker::new(script, Some(log_tx), Some(limits)).await {
+            Ok(worker) => worker,
+            Err(err) => {
+                log::error!("failed to create worker: {err}");
+                res_tx
+                    .send(
+                        http_v02::Response::builder()
+                            .status(500)
+                            .body(format!("failed to create worker: {err}").into())
+                            .unwrap(),
+                    )
+                    .unwrap();
+
+                return;
             }
+        };
 
-            // CRITICAL: Flush logs before worker is dropped to prevent log loss
-            log_handler.flush();
-        });
+        let task = Task::Fetch(Some(FetchInit::new(req, res_tx)));
 
-        let rt = tokio::runtime::Builder::new_current_thread()
-            .enable_all()
-            .build()
-            .unwrap();
+        log::debug!("exec fetch task with {}ms timeout", FETCH_TIMEOUT_MS);
 
-        match local.block_on(&rt, tasks) {
-            Ok(()) => {}
-            Err(err) => log::error!("failed to wait for end: {err}"),
+        // Wrap execution with timeout
+        let timeout_duration = Duration::from_millis(FETCH_TIMEOUT_MS);
+        match tokio::time::timeout(timeout_duration, worker.exec(task)).await {
+            Ok(Ok(())) => log::debug!("exec completed"),
+            Ok(Err(err)) => log::error!("exec did not complete: {err}"),
+            Err(_) => {
+                log::error!("exec timeout after {}ms", FETCH_TIMEOUT_MS);
+                // Note: Worker may have already sent a response via FetchInit
+                // If no response was sent, res_tx will be dropped and client gets an error
+            }
         }
-    })
+
+        // CRITICAL: Flush logs before worker is dropped to prevent log loss
+        log_handler.flush();
+
+        // Permit is automatically released here when _permit goes out of scope
+    });
 }
diff --git a/src/event_scheduled.rs b/src/event_scheduled.rs
@@ -10,6 +10,7 @@ use serde::Deserialize;
 use serde::Serialize;
 
 use crate::store;
+use crate::worker_pool::WORKER_POOL;
 
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
@@ -25,51 +26,64 @@ fn run_scheduled(
     script: Script,
     global_log_tx: std::sync::mpsc::Sender<crate::log::LogMessage>,
 ) {
-    let (res_tx, res_rx) = tokio::sync::oneshot::channel::<()>();
-
-    let task = Task::Scheduled(Some(ScheduledInit::new(res_tx, data.scheduled_time)));
+    // Try to acquire a worker slot
+    let permit = match crate::worker_pool::WORKER_SEMAPHORE
+        .clone()
+        .try_acquire_owned()
+    {
+        Ok(permit) => permit,
+        Err(_) => {
+            log::warn!(
+                "worker pool saturated, skipping scheduled task for worker: {}",
+                data.worker_id
+            );
+            return;
+        }
+    };
 
     let (log_tx, log_handler) = crate::log::create_log_handler(data.worker_id, global_log_tx);
 
-    std::thread::spawn(move || {
-        let rt = tokio::runtime::Builder::new_current_thread()
-            .enable_all()
-            .build()
-            .unwrap();
-
-        let local = tokio::task::LocalSet::new();
-
-        local.spawn_local(async move {
-            log::debug!("create worker");
-
-            let limits = RuntimeLimits {
-                max_cpu_time_ms: 100,           // 100ms CPU time for scheduled tasks
-                max_wall_clock_time_ms: 60_000, // 60s total time for scheduled tasks
-                ..Default::default()
-            };
-
-            let mut worker = Worker::new(script, Some(log_tx), Some(limits))
-                .await
-                .unwrap();
-
-            log::debug!("exec scheduled task");
-            match worker.exec(task).await {
-                Ok(()) => log::debug!("exec completed"),
-                Err(err) => log::error!("exec did not complete: {err}"),
+    // Use the global worker pool instead of spawning a new thread
+    WORKER_POOL.spawn_pinned(move || async move {
+        // Keep the permit alive for the entire worker execution
+        let _permit = permit;
+        log::debug!("create worker");
+
+        let limits = RuntimeLimits {
+            max_cpu_time_ms: 100,           // 100ms CPU time for scheduled tasks
+            max_wall_clock_time_ms: 60_000, // 60s total time for scheduled tasks
+            ..Default::default()
+        };
+
+        let mut worker = match Worker::new(script, Some(log_tx), Some(limits)).await {
+            Ok(worker) => worker,
+            Err(err) => {
+                log::error!("failed to create scheduled worker: {err}");
+                log_handler.flush();
+                return;
             }
+        };
 
-            // CRITICAL: Flush logs before worker is dropped to prevent log loss
-            log_handler.flush();
-        });
+        // Create the oneshot channel INSIDE the async block so the receiver stays alive
+        let (res_tx, res_rx) = tokio::sync::oneshot::channel::<()>();
+        let task = Task::Scheduled(Some(ScheduledInit::new(res_tx, data.scheduled_time)));
 
-        log::debug!("scheduled task listener started");
+        log::debug!("exec scheduled task");
+        match worker.exec(task).await {
+            Ok(()) => log::debug!("exec completed"),
+            Err(err) => log::error!("exec did not complete: {err}"),
+        }
 
-        match local.block_on(&rt, async { res_rx.await }) {
-            Ok(()) => {}
-            Err(err) => log::error!("failed to wait for end: {err}"),
+        // Wait for the scheduled event to complete
+        match res_rx.await {
+            Ok(()) => log::debug!("scheduled task responded"),
+            Err(err) => log::error!("scheduled task response error: {err}"),
         }
 
-        log::debug!("scheduled task listener stopped");
+        // CRITICAL: Flush logs before worker is dropped to prevent log loss
+        log_handler.flush();
+
+        // Permit is automatically released here when _permit goes out of scope
     });
 }
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -3,4 +3,5 @@ pub mod event_scheduled;
 pub mod log;
 pub mod nats;
 pub mod store;
+pub mod worker_pool;
 mod transform;
diff --git a/src/worker_pool.rs b/src/worker_pool.rs
