feat: add --max-requests for worker process recycling

robyn/__init__.py

@@ -678,6 +678,7 @@ def start(self, host: str = "127.0.0.1", port: int = 8080, _check_port: bool = T
             open_browser,
             client_timeout,
             keep_alive_timeout,
+            self.config.max_requests,
         )
 
 

robyn/argument_parser.py

@@ -2,6 +2,16 @@
 import os
 
 
+def _positive_int(value: str) -> int:
+    try:
+        ivalue = int(value)
+    except ValueError:
+        raise argparse.ArgumentTypeError(f"invalid int value: '{value}'")
+    if ivalue <= 0:
+        raise argparse.ArgumentTypeError(f"--max-requests must be a positive integer, got {ivalue}")
+    return ivalue
+
+
 class Config:
     def __init__(self) -> None:
         parser = argparse.ArgumentParser(description="Robyn, a fast async web framework with a rust runtime.")
@@ -84,6 +94,14 @@ def __init__(self) -> None:
             default=False,
             help="Fast mode. It sets the optimal values for processes, workers and log level. However, you can override them.",
         )
+        parser.add_argument(
+            "--max-requests",
+            dest="max_requests",
+            type=_positive_int,
+            default=None,
+            required=False,
+            help="Recycle worker processes after this many requests. Helps contain memory leaks. [Default: None (disabled)]",
+        )
 
         args, unknown_args = parser.parse_known_args()
         self.fast = args.fast
@@ -99,6 +117,7 @@ def __init__(self) -> None:
         self.file_path = None
         self.disable_openapi = args.disable_openapi
         self.log_level = args.log_level
+        self.max_requests = args.max_requests
 
         if self.fast:
             # doing this here before every other check

robyn/processpool.py

@@ -30,7 +30,10 @@ def run_processes(
     open_browser: bool,
     client_timeout: int = 30,
     keep_alive_timeout: int = 20,
+    max_requests: Optional[int] = None,
 ) -> List[Process]:
+    import time
+
     socket = SocketHeld(url, port)
 
     process_pool = init_processpool(
@@ -48,12 +51,24 @@ def run_processes(
         excluded_response_headers_paths,
         client_timeout,
         keep_alive_timeout,
+        max_requests,
     )
 
+    shutting_down = False
+
     def terminating_signal_handler(_sig, _frame):
-        logger.info("Terminating server!!", bold=True)
+        nonlocal shutting_down
+        shutting_down = True
+        logger.info("Gracefully shutting down server...", bold=True)
+        for process in process_pool:
+            process.terminate()
         for process in process_pool:
-            process.kill()
+            process.join(timeout=30)
+            if process.is_alive():
+                logger.warning("Process %s did not shut down in time, forcing kill.", process.pid)
+                process.kill()
+                process.join(timeout=5)
+        sys.exit(0)
 
     signal.signal(signal.SIGINT, terminating_signal_handler)
     signal.signal(signal.SIGTERM, terminating_signal_handler)
@@ -63,8 +78,39 @@ def terminating_signal_handler(_sig, _frame):
         webbrowser.open_new_tab(f"http://{url}:{port}/")
 
     logger.info("Press Ctrl + C to stop \n")
-    for process in process_pool:
-        process.join()
+
+    if max_requests and max_requests > 0 and len(process_pool) > 0:
+        while not shutting_down:
+            for i, process in enumerate(process_pool):
+                if not process.is_alive() and not shutting_down:
+                    process.join()
+                    logger.info("Worker process exited (recycling), spawning replacement.")
+                    copied_socket = socket.try_clone()
+                    new_process = Process(
+                        target=spawn_process,
+                        args=(
+                            directories,
+                            request_headers,
+                            routes,
+                            global_middlewares,
+                            route_middlewares,
+                            web_sockets,
+                            event_handlers,
+                            copied_socket,
+                            workers,
+                            response_headers,
+                            excluded_response_headers_paths,
+                            client_timeout,
+                            keep_alive_timeout,
+                            max_requests,
+                        ),
+                    )
+                    new_process.start()
+                    process_pool[i] = new_process
+            time.sleep(5)
+    else:
+        for process in process_pool:
+            process.join()
 
     return process_pool
 
@@ -84,9 +130,10 @@ def init_processpool(
     excluded_response_headers_paths: Optional[List[str]],
     client_timeout: int = 30,
     keep_alive_timeout: int = 20,
+    max_requests: Optional[int] = None,
 ) -> List[Process]:
     process_pool: List = []
-    if sys.platform.startswith("win32") or processes == 1:
+    if sys.platform.startswith("win32") or (processes == 1 and not max_requests):
         spawn_process(
             directories,
             request_headers,
@@ -101,11 +148,12 @@ def init_processpool(
             excluded_response_headers_paths,
             client_timeout,
             keep_alive_timeout,
+            max_requests,
         )
 
         return process_pool
 
-    for _ in range(processes):
+    for _ in range(max(1, processes)):
         copied_socket = socket.try_clone()
         process = Process(
             target=spawn_process,
@@ -123,6 +171,7 @@ def init_processpool(
                 excluded_response_headers_paths,
                 client_timeout,
                 keep_alive_timeout,
+                max_requests,
             ),
         )
         process.start()
@@ -161,6 +210,7 @@ def spawn_process(
     excluded_response_headers_paths: Optional[List[str]],
     client_timeout: int = 30,
     keep_alive_timeout: int = 20,
+    max_requests: Optional[int] = None,
 ):
     """
     This function is called by the main process handler to create a server runtime.
@@ -175,14 +225,13 @@ def spawn_process(
     :param socket SocketHeld: This is the main tcp socket, which is being shared across multiple processes.
     :param process_name string: This is the name given to the process to identify the process
     :param workers int: This is the name given to the process to identify the process
+    :param max_requests Optional[int]: Recycle this worker after N requests
     """
 
     loop = initialize_event_loop()
 
     server = Server()
 
-    # TODO: if we remove the dot access
-    # the startup time will improve in the server
     for directory in directories:
         server.add_directory(*directory.as_list())
 
@@ -218,8 +267,10 @@ def spawn_process(
         )
 
     try:
-        server.start(socket, workers)
+        server.start(socket, workers, max_requests)
         loop = asyncio.get_event_loop()
         loop.run_forever()
     except KeyboardInterrupt:
+        pass
+    finally:
         loop.close()

robyn/robyn.pyi

@@ -7,6 +7,10 @@ from typing import Callable, Optional, Union
 def get_version() -> str:
     pass
 
+def get_request_count() -> int:
+    """Returns the total number of HTTP requests handled by this worker process."""
+    pass
+
 class SocketHeld:
     def __init__(self, url: str, port: int):
         pass
@@ -514,7 +518,7 @@ class Server:
         use_channel: bool,
     ) -> None:
         pass
-    def start(self, socket: SocketHeld, workers: int, client_timeout: int, keep_alive_timeout: int) -> None:
+    def start(self, socket: SocketHeld, workers: int, max_requests: int | None = None) -> None:
         pass
 
 class WebSocketConnector:

src/lib.rs

@@ -34,10 +34,16 @@ fn get_version() -> String {
     env!("CARGO_PKG_VERSION").into()
 }
 
+#[pyfunction]
+fn get_request_count() -> u64 {
+    server::get_request_count()
+}
+
 #[pymodule]
 pub fn robyn(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     // the pymodule class/function to make the rustPyFunctions available
     m.add_function(wrap_pyfunction!(get_version, m)?)?;
+    m.add_function(wrap_pyfunction!(get_request_count, m)?)?;
 
     m.add_class::<Server>()?;
     m.add_class::<Headers>()?;

src/server.rs

@@ -27,7 +27,10 @@ use std::{env, thread};
 
 use actix_files::Files;
 use actix_http::KeepAlive;
+use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
 use actix_web::*;
+use futures_util::future::{ok, LocalBoxFuture, Ready};
+use std::task::{Context, Poll};
 
 use log::error;
 use once_cell::sync::OnceCell;
@@ -40,6 +43,55 @@ const MAX_PAYLOAD_SIZE: &str = "ROBYN_MAX_PAYLOAD_SIZE";
 const DEFAULT_MAX_PAYLOAD_SIZE: usize = 1_000_000; // 1Mb
 
 static STARTED: AtomicBool = AtomicBool::new(false);
+static REQUEST_COUNT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
+
+pub fn get_request_count() -> u64 {
+    REQUEST_COUNT.load(Relaxed)
+}
+
+struct RequestCounter;
+
+impl<S, B> Transform<S, ServiceRequest> for RequestCounter
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<B>;
+    type Error = Error;
+    type InitError = ();
+    type Transform = RequestCounterMiddleware<S>;
+    type Future = Ready<Result<Self::Transform, Self::InitError>>;
+
+    fn new_transform(&self, service: S) -> Self::Future {
+        ok(RequestCounterMiddleware { service })
+    }
+}
+
+struct RequestCounterMiddleware<S> {
+    service: S,
+}
+
+impl<S, B> Service<ServiceRequest> for RequestCounterMiddleware<S>
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<B>;
+    type Error = Error;
+    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.service.poll_ready(cx)
+    }
+
+    fn call(&self, req: ServiceRequest) -> Self::Future {
+        REQUEST_COUNT.fetch_add(1, Relaxed);
+        let fut = self.service.call(req);
+        Box::pin(fut)
+    }
+}
 
 #[derive(Clone)]
 struct Directory {
@@ -81,11 +133,16 @@ impl Server {
         }
     }
 
+    pub fn get_request_count(&self) -> u64 {
+        get_request_count()
+    }
+
     pub fn start(
         &mut self,
         _py: Python,
         socket: PyRef<SocketHeld>,
         workers: usize,
+        max_requests: Option<u64>,
     ) -> PyResult<()> {
         pyo3_log::init();
 
@@ -144,8 +201,8 @@ impl Server {
                     const_router.bake_global_headers(&global_response_headers);
                 }
 
-                HttpServer::new(move || {
-                    let mut app = App::new();
+                let server = HttpServer::new(move || {
+                    let mut app = App::new().wrap(RequestCounter);
 
                     let directories = directories.read().unwrap();
 
@@ -280,9 +337,33 @@ impl Server {
                 .client_request_timeout(std::time::Duration::from_secs(0))
                 .listen(raw_socket.into())
                 .unwrap()
-                .run()
-                .await
-                .unwrap();
+                .run();
+
+                let server_handle = server.handle();
+
+                if let Some(limit) = max_requests {
+                    let handle = server_handle.clone();
+                    actix_web::rt::spawn(async move {
+                        loop {
+                            actix_web::rt::time::sleep(std::time::Duration::from_secs(5)).await;
+                            if REQUEST_COUNT.load(Relaxed) >= limit {
+                                log::info!(
+                                    "Max requests ({}) reached, worker shutting down for recycling.",
+                                    limit
+                                );
+                                handle.stop(true).await;
+                                break;
+                            }
+                        }
+                    });
+                }
+
+                server.await.unwrap();
+
+                if max_requests.is_some() {
+                    log::info!("Actix server stopped, exiting worker process for recycling.");
+                    std::process::exit(0);
+                }
             });
         });