fix(websocket): configurable WS session cap (WS-09)
This commit is contained in:
@@ -88,6 +88,8 @@ pub struct HttpAdapter {
|
||||
router: Router,
|
||||
openapi_doc: CachedOpenAPIDoc,
|
||||
ws_sessions: Arc<crate::websocket::WsSessions>,
|
||||
ws_max_sessions: usize,
|
||||
ws_session_slots: Arc<tokio::sync::Semaphore>,
|
||||
}
|
||||
|
||||
impl HttpAdapter {
|
||||
@@ -113,12 +115,15 @@ impl HttpAdapter {
|
||||
let decoy = DecoyConfig::default();
|
||||
let openapi_doc = CachedOpenAPIDoc::new(®istry);
|
||||
let ws_sessions = Arc::new(crate::websocket::WsSessions::new());
|
||||
let ws_max_sessions = crate::websocket::DEFAULT_WS_MAX_SESSIONS;
|
||||
let ws_session_slots = Arc::new(tokio::sync::Semaphore::new(ws_max_sessions));
|
||||
let state = RouterState {
|
||||
registry: Arc::clone(®istry),
|
||||
identity_provider: Arc::clone(&identity_provider),
|
||||
decoy: decoy.clone(),
|
||||
openapi_doc: openapi_doc.clone(),
|
||||
ws_sessions: Arc::clone(&ws_sessions),
|
||||
ws_session_slots: Arc::clone(&ws_session_slots),
|
||||
};
|
||||
let router = build_router(state, None);
|
||||
Self {
|
||||
@@ -130,6 +135,8 @@ impl HttpAdapter {
|
||||
router,
|
||||
openapi_doc,
|
||||
ws_sessions,
|
||||
ws_max_sessions,
|
||||
ws_session_slots,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -141,6 +148,7 @@ impl HttpAdapter {
|
||||
decoy,
|
||||
openapi_doc: self.openapi_doc.clone(),
|
||||
ws_sessions: Arc::clone(&self.ws_sessions),
|
||||
ws_session_slots: Arc::clone(&self.ws_session_slots),
|
||||
};
|
||||
// `extra_routes` is borrowed, not consumed (SRV-05): a builder
|
||||
// call after `with_extra_routes` must keep the custom routes in
|
||||
@@ -157,12 +165,34 @@ impl HttpAdapter {
|
||||
decoy: self.decoy.clone(),
|
||||
openapi_doc: self.openapi_doc.clone(),
|
||||
ws_sessions: Arc::clone(&self.ws_sessions),
|
||||
ws_session_slots: Arc::clone(&self.ws_session_slots),
|
||||
};
|
||||
self.router = build_router(state, Some(routes.clone()));
|
||||
self.extra_routes = Some(routes);
|
||||
self
|
||||
}
|
||||
|
||||
/// The concurrent WS session cap (WS-09): the upgrade handler
|
||||
/// acquires one semaphore permit per upgrade, post-auth and
|
||||
/// pre-upgrade; a caller over the configured cap is rejected with
|
||||
/// **503 Service Unavailable**, and the permit is held for the
|
||||
/// session's lifetime (an ended session frees its slot).
|
||||
///
|
||||
/// Default: [`crate::websocket::DEFAULT_WS_MAX_SESSIONS`] (64).
|
||||
pub fn with_ws_max_sessions(mut self, max_sessions: usize) -> Self {
|
||||
self.ws_max_sessions = max_sessions;
|
||||
let state = RouterState {
|
||||
registry: Arc::clone(&self.registry),
|
||||
identity_provider: Arc::clone(&self.identity_provider),
|
||||
decoy: self.decoy.clone(),
|
||||
openapi_doc: self.openapi_doc.clone(),
|
||||
ws_sessions: Arc::clone(&self.ws_sessions),
|
||||
ws_session_slots: Self::rebuild_session_slots(max_sessions),
|
||||
};
|
||||
self.router = build_router(state, self.extra_routes.clone());
|
||||
self
|
||||
}
|
||||
|
||||
/// The shared WS session registry (WS-08): live sessions'
|
||||
/// [`WsPumps`](crate::websocket::WsPumps) handles, evictable via
|
||||
/// `WsSessions::abort`. The upgrade handler registers against this
|
||||
@@ -171,6 +201,12 @@ impl HttpAdapter {
|
||||
Arc::clone(&self.ws_sessions)
|
||||
}
|
||||
|
||||
/// A fresh semaphore for the new cap; the retained handles of
|
||||
/// already-open sessions are unaffected (they hold their permits).
|
||||
fn rebuild_session_slots(max_sessions: usize) -> Arc<tokio::sync::Semaphore> {
|
||||
Arc::new(tokio::sync::Semaphore::new(max_sessions))
|
||||
}
|
||||
|
||||
pub fn decoy(&self) -> &DecoyConfig {
|
||||
&self.decoy
|
||||
}
|
||||
@@ -956,6 +992,9 @@ mod tests {
|
||||
decoy: DecoyConfig::default(),
|
||||
openapi_doc: CachedOpenAPIDoc::new(&OperationRegistry::new()),
|
||||
ws_sessions: Arc::new(crate::websocket::WsSessions::new()),
|
||||
ws_session_slots: Arc::new(tokio::sync::Semaphore::new(
|
||||
crate::websocket::DEFAULT_WS_MAX_SESSIONS,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+11
-2
@@ -26,8 +26,9 @@ pub enum DecoyConfig {
|
||||
/// State embedded in the axum `Router`: the registry and identity
|
||||
/// provider every request handler reaches through the router state, plus
|
||||
/// the decoy config for the fallback, the pre-serialized
|
||||
/// `/openapi.json` projection cache (SRV-09), and the shared WS session
|
||||
/// registry the WS upgrade retains pump handles in (WS-08).
|
||||
/// `/openapi.json` projection cache (SRV-09), and the WS session
|
||||
/// lifecycle config the upgrade handler enforces: the shared pump-handle
|
||||
/// registry (WS-08) and the concurrent-session cap (WS-09).
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct RouterState {
|
||||
pub(crate) registry: Arc<OperationRegistry>,
|
||||
@@ -35,6 +36,10 @@ pub(crate) struct RouterState {
|
||||
pub(crate) decoy: DecoyConfig,
|
||||
pub(crate) openapi_doc: crate::server::adapter::CachedOpenAPIDoc,
|
||||
pub(crate) ws_sessions: Arc<crate::websocket::WsSessions>,
|
||||
/// Session cap (WS-09): one `HttpAdapter`-wide semaphore, built
|
||||
/// once at construction — the upgrade handler acquires one permit
|
||||
/// per upgrade and holds it for the session's lifetime.
|
||||
pub(crate) ws_session_slots: Arc<tokio::sync::Semaphore>,
|
||||
}
|
||||
|
||||
impl axum::extract::FromRef<RouterState> for crate::websocket::SessionState {
|
||||
@@ -42,6 +47,7 @@ impl axum::extract::FromRef<RouterState> for crate::websocket::SessionState {
|
||||
crate::websocket::SessionState::new(
|
||||
Arc::clone(&state.registry),
|
||||
Arc::clone(&state.ws_sessions),
|
||||
Arc::clone(&state.ws_session_slots),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -89,6 +95,9 @@ mod tests {
|
||||
},
|
||||
openapi_doc: crate::server::adapter::CachedOpenAPIDoc::new(&OperationRegistry::new()),
|
||||
ws_sessions: Arc::new(crate::websocket::WsSessions::new()),
|
||||
ws_session_slots: Arc::new(tokio::sync::Semaphore::new(
|
||||
crate::websocket::DEFAULT_WS_MAX_SESSIONS,
|
||||
)),
|
||||
};
|
||||
let extracted: DecoyConfig = axum::extract::FromRef::from_ref(&state);
|
||||
assert!(matches!(extracted, DecoyConfig::Redirect { .. }));
|
||||
|
||||
Reference in New Issue
Block a user