1use tokio_util::sync::CancellationToken;
6
7const NOFILE_CEILING: nix::libc::rlim_t = 4096;
16
17const NOFILE_FALLBACKS: [nix::libc::rlim_t; 3] = [2048, 1024, 512];
19
20fn nofile_candidates(
29 soft: nix::libc::rlim_t,
30 hard: nix::libc::rlim_t,
31 ceiling: nix::libc::rlim_t,
32) -> Vec<nix::libc::rlim_t> {
33 let top = hard.min(ceiling);
34 let mut candidates = vec![top];
35 candidates.extend(NOFILE_FALLBACKS.iter().copied().filter(|&c| c < top));
36 candidates.retain(|&c| c > soft);
37 candidates
38}
39
40pub fn raise_nofile_limit() {
48 raise_nofile_limit_with(read_nofile_limits, write_nofile_limits, NOFILE_CEILING);
49}
50
51fn read_nofile_limits() -> nix::Result<(nix::libc::rlim_t, nix::libc::rlim_t)> {
53 nix::sys::resource::getrlimit(nix::sys::resource::Resource::RLIMIT_NOFILE)
54}
55
56fn write_nofile_limits(soft: nix::libc::rlim_t, hard: nix::libc::rlim_t) -> nix::Result<()> {
58 nix::sys::resource::setrlimit(nix::sys::resource::Resource::RLIMIT_NOFILE, soft, hard)
59}
60
61#[derive(Debug, PartialEq, Eq)]
63enum NofileOutcome {
64 Unreadable,
66 Adequate,
68 Raised(nix::libc::rlim_t),
70 Refused,
72}
73
74fn raise_nofile_limit_with(
77 get: impl FnOnce() -> nix::Result<(nix::libc::rlim_t, nix::libc::rlim_t)>,
78 mut set: impl FnMut(nix::libc::rlim_t, nix::libc::rlim_t) -> nix::Result<()>,
79 ceiling: nix::libc::rlim_t,
80) -> NofileOutcome {
81 let (soft, hard) = match get() {
82 Ok(limits) => limits,
83 Err(e) => {
84 tracing::warn!("could not read the open-file limit: {e}");
85 return NofileOutcome::Unreadable;
86 }
87 };
88 let candidates = nofile_candidates(soft, hard, ceiling);
89 if candidates.is_empty() {
90 tracing::debug!("open-file limit {soft} (hard {hard}) needs no raising; leaving it");
91 return NofileOutcome::Adequate;
92 }
93 for target in candidates {
94 match set(target, hard) {
95 Ok(()) => {
96 tracing::info!("raised the open-file limit from {soft} to {target}");
97 return NofileOutcome::Raised(target);
98 }
99 Err(e) => {
100 tracing::warn!("could not raise the open-file limit from {soft} to {target}: {e}");
101 }
102 }
103 }
104 NofileOutcome::Refused
105}
106
107pub fn install_signal_handlers(shutdown: CancellationToken) {
116 #[cfg(unix)]
117 {
118 use tokio::signal::unix::{signal, SignalKind};
119 tokio::spawn(async move {
120 let mut term = match signal(SignalKind::terminate()) {
121 Ok(s) => s,
122 Err(e) => {
123 tracing::warn!("failed to install SIGTERM handler: {e}");
124 return;
125 }
126 };
127 let mut interrupt = match signal(SignalKind::interrupt()) {
128 Ok(s) => s,
129 Err(e) => {
130 tracing::warn!("failed to install SIGINT handler: {e}");
131 return;
132 }
133 };
134 let mut hangup = match signal(SignalKind::hangup()) {
135 Ok(s) => s,
136 Err(e) => {
137 tracing::warn!("failed to install SIGHUP handler: {e}");
138 return;
139 }
140 };
141 tokio::select! {
142 _ = term.recv() => tracing::info!("received SIGTERM; shutting down"),
143 _ = interrupt.recv() => tracing::info!("received SIGINT; shutting down"),
144 _ = hangup.recv() => tracing::info!("received SIGHUP; shutting down"),
145 }
146 shutdown.cancel();
147 });
148 }
149 #[cfg(not(unix))]
150 {
151 tokio::spawn(async move {
152 if tokio::signal::ctrl_c().await.is_ok() {
153 tracing::info!("received Ctrl-C; shutting down");
154 shutdown.cancel();
155 }
156 });
157 }
158}
159
160#[cfg(test)]
161#[allow(clippy::unwrap_used, clippy::expect_used)]
162mod tests {
163 use super::*;
164
165 #[test]
166 fn candidates_lead_with_the_ceiling_then_fall_back_to_smaller_values() {
167 assert_eq!(
169 nofile_candidates(256, nix::libc::RLIM_INFINITY, 4096),
170 vec![4096, 2048, 1024, 512]
171 );
172 }
173
174 #[test]
175 fn candidates_never_exceed_the_hard_limit() {
176 assert_eq!(nofile_candidates(256, 1024, 4096), vec![1024, 512]);
177 assert_eq!(nofile_candidates(256, 300, 4096), vec![300]);
178 }
179
180 #[test]
181 fn candidates_never_lower_or_rewrite_an_adequate_limit() {
182 assert!(nofile_candidates(4096, nix::libc::RLIM_INFINITY, 4096).is_empty());
183 assert!(nofile_candidates(65536, nix::libc::RLIM_INFINITY, 4096).is_empty());
184 assert!(nofile_candidates(1024, 1024, 4096).is_empty());
186 }
187
188 #[test]
189 fn candidates_skip_values_at_or_below_the_current_soft_limit() {
190 assert_eq!(
191 nofile_candidates(1500, nix::libc::RLIM_INFINITY, 4096),
192 vec![4096, 2048]
193 );
194 }
195
196 fn raise(
200 limits: nix::Result<(nix::libc::rlim_t, nix::libc::rlim_t)>,
201 refuse: impl Fn(nix::libc::rlim_t) -> bool,
202 ) -> (NofileOutcome, Vec<nix::libc::rlim_t>) {
203 let mut attempts = Vec::new();
204 let outcome = raise_nofile_limit_with(
205 || limits,
206 |target, _hard| {
207 attempts.push(target);
208 if refuse(target) {
209 Err(nix::errno::Errno::EINVAL)
210 } else {
211 Ok(())
212 }
213 },
214 4096,
215 );
216 (outcome, attempts)
217 }
218
219 #[test]
220 fn an_unreadable_limit_is_left_alone() {
221 let (outcome, attempts) = raise(Err(nix::errno::Errno::EPERM), |_| false);
222 assert_eq!(outcome, NofileOutcome::Unreadable);
223 assert!(attempts.is_empty());
224 }
225
226 #[test]
227 fn an_adequate_limit_is_not_touched() {
228 let (outcome, attempts) = raise(Ok((8192, nix::libc::RLIM_INFINITY)), |_| false);
229 assert_eq!(outcome, NofileOutcome::Adequate);
230 assert!(attempts.is_empty());
231 }
232
233 #[test]
234 fn the_limit_is_raised_to_the_ceiling_when_the_host_allows_it() {
235 let (outcome, attempts) = raise(Ok((256, nix::libc::RLIM_INFINITY)), |_| false);
236 assert_eq!(outcome, NofileOutcome::Raised(4096));
237 assert_eq!(attempts, vec![4096], "no further attempt after a success");
238 }
239
240 #[test]
241 fn a_refused_target_falls_back_to_the_next_smaller_one() {
242 let (outcome, attempts) = raise(Ok((256, nix::libc::RLIM_INFINITY)), |t| t > 1024);
244 assert_eq!(outcome, NofileOutcome::Raised(1024));
245 assert_eq!(attempts, vec![4096, 2048, 1024]);
246 }
247
248 #[test]
249 fn a_host_that_refuses_every_target_leaves_the_limit_as_it_was() {
250 let (outcome, attempts) = raise(Ok((256, nix::libc::RLIM_INFINITY)), |_| true);
251 assert_eq!(outcome, NofileOutcome::Refused);
252 assert_eq!(attempts, vec![4096, 2048, 1024, 512]);
253 }
254
255 #[test]
259 fn the_real_limit_accessors_round_trip() {
260 let (soft, hard) = read_nofile_limits().unwrap();
261
262 write_nofile_limits(soft, hard).unwrap();
263
264 assert_eq!(read_nofile_limits().unwrap(), (soft, hard));
265 }
266
267 #[test]
268 fn raise_nofile_limit_leaves_the_soft_limit_no_lower_than_it_found_it() {
269 use nix::sys::resource::{getrlimit, Resource};
270
271 let (before, _) = getrlimit(Resource::RLIMIT_NOFILE).unwrap();
272 raise_nofile_limit();
273 let (after, hard) = getrlimit(Resource::RLIMIT_NOFILE).unwrap();
274 assert!(after >= before, "soft limit fell from {before} to {after}");
275 assert!(after <= hard);
276 }
277}