fix: request idle timeout should be applied differently for worker kind (#636)

Author: nyannyacha

cli/src/flags.rs

@@ -208,7 +208,15 @@ fn get_start_command() -> Command {
         .value_parser(value_parser!(u64)),
     )
     .arg(
-      arg!(--"request-idle-timeout" <MILLISECONDS>)
+      arg!(--"main-worker-request-idle-timeout" <MILLISECONDS>)
+        .help(concat!(
+          "Maximum time in milliseconds that can be waited from when a ",
+          "worker takes over the request (disabled by default)"
+        ))
+        .value_parser(value_parser!(u64)),
+    )
+    .arg(
+      arg!(--"user-worker-request-idle-timeout" <MILLISECONDS>)
         .help(concat!(
           "Maximum time in milliseconds that can be waited from when a ",
           "worker takes over the request (disabled by default)"

cli/src/main.rs

@@ -13,6 +13,7 @@ use anyhow::Context;
 use anyhow::Error;
 use base::server;
 use base::server::Builder;
+use base::server::RequestIdleTimeout;
 use base::server::ServerFlags;
 use base::server::Tls;
 use base::utils::units::percentage_value;
@@ -180,8 +181,12 @@ fn main() -> Result<ExitCode, anyhow::Error> {
           sub_matches.get_one::<usize>("max-parallelism").cloned();
         let maybe_request_wait_timeout =
           sub_matches.get_one::<u64>("request-wait-timeout").cloned();
-        let maybe_request_idle_timeout =
-          sub_matches.get_one::<u64>("request-idle-timeout").cloned();
+        let maybe_main_worker_request_idle_timeout = sub_matches
+          .get_one::<u64>("main-worker-request-idle-timeout")
+          .cloned();
+        let maybe_user_worker_request_idle_timeout = sub_matches
+          .get_one::<u64>("user-worker-request-idle-timeout")
+          .cloned();
         let maybe_request_read_timeout =
           sub_matches.get_one::<u64>("request-read-timeout").cloned();
 
@@ -249,7 +254,10 @@ fn main() -> Result<ExitCode, anyhow::Error> {
           graceful_exit_keepalive_deadline_ms,
           event_worker_exit_deadline_sec,
           request_wait_timeout_ms: maybe_request_wait_timeout,
-          request_idle_timeout_ms: maybe_request_idle_timeout,
+          request_idle_timeout: RequestIdleTimeout::from_millis(
+            maybe_main_worker_request_idle_timeout,
+            maybe_user_worker_request_idle_timeout,
+          ),
           request_read_timeout_ms: maybe_request_read_timeout,
           request_buffer_size: Some(request_buffer_size),
 

crates/base/src/server.rs

@@ -22,6 +22,7 @@ use either::Either::Left;
 use either::Either::Right;
 use enum_as_inner::EnumAsInner;
 use ext_runtime::SharedMetricSource;
+use ext_workers::context::WorkerKind;
 use ext_workers::context::WorkerRequestMsg;
 use futures_util::future::poll_fn;
 use futures_util::future::BoxFuture;
@@ -278,6 +279,29 @@ pub enum OtelKind {
   Both,
 }
 
+#[derive(Debug, Clone, Copy, Default)]
+pub struct RequestIdleTimeout {
+  main: Option<Duration>,
+  user: Option<Duration>,
+}
+
+impl RequestIdleTimeout {
+  pub fn from_millis(main: Option<u64>, user: Option<u64>) -> Self {
+    Self {
+      main: main.map(Duration::from_millis),
+      user: user.map(Duration::from_millis),
+    }
+  }
+
+  pub fn get(&self, kind: WorkerKind) -> Option<Duration> {
+    match kind {
+      WorkerKind::MainWorker => self.main,
+      WorkerKind::UserWorker => self.user,
+      WorkerKind::EventsWorker => None,
+    }
+  }
+}
+
 #[derive(Debug, Default, Clone, Copy)]
 pub struct ServerFlags {
   pub otel: Option<OtelKind>,
@@ -291,8 +315,8 @@ pub struct ServerFlags {
   pub graceful_exit_keepalive_deadline_ms: Option<u64>,
   pub event_worker_exit_deadline_sec: u64,
   pub request_wait_timeout_ms: Option<u64>,
-  pub request_idle_timeout_ms: Option<u64>,
   pub request_read_timeout_ms: Option<u64>,
+  pub request_idle_timeout: RequestIdleTimeout,
   pub request_buffer_size: Option<u64>,
 
   pub beforeunload_wall_clock_pct: Option<u8>,

crates/base/src/worker/worker_surface_creation.rs

@@ -70,7 +70,7 @@ mod request {
     duplex_stream_tx: mpsc::UnboundedSender<DuplexStreamEntry>,
     msg: WorkerRequestMsg,
   ) -> Result<(), anyhow::Error> {
-    let request_idle_timeout_ms = flags.request_idle_timeout_ms;
+    let request_idle_timeout_dur = flags.request_idle_timeout.get(worker_kind);
     let request_buf_size = flags.request_buffer_size.unwrap_or_else(|| {
       const KIB: usize = 1024;
       static CHECK: Lazy<AtomicFlag> = Lazy::new(AtomicFlag::default);
@@ -125,7 +125,7 @@ mod request {
                   tokio::spawn(relay_upgraded_request_and_response(
                     req_upgrade,
                     parts,
-                    request_idle_timeout_ms,
+                    request_idle_timeout_dur,
                   ));
 
                   return;
@@ -144,8 +144,8 @@ mod request {
     tokio::task::yield_now().await;
 
     let maybe_cancel_fut = async move {
-      if let Some(timeout_ms) = request_idle_timeout_ms {
-        sleep(Duration::from_millis(timeout_ms)).await;
+      if let Some(dur) = request_idle_timeout_dur {
+        sleep(dur).await;
       } else {
         pending::<()>().await;
         unreachable!()
@@ -185,18 +185,17 @@ mod request {
       }
     }
 
-    if let Some(timeout_ms) = flags.request_idle_timeout_ms {
+    if let Some(dur) = request_idle_timeout_dur {
       let headers = res.headers();
       let is_streamed_response =
         !headers.contains_key(http_v02::header::CONTENT_LENGTH);
 
       if is_streamed_response {
-        let duration = Duration::from_millis(timeout_ms);
         let (parts, body) = res.into_parts();
 
         drop(res_tx.send(Ok(Response::from_parts(
           parts,
-          Body::wrap_stream(CancelOnWriteTimeout::new(body, duration)),
+          Body::wrap_stream(CancelOnWriteTimeout::new(body, dur)),
         ))));
 
         return Ok(());
@@ -210,14 +209,11 @@ mod request {
   async fn relay_upgraded_request_and_response(
     downstream: OnUpgrade,
     parts: http1::Parts<io::DuplexStream>,
-    maybe_idle_timeout: Option<u64>,
+    maybe_idle_timeout: Option<Duration>,
   ) {
     let upstream = Upgraded2::new(parts.io, parts.read_buf);
-    let mut upstream = if let Some(timeout_ms) = maybe_idle_timeout {
-      ReadTimeoutStream::with_timeout(
-        upstream,
-        Duration::from_millis(timeout_ms),
-      )
+    let mut upstream = if let Some(dur) = maybe_idle_timeout {
+      ReadTimeoutStream::with_timeout(upstream, dur)
     } else {
       ReadTimeoutStream::with_bypass(upstream)
     };

crates/base/tests/integration_tests.rs

@@ -21,6 +21,7 @@ use base::integration_test;
 use base::integration_test_listen_fut;
 use base::integration_test_with_server_flag;
 use base::server::Builder;
+use base::server::RequestIdleTimeout;
 use base::server::ServerEvent;
 use base::server::ServerFlags;
 use base::server::ServerHealth;
@@ -1860,7 +1861,7 @@ async fn test_request_idle_timeout_no_streamed_response(
 
   integration_test_with_server_flag!(
     ServerFlags {
-      request_idle_timeout_ms: Some(1000),
+      request_idle_timeout: RequestIdleTimeout::from_millis(None, Some(1000)),
       ..Default::default()
     },
     "./test_cases/main",
@@ -1918,7 +1919,7 @@ async fn test_request_idle_timeout_streamed_response(maybe_tls: Option<Tls>) {
 
   integration_test_with_server_flag!(
     ServerFlags {
-      request_idle_timeout_ms: Some(2000),
+      request_idle_timeout: RequestIdleTimeout::from_millis(None, Some(2000)),
       ..Default::default()
     },
     "./test_cases/main",
@@ -1991,7 +1992,7 @@ async fn test_request_idle_timeout_streamed_response_first_chunk_timeout(
 
   integration_test_with_server_flag!(
     ServerFlags {
-      request_idle_timeout_ms: Some(1000),
+      request_idle_timeout: RequestIdleTimeout::from_millis(None, Some(1000)),
       ..Default::default()
     },
     "./test_cases/main",
@@ -2079,7 +2080,7 @@ async fn test_request_idle_timeout_websocket_deno(
 
   integration_test_with_server_flag!(
     ServerFlags {
-      request_idle_timeout_ms: Some(1000),
+      request_idle_timeout: RequestIdleTimeout::from_millis(None, Some(1000)),
       ..Default::default()
     },
     "./test_cases/main",