Skip to main content

scx_arena/
lib.rs

1// Copyright (c) Meta Platforms, Inc. and affiliates.
2//
3// This software may be used and distributed according to the terms of the
4// GNU General Public License version 2.
5
6//! # SCX Arena library setup utilities
7//!
8//! Crate for setting up the BPF arena library for sched-ext schedulers.
9
10mod bpf_skel;
11
12mod arenalib;
13pub use arenalib::ArenaLib;
14
15use std::os::fd::AsFd;
16use std::os::fd::AsRawFd;
17use std::os::fd::BorrowedFd;
18use std::os::fd::FromRawFd;
19use std::os::fd::OwnedFd;
20use std::time::Duration;
21use std::time::Instant;
22
23use anyhow::bail;
24use anyhow::Context;
25use anyhow::Result;
26use libbpf_rs::libbpf_sys;
27use libbpf_rs::AsRawLibbpf as _;
28use libbpf_rs::MapCore as _;
29
30/// Cacheline size assumed by the arena allocator's alignment parameter.
31/// Mirrors scheds/include/lib/const-defs.h, keep in sync.
32#[cfg(target_arch = "s390x")]
33pub const CACHELINE_SIZE: usize = 256;
34#[cfg(target_arch = "powerpc64")]
35pub const CACHELINE_SIZE: usize = 128;
36#[cfg(not(any(target_arch = "s390x", target_arch = "powerpc64")))]
37pub const CACHELINE_SIZE: usize = 64;
38
39const MEMBARRIER_CMD_GLOBAL: libc::c_long = 1;
40const URCU_DOORBELL: &str = "scx_urcu_doorbell";
41const URCU_MIN_INTERVAL: Duration = Duration::from_millis(1);
42
43/// One background thread and the eventfd that stops it. Dropping writes the
44/// eventfd and joins the thread.
45#[derive(Debug)]
46pub(crate) struct Daemon {
47    stop: OwnedFd,
48    thread: Option<std::thread::JoinHandle<()>>,
49}
50
51impl Drop for Daemon {
52    fn drop(&mut self) {
53        let one: u64 = 1;
54        let ret = unsafe {
55            libc::write(
56                self.stop.as_raw_fd(),
57                &one as *const u64 as *const libc::c_void,
58                std::mem::size_of::<u64>(),
59            )
60        };
61        /* on a failed wakeup, leak the thread rather than hang the join */
62        if ret != std::mem::size_of::<u64>() as isize {
63            return;
64        }
65        if let Some(thread) = self.thread.take() {
66            let _ = thread.join();
67        }
68    }
69}
70
71/// Create the eventfd a Daemon is stopped through.
72fn stop_eventfd() -> Result<OwnedFd> {
73    let fd = unsafe { libc::eventfd(0, libc::EFD_CLOEXEC) };
74    if fd < 0 {
75        bail!(
76            "creating daemon stop eventfd failed: {}",
77            std::io::Error::last_os_error()
78        );
79    }
80    Ok(unsafe { OwnedFd::from_raw_fd(fd) })
81}
82
83/// Duplicate @prog's fd, None for programs that were not loaded and thus have
84/// no fd, prog.as_fd() on those would construct a BorrowedFd from an invalid
85/// fd. Errors only when a loaded program's fd cannot be duplicated.
86fn prog_fd_clone(prog: &libbpf_rs::Program<'_>) -> Result<Option<OwnedFd>> {
87    let raw_fd = unsafe { libbpf_sys::bpf_program__fd(prog.as_libbpf_object().as_ptr()) };
88    if raw_fd < 0 {
89        return Ok(None);
90    }
91    let fd = unsafe { BorrowedFd::borrow_raw(raw_fd) }
92        .try_clone_to_owned()
93        .with_context(|| format!("cloning the fd of BPF prog {:?}", prog.name()))?;
94    Ok(Some(fd))
95}
96
97/// Userspace half of the scx_urcu machinery, see the BPF side in
98/// lib/sdt_alloc.bpf.c. ArenaLib::setup() spawns the daemon when the object
99/// carries the scx_urcu doorbell. The returned ArenaLib owns it and stops and
100/// joins it on drop, so nothing else is visible to the scheduler: its whole
101/// runtime surface is calling the scx_*_free_rcu() variants from its free
102/// path hooks.
103///
104/// The daemon sleeps on the doorbell and, when woken, waits an RCU grace
105/// period, membarrier(MEMBARRIER_CMD_GLOBAL) is synchronize_rcu(), and runs
106/// the lib-provided scx_urcu_<storage>_pending/reclaim driver programs until
107/// nothing is awaiting reclaim, one grace period per cycle shared across the
108/// storages and at least URCU_MIN_INTERVAL between the side flips.
109fn urcu_run_prog(fd: &OwnedFd) -> Result<u32> {
110    let mut opts: libbpf_sys::bpf_test_run_opts = unsafe { std::mem::zeroed() };
111
112    opts.sz = std::mem::size_of::<libbpf_sys::bpf_test_run_opts>() as _;
113    let ret = unsafe { libbpf_sys::bpf_prog_test_run_opts(fd.as_raw_fd(), &mut opts) };
114    if ret != 0 {
115        bail!("urcu driver program run failed: {}", ret);
116    }
117    Ok(opts.retval)
118}
119
120fn urcu_daemon(
121    stop: OwnedFd,
122    doorbell: libbpf_rs::MapHandle,
123    pairs: Vec<(OwnedFd, OwnedFd)>,
124) -> Result<()> {
125    let mut builder = libbpf_rs::RingBufferBuilder::new();
126    builder
127        .add(&doorbell, |_| 0)
128        .context("adding urcu doorbell to ring buffer")?;
129    let rb = builder
130        .build()
131        .context("building urcu doorbell ring buffer")?;
132
133    let mut last: Option<Instant> = None;
134    loop {
135        let mut fds = [
136            libc::pollfd {
137                fd: doorbell.as_fd().as_raw_fd(),
138                events: libc::POLLIN,
139                revents: 0,
140            },
141            libc::pollfd {
142                fd: stop.as_raw_fd(),
143                events: libc::POLLIN,
144                revents: 0,
145            },
146        ];
147
148        let ret = unsafe { libc::poll(fds.as_mut_ptr(), 2, -1) };
149        if ret < 0 {
150            let err = std::io::Error::last_os_error();
151            if err.raw_os_error() == Some(libc::EINTR) {
152                continue;
153            }
154            bail!("urcu doorbell poll failed: {}", err);
155        }
156
157        /* run one final drain below before honoring a stop request */
158        let stopping = fds[1].revents != 0;
159
160        /* drain until nothing is awaiting reclaim */
161        loop {
162            if let Some(last) = last {
163                let elapsed = last.elapsed();
164                if elapsed < URCU_MIN_INTERVAL {
165                    std::thread::sleep(URCU_MIN_INTERVAL - elapsed);
166                }
167            }
168
169            /*
170             * Consume before the pending checks. A free landing afterwards is
171             * either seen by the checks or leaves its ring behind and the poll
172             * returns immediately. The reverse order can consume the ring of a
173             * free the checks missed and sleep on it.
174             */
175            rb.consume().context("consuming urcu doorbell")?;
176
177            let mut reclaims = Vec::new();
178            for (pending, reclaim) in &pairs {
179                if urcu_run_prog(pending)? != 0 {
180                    reclaims.push(reclaim);
181                }
182            }
183            if reclaims.is_empty() {
184                break;
185            }
186            last = Some(Instant::now());
187
188            let ret = unsafe { libc::syscall(libc::SYS_membarrier, MEMBARRIER_CMD_GLOBAL, 0, 0) };
189            if ret != 0 {
190                bail!(
191                    "membarrier(GLOBAL) failed: {}",
192                    std::io::Error::last_os_error()
193                );
194            }
195
196            for reclaim in reclaims {
197                while urcu_run_prog(reclaim)? != 0 {}
198            }
199        }
200
201        if stopping {
202            return Ok(());
203        }
204    }
205}
206
207/// Spawn the urcu reclaim daemon if @obj carries the scx_urcu doorbell and
208/// driver programs. Called from ArenaLib::setup(), the daemon is owned by
209/// the returned ArenaLib.
210pub(crate) fn urcu_spawn(obj: &libbpf_rs::Object) -> Result<Option<Daemon>> {
211    let Some(doorbell) = obj.maps().find(|m| m.name() == URCU_DOORBELL) else {
212        return Ok(None);
213    };
214    let doorbell =
215        libbpf_rs::MapHandle::try_from(&doorbell).context("cloning urcu doorbell handle")?;
216
217    let mut pairs = Vec::new();
218    for prog in obj.progs() {
219        let Some(name) = prog.name().to_str() else {
220            continue;
221        };
222        let Some(base) = name.strip_suffix("_pending") else {
223            continue;
224        };
225        if !name.starts_with("scx_urcu_") {
226            continue;
227        }
228
229        let reclaim_name = format!("{}_reclaim", base);
230        let reclaim = obj
231            .progs()
232            .find(|p| p.name() == reclaim_name.as_str())
233            .with_context(|| format!("urcu driver program {} not found", reclaim_name))?;
234
235        let Some(pending_fd) = prog_fd_clone(&prog)? else {
236            bail!("urcu driver program {} is not loaded", name);
237        };
238        let Some(reclaim_fd) = prog_fd_clone(&reclaim)? else {
239            bail!("urcu driver program {} is not loaded", reclaim_name);
240        };
241        pairs.push((pending_fd, reclaim_fd));
242    }
243    if pairs.is_empty() {
244        return Ok(None);
245    }
246
247    let stop = stop_eventfd()?;
248    let daemon_stop = stop.try_clone().context("cloning urcu stop eventfd")?;
249    let thread = std::thread::Builder::new()
250        .name("scx-urcu".into())
251        .spawn(move || {
252            if let Err(e) = urcu_daemon(daemon_stop, doorbell, pairs) {
253                eprintln!("scx-urcu daemon exiting on error: {:#}", e);
254            }
255        })
256        .context("spawning urcu daemon thread")?;
257
258    Ok(Some(Daemon {
259        stop,
260        thread: Some(thread),
261    }))
262}
263
264const BPF_STDOUT: u32 = 1;
265const BPF_STDERR: u32 = 2;
266const BPF_STREAMS: [(u32, &str); 2] = [(BPF_STDOUT, "stdout"), (BPF_STDERR, "stderr")];
267const STREAM_POLL_INTERVAL: Duration = Duration::from_secs(1);
268
269fn stream_read(fd: &OwnedFd, stream_id: u32, buf: &mut [u8]) -> i32 {
270    unsafe {
271        libbpf_sys::bpf_prog_stream_read(
272            fd.as_raw_fd(),
273            stream_id,
274            buf.as_mut_ptr() as *mut _,
275            buf.len() as u32,
276            std::ptr::null_mut(),
277        )
278    }
279}
280
281/// The kernel prefixes the errors it reports on a program's BPF stderr
282/// stream with "ERROR: ".
283fn stream_is_fatal(lines: &[String]) -> bool {
284    lines.iter().any(|l| l.starts_with("ERROR: "))
285}
286
287/// Split the complete lines out of @carry, keeping a trailing partial for the
288/// next read. Stream elements are concatenated without separators and reads
289/// split lines arbitrarily, so prefixes may only be matched on complete lines.
290/// @new_data says whether this poll read anything: a partial with no
291/// continuation after a full poll interval is flushed as a line of its own.
292fn stream_extract_lines(carry: &mut String, new_data: bool) -> Vec<String> {
293    let mut lines: Vec<String> = Vec::new();
294
295    if let Some(pos) = carry.rfind('\n') {
296        lines = carry[..pos].split('\n').map(str::to_string).collect();
297        carry.drain(..=pos);
298    }
299    if !new_data && !carry.is_empty() {
300        lines.push(std::mem::take(carry));
301    }
302
303    lines
304}
305
306/// Forward the BPF stdout/stderr streams of every program in the object to
307/// the scheduler's stdout and stderr respectively, and abort when the
308/// kernel reports an error on a stderr stream. Nothing reads these streams
309/// otherwise, so messages would be silently dropped. Lines prefixed "IGN: "
310/// are dropped instead of forwarded, for prints with side effects that
311/// nobody needs to see, the arena association print in
312/// scx_arena_subprog_init() for example. Called from ArenaLib::setup(), the
313/// watcher is owned by the returned ArenaLib.
314pub(crate) fn stream_watcher_spawn(obj: &libbpf_rs::Object) -> Result<Daemon> {
315    let mut progs = Vec::new();
316
317    for prog in obj.progs() {
318        let Some(name) = prog.name().to_str() else {
319            continue;
320        };
321        /* feature-gated programs may not be loaded and then have no fd */
322        let Some(fd) = prog_fd_clone(&prog)? else {
323            continue;
324        };
325        progs.push((name.to_string(), fd));
326    }
327
328    let stop = stop_eventfd()?;
329    let daemon_stop = stop
330        .try_clone()
331        .context("cloning stream watcher stop eventfd")?;
332    let thread = std::thread::Builder::new()
333        .name("scx-bpf-stream".into())
334        .spawn(move || {
335            let mut buf = vec![0u8; 65536];
336            let mut carries: Vec<[String; 2]> = (0..progs.len())
337                .map(|_| [String::new(), String::new()])
338                .collect();
339
340            loop {
341                let mut fds = [libc::pollfd {
342                    fd: daemon_stop.as_raw_fd(),
343                    events: libc::POLLIN,
344                    revents: 0,
345                }];
346                let ret = unsafe {
347                    libc::poll(fds.as_mut_ptr(), 1, STREAM_POLL_INTERVAL.as_millis() as i32)
348                };
349                /* run one final scan below before honoring a stop request */
350                let stopping = if ret < 0 {
351                    std::io::Error::last_os_error().raw_os_error() != Some(libc::EINTR)
352                } else {
353                    ret > 0
354                };
355
356                for (pidx, (name, fd)) in progs.iter().enumerate() {
357                    for (sidx, (stream_id, label)) in BPF_STREAMS.iter().enumerate() {
358                        let carry = &mut carries[pidx][sidx];
359                        let mut new_data = false;
360
361                        /* drain fully, a backlog can exceed the buffer */
362                        loop {
363                            let n = stream_read(fd, *stream_id, &mut buf);
364                            if n <= 0 {
365                                break;
366                            }
367                            carry.push_str(&String::from_utf8_lossy(&buf[..n as usize]));
368                            new_data = true;
369                            if (n as usize) < buf.len() {
370                                break;
371                            }
372                        }
373
374                        let lines = stream_extract_lines(carry, new_data);
375                        if lines.is_empty() {
376                            continue;
377                        }
378
379                        let kept: Vec<&str> = lines
380                            .iter()
381                            .map(String::as_str)
382                            .filter(|l| !l.starts_with("IGN: "))
383                            .collect();
384                        if !kept.is_empty() {
385                            let msg = kept.join("\n");
386                            let out = format!("BPF {} of prog {}:\n{}\n", label, name, msg);
387                            if *stream_id == BPF_STDOUT {
388                                print!("{}", out);
389                                let _ = std::io::Write::flush(&mut std::io::stdout());
390                            } else {
391                                eprint!("{}", out);
392                            }
393                        }
394                        if *stream_id == BPF_STDERR && stream_is_fatal(&lines) {
395                            eprintln!("FATAL: aborting on BPF error report");
396                            std::process::exit(1);
397                        }
398                    }
399                }
400
401                if stopping {
402                    break;
403                }
404            }
405        })
406        .context("spawning BPF stream watcher")?;
407
408    Ok(Daemon {
409        stop,
410        thread: Some(thread),
411    })
412}
413
414#[cfg(test)]
415mod tests {
416    use super::stream_extract_lines;
417    use super::stream_is_fatal;
418
419    #[test]
420    fn test_stream_extract_lines() {
421        let mut carry = String::new();
422
423        /* empty carry, quiet tick */
424        assert!(stream_extract_lines(&mut carry, false).is_empty());
425
426        /* one complete line plus a partial */
427        carry.push_str("ERROR: whole line\npartial");
428        assert_eq!(
429            stream_extract_lines(&mut carry, true),
430            ["ERROR: whole line"]
431        );
432        assert_eq!(carry, "partial");
433
434        /* the partial is not flushed while data keeps arriving */
435        assert!(stream_extract_lines(&mut carry, true).is_empty());
436        assert_eq!(carry, "partial");
437
438        /* a line split across reads is reassembled */
439        carry.push_str(" continued\nnext");
440        assert_eq!(
441            stream_extract_lines(&mut carry, true),
442            ["partial continued"]
443        );
444        assert_eq!(carry, "next");
445
446        /* a stale partial flushes on a quiet tick, exactly once */
447        assert_eq!(stream_extract_lines(&mut carry, false), ["next"]);
448        assert!(carry.is_empty());
449        assert!(stream_extract_lines(&mut carry, false).is_empty());
450
451        /* multiple lines per chunk, blank lines forwarded */
452        carry.push_str("a\n\nb\ntail");
453        assert_eq!(stream_extract_lines(&mut carry, true), ["a", "", "b"]);
454        assert_eq!(carry, "tail");
455        carry.clear();
456
457        /* multi-byte UTF-8 split across reads stays on char boundaries */
458        carry.push_str("caf");
459        assert!(stream_extract_lines(&mut carry, true).is_empty());
460        carry.push_str("\u{e9}\n");
461        assert_eq!(stream_extract_lines(&mut carry, true), ["caf\u{e9}"]);
462        assert!(carry.is_empty());
463    }
464
465    #[test]
466    fn test_stream_is_fatal() {
467        let fatal = ["noise".to_string(), "ERROR: Arena READ access".to_string()];
468        assert!(stream_is_fatal(&fatal));
469
470        /* the prefix only counts at the start of a line */
471        let glued = ["no task dataERROR: Arena READ access".to_string()];
472        assert!(!stream_is_fatal(&glued));
473        assert!(!stream_is_fatal(&[]));
474    }
475}