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