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
126
127
128
129
130
131
132
133
134
135
//! The `runtime:` section of the config file: the async runtime the binary
//! runs the proxy on. Only the binary reads it; the library runs on whatever
//! runtime its embedder provides.
use std::ffi::OsStr;
use std::fmt;
use std::num::NonZeroUsize;
use serde::Deserialize;
/// The variable tokio reads for its default worker count.
const WORKER_THREADS_ENV: &str = "TOKIO_WORKER_THREADS";
/// Threads the runtime may add for blocking work (tokio's default), set
/// explicitly because tokio adds it to the worker count unchecked.
const MAX_BLOCKING_THREADS: usize = 512;
/// The keys of the config file this binary reads itself; the library reads
/// the rest and ignores these.
#[derive(Debug, Default, Deserialize)]
pub struct FileConfig {
#[serde(default)]
pub runtime: RuntimeConfig,
}
/// `runtime:`.
#[derive(Debug, Default, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RuntimeConfig {
/// Worker threads of the async runtime, so the CPU cores the proxy keeps
/// busy. Left out: tokio's default. Present, it must be a positive
/// integer: an explicit null or empty value is refused, not taken for
/// "left out".
#[serde(default, deserialize_with = "present_count")]
pub worker_threads: Option<NonZeroUsize>,
}
/// A key that is present holds a count; only an absent one (`default`) is
/// `None`.
fn present_count<'de, D: serde::Deserializer<'de>>(
de: D,
) -> Result<Option<NonZeroUsize>, D::Error> {
NonZeroUsize::deserialize(de).map(Some)
}
/// Where the worker count came from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WorkerSource {
/// `runtime.worker_threads`.
Config,
/// The `TOKIO_WORKER_THREADS` environment variable.
Environment,
/// The available parallelism (on Linux, the cgroup CPU quota).
AvailableParallelism,
}
impl fmt::Display for WorkerSource {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(match self {
Self::Config => "runtime.worker_threads",
Self::Environment => WORKER_THREADS_ENV,
Self::AvailableParallelism => "available parallelism",
})
}
}
impl FileConfig {
/// The binary's own settings in a config file.
///
/// # Errors
///
/// A `runtime:` section that is not a map, holds an unknown key, or sets
/// `worker_threads` to anything but a positive integer; the message names
/// the key.
pub fn from_yaml_str(yaml: &str) -> Result<Self, serde_yaml::Error> {
serde_yaml::from_str(yaml)
}
}
impl RuntimeConfig {
/// A multi-thread runtime with the configured worker count, or tokio's
/// default, and where that count came from.
///
/// # Errors
///
/// The runtime cannot be created, `TOKIO_WORKER_THREADS` is used and is
/// not a positive integer, or the worker count is so large that tokio's
/// thread limit overflows: tokio would panic on either, so they are
/// refused here with the name of the setting.
pub fn build(&self) -> std::io::Result<(tokio::runtime::Runtime, WorkerSource)> {
self.build_with(std::env::var_os(WORKER_THREADS_ENV).as_deref())
}
/// [`build`](Self::build) with the value of `TOKIO_WORKER_THREADS` given
/// rather than read, so the count depends on the arguments alone.
fn build_with(
&self,
env: Option<&OsStr>,
) -> std::io::Result<(tokio::runtime::Runtime, WorkerSource)> {
let (workers, source) = match (self.worker_threads, env) {
(Some(n), _) => (n.get(), WorkerSource::Config),
(None, Some(raw)) => {
let n = raw
.to_str()
.and_then(|s| s.trim().parse::<NonZeroUsize>().ok())
.ok_or_else(|| {
std::io::Error::other(format!(
"{WORKER_THREADS_ENV} must be a positive integer, got {raw:?}"
))
})?;
(n.get(), WorkerSource::Environment)
}
// What tokio falls back to itself; set here so tokio does not
// read the variable a second time.
(None, None) => (
std::thread::available_parallelism().map_or(1, NonZeroUsize::get),
WorkerSource::AvailableParallelism,
),
};
if workers.checked_add(MAX_BLOCKING_THREADS).is_none() {
return Err(std::io::Error::other(format!(
"{source} is too large: {workers} worker threads"
)));
}
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(workers)
.max_blocking_threads(MAX_BLOCKING_THREADS)
.enable_all()
.build()?;
Ok((runtime, source))
}
}
#[cfg(test)]
mod tests;