Skip to main content

omni_dev/daemon/
lifecycle.rs

1//! Process-lifecycle wiring: translate OS termination signals into a graceful
2//! shutdown of the daemon's [`CancellationToken`], and lift the descriptor limit
3//! the service manager starts it with.
4
5use tokio_util::sync::CancellationToken;
6
7/// The most descriptors the daemon asks for (#2111).
8///
9/// Far above the daemon's real need — two push streams per VS Code window plus
10/// the requests in flight — yet below macOS's per-process cap (`OPEN_MAX`, 10240),
11/// above which `setrlimit` fails outright. It stays modest on purpose: the limit
12/// is inherited by every `gh`/`git`/`code` child the daemon spawns, and a very
13/// large one makes close-all-descriptors loops and `select`-based tools slow or
14/// wrong.
15const NOFILE_CEILING: nix::libc::rlim_t = 4096;
16
17/// Smaller soft limits to fall back to when the first choice is refused.
18const NOFILE_FALLBACKS: [nix::libc::rlim_t; 3] = [2048, 1024, 512];
19
20/// The soft `RLIMIT_NOFILE` values to try raising to, best first; empty when the
21/// limit should be left alone.
22///
23/// Aims for the lower of the hard limit and `ceiling`, then the
24/// [`NOFILE_FALLBACKS`] below it — some hosts cap the soft limit under the hard one
25/// (macOS's `kern.maxfilesperproc`), and a partial raise still beats none. A soft
26/// limit already at or above a candidate never gets it, so nothing is ever lowered:
27/// an operator who set a higher one meant it.
28fn 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
40/// Raises the soft `RLIMIT_NOFILE` toward the hard limit, up to `NOFILE_CEILING`
41/// (4096), settling for a smaller value if the host refuses that one.
42///
43/// launchd starts the daemon with macOS's default soft limit of 256, which a busy
44/// day of open windows and in-flight requests crosses — and once `accept` fails
45/// with `EMFILE` the daemon cannot hear heartbeats (#2111). Best-effort: a failure
46/// is logged and the daemon carries on with the limit it has.
47pub fn raise_nofile_limit() {
48    raise_nofile_limit_with(read_nofile_limits, write_nofile_limits, NOFILE_CEILING);
49}
50
51/// The process's current `(soft, hard)` `RLIMIT_NOFILE`.
52fn 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
56/// Sets the process's `RLIMIT_NOFILE` to `(soft, hard)`.
57fn 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/// What [`raise_nofile_limit_with`] did.
62#[derive(Debug, PartialEq, Eq)]
63enum NofileOutcome {
64    /// The current limits could not be read, so nothing was changed.
65    Unreadable,
66    /// The soft limit already met every candidate, so nothing was changed.
67    Adequate,
68    /// The soft limit was raised to this value.
69    Raised(nix::libc::rlim_t),
70    /// The host refused every candidate; the soft limit is as it was.
71    Refused,
72}
73
74/// [`raise_nofile_limit`] over injected limit accessors, so the fallback ladder
75/// can be driven without touching the test process's real limit.
76fn 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
107/// Spawns a task that cancels `shutdown` when the process is asked to stop.
108///
109/// On Unix this listens for `SIGTERM` (what `launchctl bootout` and service
110/// managers send), `SIGINT` (Ctrl-C in a foreground `daemon run`), and
111/// `SIGHUP` (the default disposition would hard-kill; treating it as a
112/// graceful stop keeps the socket unlinked and services drained even though a
113/// `daemon start`-launched daemon sits in its own session and never sees a
114/// terminal hangup). Elsewhere it listens for Ctrl-C only.
115pub 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        // macOS under launchd: soft 256, hard unlimited.
168        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        // Already at the hard limit: nothing to gain.
185        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    /// Runs [`raise_nofile_limit_with`] over `limits`, with `refuse` deciding which
197    /// targets the "host" rejects, and returns the outcome plus every target it
198    /// was asked to set.
199    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        // macOS refuses a soft limit above `kern.maxfilesperproc`.
243        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    /// The real accessors round-trip. Re-applying the limits the process already
256    /// has changes nothing, so this cannot starve a concurrently running test —
257    /// which lowering the soft limit to exercise a real raise would.
258    #[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}