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

@@ -84,6 +84,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=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 +107,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

@@ -8,7 +8,7 @@
 
 from robyn.events import Events
 from robyn.logger import logger
-from robyn.robyn import FunctionInfo, Headers, Server, SocketHeld
+from robyn.robyn import FunctionInfo, Headers, Server, SocketHeld, get_request_count
 from robyn.router import GlobalMiddleware, Route, RouteMiddleware
 from robyn.types import Directory
 
@@ -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,38 @@ 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:
+                    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,6 +129,7 @@ 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:
@@ -101,6 +147,7 @@ def init_processpool(
             excluded_response_headers_paths,
             client_timeout,
             keep_alive_timeout,
+            max_requests,
         )
 
         return process_pool
@@ -123,6 +170,7 @@ def init_processpool(
                 excluded_response_headers_paths,
                 client_timeout,
                 keep_alive_timeout,
+                max_requests,
             ),
         )
         process.start()
@@ -161,6 +209,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 +224,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())
 
@@ -220,6 +268,20 @@ def spawn_process(
     try:
         server.start(socket, workers)
         loop = asyncio.get_event_loop()
+
+        if max_requests and max_requests > 0:
+
+            def _check_max_requests():
+                if get_request_count() >= max_requests:
+                    logger.info("Max requests (%d) reached, worker shutting down for recycling.", max_requests)
+                    loop.stop()
+                else:
+                    loop.call_later(5, _check_max_requests)
+
+            loop.call_later(5, _check_max_requests)
+
         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

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

@@ -40,6 +40,11 @@ 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)
+}
 
 #[derive(Clone)]
 struct Directory {
@@ -482,6 +487,8 @@ async fn index(
     excluded_response_headers_paths: web::Data<Option<Vec<String>>>,
     req: HttpRequest,
 ) -> ResponseType {
+    REQUEST_COUNT.fetch_add(1, Relaxed);
+
     if !HttpMethod::is_supported(req.method()) {
         return ResponseType::Standard(Response::method_not_allowed(None));
     }