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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
//! 待ち受けて、届いた接続を mon のルータへ渡す。
//!
//! axum::serve は、ヘッダを受け取り終えるまでの時間に上限を掛けない。
//! hyper の header_read_timeout にはタイマーが要るが、axum::serve はそれを渡していないためである。
//! 上限がないと、ヘッダを少しずつ送る相手や、つないだまま何も送らない相手に、接続を持たれ続ける。
//! ここではタイマーを渡し、REQUEST_TIMEOUT の秒数でヘッダの受け取りを打ち切る。
//!
//! 受け付けるのは HTTP/1 だけで、axum::serve の既定と同じである。
use axum::Router;
use axum::extract::ConnectInfo;
use axum::http::Request;
use crate::conf::Conf;
use hyper::body::Incoming;
use hyper::server::conn::http1;
use hyper_util::rt::{TokioIo, TokioTimer};
use hyper_util::service::TowerToHyperService;
use std::future::Future;
use std::io::ErrorKind;
use std::pin::pin;
use std::time::Duration;
use tokio::net::TcpListener;
use tokio::sync::watch;
use tower::ServiceExt;
/// 受け付けに失敗したとき、次を試すまで待つ時間。
/// ファイル記述子が尽きたときに、同じ失敗を休みなく繰り返さないようにする
const ACCEPT_RETRY: Duration = Duration::from_secs(1);
/// `listener` で待ち受け、届いた接続を `app` に渡す。`app` は [`crate::route::wrap`] の戻り値である。
///
/// ヘッダを受け取り終えるまでの上限には `conf.timeout`（REQUEST_TIMEOUT）を使う。
/// 相手のアドレスを `ConnectInfo<SocketAddr>` として入れるので、[`crate::ClientIp`] は TCP の相手を使える。
///
/// `shutdown` が終わると受け付けをやめ、処理中のリクエストを終えてから戻る。
pub async fn run<F>(listener: TcpListener, app: Router, conf: &Conf, shutdown: F) -> std::io::Result<()>
where
F: Future<Output = ()>,
{
let header_timeout = conf.timeout;
// 接続に、閉じるよう知らせる。送り手を落とすと、受け手の changed() が戻る
let (signal_sender, signal_receiver) = watch::channel(());
// 接続がすべて終わったかを知る。受け手が 1 つ残らず落ちると、送り手の closed() が戻る
let (close_sender, close_receiver) = watch::channel(());
let mut shutdown = pin!(shutdown);
loop {
let (stream, peer) = tokio::select! {
accepted = listener.accept() => match accepted {
Ok(pair) => pair,
Err(error) => {
wait_after(error).await;
continue;
}
},
_ = &mut shutdown => break,
};
let service = app.clone().map_request(move |mut request: Request<Incoming>| {
request.extensions_mut().insert(ConnectInfo(peer));
request
});
let mut signal = signal_receiver.clone();
let close = close_receiver.clone();
tokio::spawn(async move {
let mut builder = http1::Builder::new();
builder.timer(TokioTimer::new()).header_read_timeout(header_timeout);
let connection = builder
.serve_connection(TokioIo::new(stream), TowerToHyperService::new(service))
.with_upgrades();
let mut connection = pin!(connection);
tokio::select! {
result = connection.as_mut() => {
if let Err(error) = result {
tracing::debug!(error = ?error, "接続が途中で終わりました");
}
}
_ = signal.changed() => {
connection.as_mut().graceful_shutdown();
if let Err(error) = connection.as_mut().await {
tracing::debug!(error = ?error, "接続が途中で終わりました");
}
}
}
drop(close);
});
}
// 受け付けをやめ、接続に閉じるよう知らせてから、すべて終わるのを待つ
drop(listener);
drop(signal_sender);
drop(signal_receiver);
drop(close_receiver);
close_sender.closed().await;
Ok(())
}
/// 相手が途中で切っただけなら、すぐ次を受け付ける。それ以外は少し待つ
async fn wait_after(error: std::io::Error) {
let dropped = matches!(
error.kind(),
ErrorKind::ConnectionAborted | ErrorKind::ConnectionReset | ErrorKind::ConnectionRefused
);
if dropped {
return;
}
tracing::warn!(error = ?error, "接続を受け付けられません");
tokio::time::sleep(ACCEPT_RETRY).await;
}