sova_core/
upgrade.rs

1//! HTTP/1 connection upgrade (WebSocket handshake, …).
2
3use crate::response::Response;
4use hyper::upgrade::{OnUpgrade as HyperOnUpgrade, Upgraded};
5use std::sync::Arc;
6use tokio::sync::{OwnedSemaphorePermit, Semaphore};
7
8/// Cap for concurrent upgraded connections (`App::max_upgraded_connections`).
9#[derive(Clone)]
10pub(crate) struct UpgradeBudget(pub Arc<Semaphore>);
11
12/// Pending upgrade extracted from the Hyper request (before `into_body`).
13pub(crate) struct PendingUpgrade {
14    pub(crate) on_upgrade: HyperOnUpgrade,
15    pub(crate) budget: UpgradeBudget,
16}
17
18/// Holds one slot in `max_upgraded_connections` until dropped.
19#[allow(dead_code)]
20pub struct UpgradePermit(OwnedSemaphorePermit);
21
22/// Handle to finish an HTTP upgrade (e.g. WebSocket).
23///
24/// Prefer ordinary routes + [`crate::Request::on_upgrade`] over [`crate::Router::raw`].
25pub struct OnUpgrade {
26    inner: HyperOnUpgrade,
27    permit: OwnedSemaphorePermit,
28}
29
30impl OnUpgrade {
31    /// Complete the protocol upgrade.
32    ///
33    /// Keep the returned [`UpgradePermit`] (or the whole tuple) alive for the
34    /// lifetime of the upgraded connection so the budget slot stays reserved.
35    pub async fn upgrade(self) -> Result<(Upgraded, UpgradePermit), hyper::Error> {
36        let io = self.inner.await?;
37        Ok((io, UpgradePermit(self.permit)))
38    }
39}
40
41/// Try to take the upgrade; on budget exhaustion returns **503** + `Retry-After`.
42pub(crate) fn take_upgrade(pending: PendingUpgrade) -> Result<OnUpgrade, Box<Response>> {
43    let permit = match pending.budget.0.clone().try_acquire_owned() {
44        Ok(p) => p,
45        Err(_) => {
46            return Err(Box::new(
47                Response::text("Service Unavailable")
48                    .status(503)
49                    .header("retry-after", "5"),
50            ));
51        }
52    };
53    Ok(OnUpgrade {
54        inner: pending.on_upgrade,
55        permit,
56    })
57}
sova_core/upgrade.rs

sova_core/
upgrade.rs
