Skip to main content

scx_characterize/
main.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
6use anyhow::{bail, Context as _, Result};
7use clap::{Parser, Subcommand};
8use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd};
9
10mod extract;
11mod process;
12mod record;
13mod sched_util;
14
15pub mod bpf_intf;
16pub mod bpf_skel;
17pub use bpf_skel as bpf;
18
19fn create_eventfd() -> Result<OwnedFd> {
20    let fd = unsafe { libc::eventfd(0, libc::EFD_NONBLOCK | libc::EFD_CLOEXEC) };
21    if fd < 0 {
22        bail!("eventfd failed: {}", std::io::Error::last_os_error());
23    }
24    Ok(unsafe { OwnedFd::from_raw_fd(fd) })
25}
26
27fn signal_eventfd(fd: RawFd) {
28    let val: u64 = 1;
29    unsafe {
30        libc::write(fd, &val as *const u64 as *const libc::c_void, 8);
31    }
32}
33
34pub struct Context {
35    shutdown_fd: OwnedFd,
36}
37
38impl Context {
39    fn new() -> Result<Self> {
40        let shutdown_fd = create_eventfd()?;
41        let raw_fd = shutdown_fd.as_raw_fd();
42
43        ctrlc::set_handler(move || {
44            signal_eventfd(raw_fd);
45        })
46        .context("failed to set Ctrl+C handler")?;
47
48        Ok(Self { shutdown_fd })
49    }
50
51    pub fn shutdown_fd(&self) -> RawFd {
52        self.shutdown_fd.as_raw_fd()
53    }
54}
55
56#[derive(Debug, Parser)]
57#[command(
58    name = "scx_characterize",
59    about = "Workload characterization tool for sched_ext schedulers"
60)]
61struct Opts {
62    #[clap(subcommand)]
63    command: Commands,
64}
65
66#[derive(Debug, Subcommand)]
67enum Commands {
68    /// Record workload profile using perf
69    Record(record::RecordOpts),
70    /// Process a recorded profile
71    Process(process::ProcessOpts),
72    /// Extract data from processed profile
73    Extract(extract::ExtractOpts),
74}
75
76fn main() -> Result<()> {
77    let ctx = Context::new()?;
78    let opts = Opts::parse();
79
80    match opts.command {
81        Commands::Record(record_opts) => record::cmd_record(&ctx, record_opts),
82        Commands::Process(process_opts) => process::cmd_process(process_opts),
83        Commands::Extract(extract_opts) => extract::cmd_extract(extract_opts),
84    }
85}