More tests; recv header and keep-alive timeouts.

Author: alistairjevans

bin/run-server.rb

@@ -0,0 +1,60 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+$LOAD_PATH.unshift File.expand_path("../lib", __dir__)
+
+puts "Loading hyper_ruby"
+
+require "hyper_ruby"
+require "json"
+
+# Create and configure the server
+server = HyperRuby::Server.new
+config = {
+  bind_address: ENV.fetch("BIND_ADDRESS", "127.0.0.1:3000"),
+  tokio_threads: ENV.fetch("TOKIO_THREADS", "1").to_i,
+  debug: ENV.fetch("DEBUG", "false") == "true",
+  recv_timeout: ENV.fetch("RECV_TIMEOUT", "30000").to_i
+}
+server.configure(config)
+
+puts "Starting server with config: #{config}"
+
+# Start the server
+server.start
+
+puts "Server started"
+
+# Create a worker thread to handle requests
+worker = Thread.new do
+  server.run_worker do |request|
+    buffer = String.new(capacity: 1024)
+    request.fill_body(buffer)
+
+    # Create a response that echoes back request details
+    response_data = {
+      method: request.http_method,
+      path: request.path,
+      headers: request.headers,
+      body: buffer
+    }
+
+    HyperRuby::Response.new(
+      200,
+      { "Content-Type" => "application/json" },
+      JSON.pretty_generate(response_data)
+    )
+  end
+end
+
+puts "Server running at #{config[:bind_address]}"
+puts "Press Ctrl+C to stop"
+
+# Wait for Ctrl+C
+begin
+  sleep
+rescue Interrupt
+  puts "\nShutting down..."
+  server.stop
+  worker.join
+end 

ext/hyper_ruby/Cargo.toml

@@ -17,7 +17,7 @@ tokio-stream = { version = "0.1", features = ["net"] }
 crossbeam-channel = "0.5.14"
 rb-sys = "0.9.110"
 hyper = { version = "1.0", features = ["http1", "http2", "server"] }
-hyper-util = { version = "0.1", features = ["tokio", "server", "http1", "http2"] }
+hyper-util = { version = "0.1", features = ["tokio", "server", "server-auto", "http1", "http2"] }
 http-body-util = "0.1.2"
 jemallocator = { version = "0.5.4", features = ["disable_initial_exec_tls"] }
 futures = "0.3.31"

ext/hyper_ruby/src/grpc.rs

@@ -115,8 +115,7 @@ pub fn create_grpc_error_response(http_status: u16, grpc_status: u32, message: &
     // Create trailers
     let mut trailers = HeaderMap::new();
     trailers.insert("grpc-status", grpc_status.to_string().parse().unwrap());
-    trailers.insert("grpc-accept-encoding", "identity".parse().unwrap());
-    trailers.insert("accept-encoding", "identity".parse().unwrap());
+    trailers.insert("grpc-accept-encoding", "identity,gzip,deflate,zstd".parse().unwrap());
 
     // Add grpc-message if provided
     if !message.is_empty() {

ext/hyper_ruby/src/lib.rs

@@ -36,6 +36,11 @@ use log::{debug, info, warn};
 use env_logger;
 use crate::response::BodyWithTrailers;
 use std::sync::Once;
+use tokio::time::timeout;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tokio::time::{Sleep, sleep};
 
 static LOGGER_INIT: Once = Once::new();
 
@@ -47,6 +52,7 @@ struct ServerConfig {
     bind_address: String,
     tokio_threads: Option<usize>,
     debug: bool,
+    recv_timeout: u64,
 }
 
 impl ServerConfig {
@@ -55,6 +61,7 @@ impl ServerConfig {
             bind_address: String::from("127.0.0.1:3000"),
             tokio_threads: None,
             debug: false,
+            recv_timeout: 30000, // Default 30 second timeout
         }
     }
 }
@@ -101,6 +108,10 @@ impl Server {
             server_config.debug = bool::try_convert(debug)?;
         }
 
+        if let Some(recv_timeout) = config.get(magnus::Symbol::new("recv_timeout")) {
+            server_config.recv_timeout = u64::try_convert(recv_timeout)?;
+        }
+
         // Initialize logging if debug is enabled, but only do it once
         if server_config.debug {
             LOGGER_INIT.call_once(|| {
@@ -215,16 +226,19 @@ impl Server {
             let work_tx = work_tx.clone();
 
             let server_task = tokio::spawn(async move {
+                let timer = hyper_util::rt::TokioTimer::new();
+
                 if config.bind_address.starts_with("unix:") {
                     let path = config.bind_address.trim_start_matches("unix:");
                     let listener = UnixListener::bind(path).unwrap();
 
                     loop {
                         let (stream, _) = listener.accept().await.unwrap();
                         let work_tx = work_tx.clone();
+                        let timer = timer.clone();
 
                         tokio::task::spawn(async move {
-                            handle_connection(stream, work_tx).await;
+                            handle_connection(stream, work_tx, config.recv_timeout, timer).await;
                         });
                     }
                 } else {
@@ -235,9 +249,10 @@ impl Server {
                     loop {
                         let (stream, _) = listener.accept().await.unwrap();
                         let work_tx = work_tx.clone();
+                        let timer = timer.clone();
 
                         tokio::task::spawn(async move {
-                            handle_connection(stream, work_tx).await;
+                            handle_connection(stream, work_tx, config.recv_timeout, timer).await;
                         });
                     }
                 }
@@ -282,19 +297,27 @@ impl Server {
 async fn handle_request(
     req: HyperRequest<Incoming>,
     work_tx: Arc<crossbeam_channel::Sender<RequestWithCompletion>>,
+    recv_timeout: u64,
 ) -> Result<HyperResponse<BodyWithTrailers>, Error> {
     debug!("Received request: {:?}", req);
     debug!("HTTP version: {:?}", req.version());
     debug!("Headers: {:?}", req.headers());
 
     let (parts, body) = req.into_parts();
 
-    // Collect the body
-    let body_bytes = match body.collect().await {
-        Ok(collected) => collected.to_bytes(),
-        Err(e) => {
+    // Collect the body with timeout
+    let body_bytes = match timeout(
+        std::time::Duration::from_millis(recv_timeout),
+        body.collect()
+    ).await {
+        Ok(Ok(collected)) => collected.to_bytes(),
+        Ok(Err(e)) => {
             debug!("Error collecting body: {:?}", e);
             return Err(e);
+        },
+        Err(_) => {
+            debug!("Timeout collecting body");
+            return Ok(create_timeout_response());
         }
     };
 
@@ -336,23 +359,42 @@ async fn handle_request(
     }
 }
 
+fn create_timeout_response() -> HyperResponse<BodyWithTrailers> {
+    let builder = HyperResponse::builder()
+        .status(StatusCode::REQUEST_TIMEOUT)
+        .header("content-type", "text/plain");
+    
+    builder.body(BodyWithTrailers::new(Bytes::from("Request timed out while receiving body"), None))
+        .unwrap()
+}
+
 async fn handle_connection(
     stream: impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
     work_tx: Arc<crossbeam_channel::Sender<RequestWithCompletion>>,
+    recv_timeout: u64,
+    timer: hyper_util::rt::TokioTimer,
 ) {
     info!("New connection established");
 
     let service = service_fn(move |req: HyperRequest<Incoming>| {
         debug!("Service handling request");
         let work_tx = work_tx.clone();
-        handle_request(req, work_tx)
+        handle_request(req, work_tx, recv_timeout)
     });
 
     let io = TokioIo::new(stream);
 
-    debug!("Setting up HTTP/2 connection");
-    let builder = auto::Builder::new(hyper_util::rt::TokioExecutor::new());
-    
+    debug!("Setting up connection");
+    let mut builder = auto::Builder::new(hyper_util::rt::TokioExecutor::new());
+
+    builder.http1()
+           .header_read_timeout(std::time::Duration::from_millis(recv_timeout))
+           .timer(timer.clone());
+
+    builder.http2()
+           .keep_alive_interval(std::time::Duration::from_secs(10))
+           .timer(timer);
+
     if let Err(err) = builder
         .serve_connection(io, service)
         .await

ext/hyper_ruby/src/response.rs

@@ -20,6 +20,7 @@ impl std::error::Error for ResponseError {}
 #[derive(Debug, Clone)]
 pub struct BodyWithTrailers {
     data: Bytes,
+    data_sent: bool,
     trailers_sent: bool,
     trailers: Option<HeaderMap>,
 }
@@ -28,6 +29,7 @@ impl BodyWithTrailers {
     pub fn new(data: Bytes, trailers: Option<HeaderMap>) -> Self {
         Self {
             data,
+            data_sent: false,
             trailers_sent: false,
             trailers,
         }
@@ -46,9 +48,9 @@ impl Body for BodyWithTrailers {
         mut self: Pin<&mut Self>,
         _cx: &mut std::task::Context<'_>,
     ) -> std::task::Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
-        if !self.data.is_empty() {
+        if !self.data_sent && !self.data.is_empty() {
+            self.data_sent = true;
             let data = self.data.clone();
-            self.data = Bytes::new();
             return std::task::Poll::Ready(Some(Ok(Frame::data(data))));
         }
 
@@ -90,12 +92,20 @@ impl Response {
         if body.len() > 0 {
             unsafe {
                 let rust_body = Bytes::copy_from_slice(body.as_slice());
+                builder_headers.insert(
+                    HeaderName::from_static("content-length"),
+                    rust_body.len().to_string().try_into().unwrap()
+                );
                 match builder.body(BodyWithTrailers::new(rust_body, None)) {
                     Ok(response) => Ok(Self { response }),
                     Err(_) => Err(MagnusError::new(magnus::exception::runtime_error(), "Failed to create response"))
                 }
             }
         } else {
+            builder_headers.insert(
+                HeaderName::from_static("content-length"),
+                "0".try_into().unwrap()
+            );
             match builder.body(BodyWithTrailers::new(Bytes::new(), None)) {
                 Ok(response) => Ok(Self { response }),
                 Err(_) => Err(MagnusError::new(magnus::exception::runtime_error(), "Failed to create response"))
@@ -137,8 +147,7 @@ impl GrpcResponse {
 
         let mut trailers = HeaderMap::new();
         trailers.insert("grpc-status", status.to_string().parse().unwrap());
-        trailers.insert("grpc-accept-encoding", "identity".parse().unwrap());
-        trailers.insert("accept-encoding", "identity".parse().unwrap());
+        trailers.insert("grpc-accept-encoding", "identity,gzip,deflate,zstd".parse().unwrap());
 
         Ok(Self { response: builder.body(BodyWithTrailers::new(framed_message, Some(trailers))).unwrap() })
     }
@@ -153,8 +162,7 @@ impl GrpcResponse {
 
         let mut trailers = HeaderMap::new();
         trailers.insert("grpc-status", status_num.to_string().parse().unwrap());
-        trailers.insert("grpc-accept-encoding", "identity".parse().unwrap());
-        trailers.insert("accept-encoding", "identity".parse().unwrap());
+        trailers.insert("grpc-accept-encoding", "identity,gzip,deflate,zstd".parse().unwrap());
 
         if !message_str.is_empty() {
             trailers.insert("grpc-message", message_str.parse().unwrap());