perf(cache): consolidate HotSnapshot atomic refcounts into single Arc<FrozenState> (#149)

Author: ZhuchkaTriplesix

src/dispatch.rs

@@ -316,6 +316,7 @@ fn ensure_compiled_snapshot(state: &Arc<RwLock<AppState>>) -> Arc<CompiledRouter
     let mut st = state.write();
     if st.compiled.is_none() {
         st.compiled = Some(Arc::new(st.snapshot_routers()));
+        st.rebuild_snapshot();
     }
     Arc::clone(st.compiled.as_ref().expect("just populated"))
 }
@@ -472,20 +473,6 @@ pub async fn run_rsgi(
     let Some((method, path, query_string, is_head, snapshot)) = prelim else {
         return Ok(Python::with_gil(|py| py.None()));
     };
-    type CfgClones = (
-        Option<Py<PyAny>>,
-        Option<Py<PyAny>>,
-        Arc<Vec<Py<PyAny>>>,
-        Arc<Vec<Py<PyAny>>>,
-    );
-    let (cors_cfg, security_cfg, req_mw, res_mw): CfgClones = Python::with_gil(|_py| {
-        (
-            snapshot.cors.clone(),
-            snapshot.security_headers.clone(),
-            snapshot.request_middleware.clone(),
-            snapshot.response_middleware.clone(),
-        )
-    });
     if (method == "GET" || method == "HEAD") && path == "/openapi.json" && snapshot.include_openapi
     {
         let _ = Python::with_gil(|py| {
@@ -539,7 +526,7 @@ pub async fn run_rsgi(
             }
         }
     }
-    for mw in req_mw.iter() {
+    for mw in snapshot.request_middleware.iter() {
         let out: Py<PyAny> = match Python::with_gil(|py| {
             let f = mw.bind(py);
             f.call1((scope.bind(py), protocol.bind(py)))
@@ -556,10 +543,12 @@ pub async fn run_rsgi(
             return match Python::with_gil(|py| -> PyResult<()> {
                 let mut mapped = map_handler_return(py, &out)?;
 
-                if !res_mw.is_empty() && !matches!(mapped, HandlerMap::AlreadySent) {
+                if !snapshot.response_middleware.is_empty()
+                    && !matches!(mapped, HandlerMap::AlreadySent)
+                {
                     let response_module = py.import("oxyroute.response")?;
                     let response_class = response_module.getattr("Response")?;
-                    for res_m in res_mw.iter() {
+                    for res_m in snapshot.response_middleware.iter() {
                         let kwargs = pyo3::types::PyDict::new(py);
                         match &mapped {
                             HandlerMap::Simple {
@@ -605,16 +594,16 @@ pub async fn run_rsgi(
                     }
                 }
 
-                let mapped = if security_cfg.is_some() || cors_cfg.is_some() {
+                let mapped = if snapshot.security_headers.is_some() || snapshot.cors.is_some() {
                     let scope_bound = scope.bind(py).clone();
                     let mapped = merge_config_response_headers(
                         py,
-                        &security_cfg,
+                        &snapshot.security_headers,
                         scope_bound.clone(),
                         mapped,
                         true,
                     )?;
-                    merge_config_response_headers(py, &cors_cfg, scope_bound, mapped, false)?
+                    merge_config_response_headers(py, &snapshot.cors, scope_bound, mapped, false)?
                 } else {
                     mapped
                 };
@@ -1240,10 +1229,10 @@ pub async fn run_rsgi(
     match Python::with_gil(|py| -> PyResult<()> {
         let mut mapped = map_handler_return(py, &handler_out)?;
 
-        if !res_mw.is_empty() && !matches!(mapped, HandlerMap::AlreadySent) {
+        if !snapshot.response_middleware.is_empty() && !matches!(mapped, HandlerMap::AlreadySent) {
             let response_module = py.import("oxyroute.response")?;
             let response_class = response_module.getattr("Response")?;
-            for res_m in res_mw.iter() {
+            for res_m in snapshot.response_middleware.iter() {
                 let kwargs = pyo3::types::PyDict::new(py);
                 match &mapped {
                     HandlerMap::Simple {
@@ -1287,16 +1276,16 @@ pub async fn run_rsgi(
             }
         }
 
-        let mapped = if security_cfg.is_some() || cors_cfg.is_some() {
+        let mapped = if snapshot.security_headers.is_some() || snapshot.cors.is_some() {
             let scope_bound = scope.bind(py).clone();
             let mapped = merge_config_response_headers(
                 py,
-                &security_cfg,
+                &snapshot.security_headers,
                 scope_bound.clone(),
                 mapped,
                 true,
             )?;
-            merge_config_response_headers(py, &cors_cfg, scope_bound, mapped, false)?
+            merge_config_response_headers(py, &snapshot.cors, scope_bound, mapped, false)?
         } else {
             mapped
         };
@@ -1769,7 +1758,6 @@ async fn run_rsgi_websocket(
         Some(c) => c,
         None => ensure_compiled_snapshot(&state),
     };
-    let ws_routes = Arc::clone(&snapshot.websocket_routes);
     let route_match = match_ws_route_compiled(&compiled, &path);
     let Some((route_idx, params)) = route_match else {
         // No route → polite close. ``close`` is sync on RSGIWebsocketProtocol.
@@ -1781,7 +1769,8 @@ async fn run_rsgi_websocket(
         return Ok(Python::with_gil(|py| py.None()));
     };
     let (handler, is_async) = Python::with_gil(|_py| -> PyResult<(Py<PyAny>, bool)> {
-        let e = ws_routes
+        let e = snapshot
+            .websocket_routes
             .get(route_idx)
             .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("ws route index"))?;
         Ok((e.handler.clone(), e.is_async))

src/lib.rs

@@ -447,6 +447,7 @@ impl App {
         }
         // Keep auto-compiled routing snapshots fresh when routes are added before explicit freeze().
         st.compiled = None;
+        st.rebuild_snapshot();
         Ok(())
     }
 
@@ -480,6 +481,7 @@ impl App {
                 .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
         }
         st.compiled = None;
+        st.rebuild_snapshot();
         Ok(())
     }
 
@@ -490,12 +492,14 @@ impl App {
         if st.compiled.is_none() {
             st.compiled = Some(Arc::new(st.snapshot_routers()));
         }
+        st.rebuild_snapshot();
         Ok(())
     }
 
     fn set_openapi_served(&self, enabled: bool) -> PyResult<()> {
         let mut st = self.state.write();
         st.include_openapi = enabled;
+        st.rebuild_snapshot();
         Ok(())
     }
 
@@ -568,6 +572,7 @@ impl App {
         })
         .unwrap_or(false);
         Arc::make_mut(&mut st.exception_handlers).push((exc_type.unbind(), handler, is_async));
+        st.rebuild_snapshot();
         Ok(())
     }
 
@@ -578,6 +583,7 @@ impl App {
         } else {
             st.request_middleware = Arc::new(Vec::new());
         }
+        st.rebuild_snapshot();
         Ok(())
     }
 
@@ -596,13 +602,15 @@ impl App {
                 "phase must be 'request', 'response', or 'both'",
             ));
         }
+        st.rebuild_snapshot();
         Ok(())
     }
 
     /// Optional Python CORS config with ``response_header_pairs(scope)`` (see ``oxyroute.cors``).
     fn set_cors(&self, config: Option<Py<PyAny>>) -> PyResult<()> {
         let mut st = self.state.write();
         st.cors = config;
+        st.rebuild_snapshot();
         Ok(())
     }
 
@@ -611,6 +619,7 @@ impl App {
     fn set_security_headers(&self, config: Option<Py<PyAny>>) -> PyResult<()> {
         let mut st = self.state.write();
         st.security_headers = config;
+        st.rebuild_snapshot();
         Ok(())
     }
 
@@ -633,6 +642,7 @@ impl App {
                 })?;
             let mut st = state.write();
             st.db_pool = Some(pool);
+            st.rebuild_snapshot();
             Ok(())
         })
     }
@@ -643,7 +653,9 @@ impl App {
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             let pool = {
                 let mut st = state.write();
-                st.db_pool.take()
+                let p = st.db_pool.take();
+                st.rebuild_snapshot();
+                p
             };
             if let Some(p) = pool {
                 p.close().await;

src/state.rs

@@ -169,6 +169,7 @@ pub struct AppState {
     pub db_pool: Option<sqlx::PgPool>,
     /// Bitmask of allowed HTTP methods per registered path template.
     pub path_method_masks: Mutex<std::collections::HashMap<String, MethodMask>>,
+    pub snapshot: Arc<FrozenState>,
 }
 
 impl AppState {
@@ -178,9 +179,26 @@ impl AppState {
             "info": { "title": "OxyRoute", "version": "0.5.0" },
             "paths": {}
         });
+        let routes = Arc::new(Vec::new());
+        let websocket_routes = Arc::new(Vec::new());
+        let request_middleware = Arc::new(Vec::new());
+        let response_middleware = Arc::new(Vec::new());
+        let exception_handlers = Arc::new(Vec::new());
+        let snapshot = Arc::new(FrozenState {
+            routes: Arc::clone(&routes),
+            websocket_routes: Arc::clone(&websocket_routes),
+            compiled: None,
+            cors: None,
+            security_headers: None,
+            request_middleware: Arc::clone(&request_middleware),
+            response_middleware: Arc::clone(&response_middleware),
+            exception_handlers: Arc::clone(&exception_handlers),
+            include_openapi: true,
+            db_pool: None,
+        });
         Self {
-            routes: Arc::new(Vec::new()),
-            websocket_routes: Arc::new(Vec::new()),
+            routes,
+            websocket_routes,
             get: Mutex::new(Router::new()),
             post: Mutex::new(Router::new()),
             put: Mutex::new(Router::new()),
@@ -192,24 +210,19 @@ impl AppState {
             compiled: None,
             frozen: false,
             include_openapi: true,
-            request_middleware: Arc::new(Vec::new()),
-            response_middleware: Arc::new(Vec::new()),
-            exception_handlers: Arc::new(Vec::new()),
+            request_middleware,
+            response_middleware,
+            exception_handlers,
             cors: None,
             security_headers: None,
             db_pool: None,
             path_method_masks: Mutex::new(std::collections::HashMap::new()),
+            snapshot,
         }
     }
 
-    /// Cheap read-side snapshot of the fields the request hot path touches: the
-    /// returned [`HotSnapshot`] is built **inside one** `state.read()` so the request
-    /// dispatch can release the `RwLock` immediately and avoid further reads.
-    ///
-    /// Cheap because every cloned field is `Arc::clone` / `Option<Py<PyAny>>::clone`
-    /// (both refcount bumps), not deep clones.
-    pub fn hot_snapshot(&self) -> HotSnapshot {
-        HotSnapshot {
+    pub fn rebuild_snapshot(&mut self) {
+        self.snapshot = Arc::new(FrozenState {
             routes: Arc::clone(&self.routes),
             websocket_routes: Arc::clone(&self.websocket_routes),
             compiled: self.compiled.as_ref().map(Arc::clone),
@@ -220,7 +233,13 @@ impl AppState {
             exception_handlers: Arc::clone(&self.exception_handlers),
             include_openapi: self.include_openapi,
             db_pool: self.db_pool.clone(),
-        }
+        });
+    }
+
+    /// Read-side snapshot of the fields the request hot path touches: only 1 atomic
+    /// pointer clone (`Arc::clone(&self.snapshot)`).
+    pub fn hot_snapshot(&self) -> Arc<FrozenState> {
+        Arc::clone(&self.snapshot)
     }
 
     /// Clone current mutex-protected [`Router`]s into a snapshot (used at freeze / tests).
@@ -242,10 +261,9 @@ impl AppState {
     }
 }
 
-/// One-shot read-side view of [`AppState`] for [`run_rsgi`]. All fields are cheap to clone
-/// (`Arc`/`Option<Py<PyAny>>` refcount bumps) so the hot path can drop the `RwLock` after a
-/// single `read()`. See [`AppState::hot_snapshot`].
-pub struct HotSnapshot {
+/// Consolidated immutable read-side view of [`AppState`] for request dispatching.
+/// Only 1 atomic refcount increment is needed per request.
+pub struct FrozenState {
     pub routes: Arc<Vec<RouteEntry>>,
     pub websocket_routes: Arc<Vec<WebsocketRoute>>,
     pub compiled: Option<Arc<CompiledRouters>>,
@@ -258,6 +276,8 @@ pub struct HotSnapshot {
     pub db_pool: Option<sqlx::PgPool>,
 }
 
+pub type HotSnapshot = Arc<FrozenState>;
+
 /// Lookup a WebSocket route in a precomputed [`CompiledRouters`] (lock-free).
 pub fn match_ws_route_compiled<'a, 'b>(
     compiled: &'a CompiledRouters,