Handle termination reason

Author: max-lt

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "openworkers-runner"
-version = "0.1.10"
+version = "0.1.11"
 edition = "2021"
 default-run = "openworkers-runner"
 
@@ -16,8 +16,8 @@ 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" ] }
-openworkers-runtime ={ git = "https://github.com/openworkers/openworkers-runtime", tag = "v0.1.10"}
-# openworkers-runtime =  { path = "../openworkers-runtime" }
+# openworkers-runtime ={ git = "https://github.com/openworkers/openworkers-runtime", tag = "v0.1.10"}
+openworkers-runtime =  { path = "../openworkers-runtime" }
 async-nats = "0.37"
 futures = "0.3"
 serde_json = "1.0.137"

Cargo.toml

@@ -2,6 +2,9 @@ use bytes::Bytes;
 
 use log::debug;
 use log::error;
+use log::warn;
+
+use std::time::Duration;
 
 use tokio::sync::oneshot::channel;
 
@@ -41,7 +44,14 @@ async fn handle_request(data: Data<AppState>, req: HttpRequest, body: Bytes) ->
 
     // Expect x-request-id header
     let request_id = match req.headers().get("x-request-id") {
-        Some(value) => value.to_str().unwrap(),
+        Some(value) => match value.to_str() {
+            Ok(s) => s,
+            Err(_) => {
+                return HttpResponse::BadRequest()
+                    .content_type("text/plain")
+                    .body("Invalid x-request-id header encoding");
+            }
+        },
         None => {
             return HttpResponse::BadRequest()
                 .content_type("text/plain")
@@ -50,17 +60,35 @@ async fn handle_request(data: Data<AppState>, req: HttpRequest, body: Bytes) ->
     };
 
     let host = match req.headers().get("host") {
-        Some(value) => Some(value.to_str().unwrap().to_string()),
+        Some(value) => match value.to_str() {
+            Ok(s) => Some(s.to_string()),
+            Err(_) => {
+                error!("Invalid host header encoding");
+                None
+            }
+        },
         None => None,
     };
 
     let mut worker_id = match req.headers().get("x-worker-id") {
-        Some(value) => Some(value.to_str().unwrap().to_string()),
+        Some(value) => match value.to_str() {
+            Ok(s) => Some(s.to_string()),
+            Err(_) => {
+                error!("Invalid x-worker-id header encoding");
+                None
+            }
+        },
         None => None,
     };
 
     let mut worker_name = match req.headers().get("x-worker-name") {
-        Some(value) => Some(value.to_str().unwrap().to_string()),
+        Some(value) => match value.to_str() {
+            Ok(s) => Some(s.to_string()),
+            Err(_) => {
+                error!("Invalid x-worker-name header encoding");
+                None
+            }
+        },
         None => None,
     };
 
@@ -117,7 +145,7 @@ async fn handle_request(data: Data<AppState>, req: HttpRequest, body: Bytes) ->
 
     // Create a new request to forward to the worker.
     let request = {
-        let mut request: http_v02::Request<Bytes> = http_v02::Request::builder()
+        let mut request: http_v02::Request<Bytes> = match http_v02::Request::builder()
             .uri(format!(
                 "{}://{}{}",
                 req.connection_info().scheme(),
@@ -126,7 +154,15 @@ async fn handle_request(data: Data<AppState>, req: HttpRequest, body: Bytes) ->
             ))
             .method(req.method())
             .body(body)
-            .unwrap();
+        {
+            Ok(r) => r,
+            Err(e) => {
+                error!("Failed to build request: {}", e);
+                return HttpResponse::InternalServerError()
+                    .content_type("text/plain")
+                    .body("Failed to build forwarded request");
+            }
+        };
 
         // Copy headers from the incoming request to the forwarded request.
         let headers = request.headers_mut();
@@ -136,18 +172,26 @@ async fn handle_request(data: Data<AppState>, req: HttpRequest, body: Bytes) ->
 
         // If the worker id is not provided, we add it to the headers.
         if req.headers().get("x-worker-id").is_none() {
-            headers.insert(
-                "x-worker-id",
-                http_v02::HeaderValue::from_str(&worker.id).unwrap(),
-            );
+            match http_v02::HeaderValue::from_str(&worker.id) {
+                Ok(header_value) => {
+                    headers.insert("x-worker-id", header_value);
+                }
+                Err(e) => {
+                    error!("Invalid worker id for header: {}", e);
+                }
+            }
         }
 
         // If the worker name is not provided, we add it to the headers.
         if req.headers().get("x-worker-name").is_none() {
-            headers.insert(
-                "x-worker-name",
-                http_v02::HeaderValue::from_str(&worker.name).unwrap(),
-            );
+            match http_v02::HeaderValue::from_str(&worker.name) {
+                Ok(header_value) => {
+                    headers.insert("x-worker-name", header_value);
+                }
+                Err(e) => {
+                    error!("Invalid worker name for header: {}", e);
+                }
+            }
         }
 
         request
@@ -179,8 +223,9 @@ async fn handle_request(data: Data<AppState>, req: HttpRequest, body: Bytes) ->
     };
 
     let (res_tx, res_rx) = channel::<http_v02::Response<Bytes>>();
+    let (termination_tx, termination_rx) = channel::<openworkers_runner::TerminationReason>();
 
-    openworkers_runner::event_fetch::run_fetch(worker, request, res_tx, data.log_tx.clone(), permit);
+    openworkers_runner::event_fetch::run_fetch(worker, request, res_tx, termination_tx, data.log_tx.clone(), permit);
 
     let response = match res_rx.await {
         Ok(res) => {
@@ -192,9 +237,44 @@ async fn handle_request(data: Data<AppState>, req: HttpRequest, body: Bytes) ->
 
             rb.body(res.body().clone())
         }
-        Err(err) => {
-            error!("worker fetch error: {}, ensure the worker registered a listener for the 'fetch' event", err);
-            HttpResponse::InternalServerError().body(err.to_string())
+        Err(_) => {
+            // Worker didn't send a response, check termination reason
+            use openworkers_runner::TerminationReason;
+
+            let reason = termination_rx.await.unwrap_or(TerminationReason::Exception);
+
+            error!("worker terminated without sending response: {:?}", reason);
+
+            let status = reason.http_status();
+            let body = match reason {
+                TerminationReason::Success => {
+                    // This shouldn't happen - worker completed but didn't send response
+                    "Worker completed but did not send a response (missing fetch event listener?)"
+                }
+                TerminationReason::CpuTimeLimit => {
+                    "Worker exceeded CPU time limit (100ms)"
+                }
+                TerminationReason::WallClockTimeout => {
+                    "Worker exceeded wall-clock time limit (60s)"
+                }
+                TerminationReason::MemoryLimit => {
+                    "Worker exceeded memory limit (128MB)"
+                }
+                TerminationReason::Exception => {
+                    "Worker threw an uncaught exception"
+                }
+                TerminationReason::InitializationError => {
+                    "Worker failed to initialize"
+                }
+                TerminationReason::Terminated => {
+                    "Worker was terminated"
+                }
+            };
+
+            HttpResponse::build(actix_web::http::StatusCode::from_u16(status).unwrap())
+                .content_type("text/plain")
+                .insert_header(("X-Termination-Reason", format!("{:?}", reason)))
+                .body(body)
         }
     };
 
@@ -231,26 +311,65 @@ async fn main() -> std::io::Result<()> {
     }
 
     let db_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set");
-    let pool = PgPoolOptions::new()
-        .max_connections(4)
-        .connect(&db_url)
-        .await
-        .expect("Failed to connect to Postgres");
-
-    // Check postgres connection
-    sqlx::query("SELECT 1")
-        .fetch_one(&pool)
-        .await
-        .expect("Failed to query Postgres");
-    debug!("connected to Postgres");
-
-    // Check NATS connection
-    let nats_client = openworkers_runner::nats::nats_connect().await;
-    nats_client
-        .publish("boot", "0".into())
-        .await
-        .expect("Failed to connect to NATS");
-    debug!("connected to NATS");
+
+    // Retry database connection with exponential backoff
+    let mut retry_count = 0;
+    let max_retries = 5;
+    let pool = loop {
+        match PgPoolOptions::new()
+            .max_connections(20) // Increased from 4
+            .acquire_timeout(Duration::from_secs(5))
+            .connect(&db_url)
+            .await
+        {
+            Ok(pool) => {
+                // Test the connection
+                match sqlx::query("SELECT 1").fetch_one(&pool).await {
+                    Ok(_) => {
+                        debug!("connected to Postgres");
+                        break pool;
+                    }
+                    Err(e) => {
+                        error!("Database connection test failed: {}", e);
+                        if retry_count >= max_retries {
+                            panic!("Failed to connect to database after {} retries", max_retries);
+                        }
+                    }
+                }
+            }
+            Err(e) => {
+                retry_count += 1;
+                if retry_count > max_retries {
+                    panic!("Failed to connect to database after {} retries: {}", max_retries, e);
+                }
+                let wait_time = Duration::from_secs(2u64.pow(retry_count.min(5)));
+                warn!("Database connection attempt {} failed: {}. Retrying in {:?}...",
+                      retry_count, e, wait_time);
+                tokio::time::sleep(wait_time).await;
+            }
+        }
+    };
+
+    // Connect to NATS with retries
+    let mut retry_count = 0;
+    loop {
+        match openworkers_runner::nats::nats_connect().await.publish("boot", "0".into()).await {
+            Ok(_) => {
+                debug!("connected to NATS");
+                break;
+            }
+            Err(e) => {
+                retry_count += 1;
+                if retry_count > max_retries {
+                    panic!("Failed to connect to NATS after {} retries: {}", max_retries, e);
+                }
+                let wait_time = Duration::from_secs(2u64.pow(retry_count.min(5)));
+                warn!("NATS connection attempt {} failed: {}. Retrying in {:?}...",
+                      retry_count, e, wait_time);
+                tokio::time::sleep(wait_time).await;
+            }
+        }
+    }
 
     // Start global log publisher
     let log_tx = openworkers_runner::log::start_log_publisher();

bin/main.rs

@@ -6,13 +6,15 @@ use openworkers_runtime::FetchInit;
 use openworkers_runtime::RuntimeLimits;
 use openworkers_runtime::Script;
 use openworkers_runtime::Task;
+use openworkers_runtime::TerminationReason;
 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>>;
+type TerminationTx = tokio::sync::oneshot::Sender<TerminationReason>;
 
 // Default timeout for fetch events
 const FETCH_TIMEOUT_MS: u64 = 64_000; // 64 seconds
@@ -21,13 +23,31 @@ pub fn run_fetch(
     worker: WorkerData,
     req: http_v02::Request<Bytes>,
     res_tx: ResTx,
+    termination_tx: TerminationTx,
     global_log_tx: std::sync::mpsc::Sender<crate::log::LogMessage>,
     permit: OwnedSemaphorePermit,
 ) {
     let (log_tx, log_handler) = crate::log::create_log_handler(worker.id.clone(), global_log_tx);
 
+    let code = match crate::transform::parse_worker_code(&worker) {
+        Ok(code) => code,
+        Err(e) => {
+            log::error!("Failed to parse worker code: {}", e);
+            res_tx
+                .send(
+                    http_v02::Response::builder()
+                        .status(500)
+                        .body(format!("Failed to parse worker code: {}", e).into())
+                        .unwrap(),
+                )
+                .ok(); // Ignore send error
+            termination_tx.send(TerminationReason::InitializationError).ok();
+            return;
+        }
+    };
+
     let script = Script {
-        code: crate::transform::parse_worker_code(&worker),
+        code,
         env: match worker.env {
             Some(env) => Some(env.deref().to_owned()),
             None => None,
@@ -71,15 +91,23 @@ pub fn run_fetch(
 
         // 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}"),
+        let termination_reason = match tokio::time::timeout(timeout_duration, worker.exec(task)).await {
+            Ok(Ok(reason)) => {
+                log::debug!("worker exec completed: reason={:?}", reason);
+                reason
+            }
+            Ok(Err(err)) => {
+                log::error!("worker exec error: {err}");
+                TerminationReason::Exception
+            }
             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
+                log::error!("worker exec timeout after {}ms (outer timeout)", FETCH_TIMEOUT_MS);
+                TerminationReason::WallClockTimeout
             }
-        }
+        };
+
+        // Send termination reason back to the main thread
+        let _ = termination_tx.send(termination_reason);
 
         // CRITICAL: Flush logs before worker is dropped to prevent log loss
         log_handler.flush();