tuwunel_router/
request.rs

1use std::{
2	convert::Infallible,
3	fmt::Debug,
4	sync::{Arc, atomic::Ordering},
5	time::Duration,
6};
7
8use axum::{
9	extract::Request,
10	response::{IntoResponse, Response},
11};
12use futures::FutureExt;
13use http::{Method, StatusCode, Uri};
14use ruma::api::error::ErrorKind;
15use tokio::{sync::Notify, task, time::sleep};
16use tower::{Service, ServiceExt};
17use tracing::Span;
18use tuwunel_core::{Error, Result, debug, debug_error, debug_warn, defer, error, trace};
19use tuwunel_service::Services;
20
21#[tracing::instrument(
22	name = "request",
23	level = "debug",
24	skip_all,
25	err(Debug, level = "debug")
26	fields(
27		task = %task::id(),
28		id = %services
29			.server
30			.metrics
31			.requests_count
32			.fetch_add(1, Ordering::Relaxed)
33	)
34)]
35pub(crate) async fn handle<S>(
36	services: Arc<Services>,
37	req: Request,
38	inner: S,
39) -> Result<Response, StatusCode>
40where
41	S: Service<Request, Error = Infallible> + Send + 'static,
42	S::Response: IntoResponse,
43	S::Future: Send + 'static,
44{
45	if !services.server.is_running() {
46		debug_warn!(
47			method = %req.method(),
48			uri = %req.uri(),
49			"unavailable pending shutdown"
50		);
51
52		return Err(StatusCode::SERVICE_UNAVAILABLE);
53	}
54
55	let uri = req.uri().clone();
56	let method = req.method().clone();
57	let parent = Span::current();
58	let response = match method {
59		| Method::PUT | Method::POST | Method::DELETE | Method::PATCH =>
60			spawn_execute(services, req, inner, parent).await?,
61		| _ => execute(&services, req, inner, &parent).await,
62	};
63
64	handle_result(&method, &uri, response)
65}
66
67async fn spawn_execute<S>(
68	services: Arc<Services>,
69	mut req: Request,
70	inner: S,
71	parent: Span,
72) -> Result<Response, StatusCode>
73where
74	S: Service<Request, Error = Infallible> + Send + 'static,
75	S::Response: IntoResponse,
76	S::Future: Send + 'static,
77{
78	let detached = Arc::new(Notify::new());
79	req.extensions_mut().insert(detached.clone());
80
81	let task = services
82		.clone()
83		.server
84		.runtime()
85		.spawn(async move {
86			tokio::select! {
87				response = execute(&services, req, inner, &parent) => response,
88				response = services.server.until_shutdown()
89					.then(|()| {
90						let timeout = services.config.client_shutdown_timeout;
91						sleep(Duration::from_secs(timeout))
92					})
93					.map(|()| StatusCode::SERVICE_UNAVAILABLE)
94					.map(IntoResponse::into_response) => response,
95			}
96		});
97
98	let abort = task.abort_handle();
99	defer! {{
100		if !abort.is_finished() {
101			debug_warn!(
102				task = ?abort.id(),
103				"Client disconnected; detached request."
104			);
105
106			detached.notify_one();
107		}
108	}};
109
110	task.await.map_err(unhandled)
111}
112
113#[tracing::instrument(
114	name = "handle",
115	level = "debug",
116	parent = parent,
117	skip_all,
118	ret(level = "trace"),
119	fields(
120		task = %task::id(),
121	)
122)]
123#[cfg_attr(not(debug_assertions), expect(unused_variables))]
124async fn execute<S>(
125	// we made a safety contract that Services will not go out of scope
126	// during the request; this ensures a reference is accounted for at
127	// the base frame of the task regardless of its detachment.
128	services: &Arc<Services>,
129	req: Request,
130	inner: S,
131	parent: &Span,
132) -> Response
133where
134	S: Service<Request, Error = Infallible>,
135	S::Response: IntoResponse,
136{
137	#[cfg(debug_assertions)]
138	services
139		.server
140		.metrics
141		.requests_handle_active
142		.fetch_add(1, Ordering::Relaxed);
143
144	#[cfg(debug_assertions)]
145	defer! {{
146		_ = services.server
147			.metrics
148			.requests_handle_finished
149			.fetch_add(1, Ordering::Relaxed);
150		_ = services.server
151			.metrics
152			.requests_handle_active
153			.fetch_sub(1, Ordering::Relaxed);
154	}};
155
156	inner
157		.oneshot(req)
158		.map(IntoResponse::into_response)
159		.await
160}
161
162fn handle_result(method: &Method, uri: &Uri, result: Response) -> Result<Response, StatusCode> {
163	let status = result.status();
164	let code = status.as_u16();
165	let reason = status
166		.canonical_reason()
167		.unwrap_or("Unknown Reason");
168
169	if status.is_server_error() {
170		error!(method = ?method, uri = ?uri, "{code} {reason}");
171	} else if status.is_client_error() {
172		debug_error!(method = ?method, uri = ?uri, "{code} {reason}");
173	} else if status.is_redirection() {
174		debug!(method = ?method, uri = ?uri, "{code} {reason}");
175	} else {
176		trace!(method = ?method, uri = ?uri, "{code} {reason}");
177	}
178
179	if status == StatusCode::METHOD_NOT_ALLOWED {
180		return Ok(Error::Request(
181			ErrorKind::Unrecognized,
182			"Method Not Allowed".into(),
183			StatusCode::METHOD_NOT_ALLOWED,
184		)
185		.into_response());
186	}
187
188	Ok(result)
189}
190
191#[cold]
192fn unhandled<Error: Debug>(e: Error) -> StatusCode {
193	error!("unhandled error or panic during request: {e:?}");
194
195	StatusCode::INTERNAL_SERVER_ERROR
196}
tuwunel_router/request.rs

tuwunel_router/
request.rs
