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(
43    pending: PendingUpgrade,
44) -> Result<OnUpgrade, Box<Response>> {
45    let permit = match pending.budget.0.clone().try_acquire_owned() {
46        Ok(p) => p,
47        Err(_) => {
48            return Err(Box::new(
49                Response::text("Service Unavailable")
50                    .status(503)
51                    .header("retry-after", "5"),
52            ));
53        }
54    };
55    Ok(OnUpgrade {
56        inner: pending.on_upgrade,
57        permit,
58    })
59}
sova_core/upgrade.rs

sova_core/
upgrade.rs
