scxtop/
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 scx_utils::compat;
7use scxtop::Action;
8use scxtop::App;
9use scxtop::Event;
10use scxtop::Key;
11use scxtop::KeyMap;
12use scxtop::PerfettoTraceManager;
13use scxtop::SCHED_NAME_PATH;
14use scxtop::Tui;
15use scxtop::bpf_skel::types::bpf_event;
16use scxtop::bpf_skel::*;
17use scxtop::cli::{Cli, Commands, TraceArgs, TuiArgs, generate_completions};
18use scxtop::config::Config;
19use scxtop::edm::{ActionHandler, BpfEventActionPublisher, BpfEventHandler, EventDispatchManager};
20use scxtop::mangoapp::poll_mangoapp;
21use scxtop::read_file_string;
22use scxtop::tracer::Tracer;
23
24use anyhow::Result;
25use anyhow::anyhow;
26use clap::{CommandFactory, Parser};
27use libbpf_rs::Link;
28use libbpf_rs::ProgramInput;
29use libbpf_rs::RingBufferBuilder;
30use libbpf_rs::UprobeOpts;
31use libbpf_rs::skel::OpenSkel;
32use libbpf_rs::skel::SkelBuilder;
33use log::debug;
34use log::info;
35use ratatui::crossterm::event::KeyCode::Char;
36use simplelog::{
37    ColorChoice, Config as SimplelogConfig, LevelFilter, TermLogger, TerminalMode, WriteLogger,
38};
39use std::sync::atomic::AtomicBool;
40use tokio::sync::mpsc;
41
42use std::ffi::CString;
43use std::fs::File;
44use std::mem::MaybeUninit;
45use std::str::FromStr;
46use std::sync::Arc;
47use std::sync::atomic::Ordering;
48use std::time::Duration;
49
50fn get_action(_app: &App, keymap: &KeyMap, event: Event) -> Action {
51    match event {
52        Event::Error => Action::None,
53        Event::Tick => Action::Tick,
54        Event::TickRateChange(tick_rate_ms) => {
55            Action::TickRateChange(std::time::Duration::from_millis(tick_rate_ms))
56        }
57        Event::Key(key) => match key.code {
58            Char(c) => keymap.action(&Key::Char(c)),
59            _ => keymap.action(&Key::Code(key.code)),
60        },
61        _ => Action::None,
62    }
63}
64
65/// Attaches BPF programs to the skel.
66fn attach_progs(skel: &mut BpfSkel) -> Result<Vec<Link>> {
67    // Attach probes
68    let mut links = vec![
69        skel.progs.on_sched_cpu_perf.attach()?,
70        skel.progs.scx_sched_reg.attach()?,
71        skel.progs.scx_sched_unreg.attach()?,
72        skel.progs.on_sched_switch.attach()?,
73        skel.progs.on_sched_wakeup.attach()?,
74        skel.progs.on_sched_wakeup_new.attach()?,
75        skel.progs.on_sched_waking.attach()?,
76    ];
77
78    // 6.13 compatibility
79    if compat::ksym_exists("scx_bpf_dsq_insert_vtime")? {
80        if let Ok(link) = skel.progs.scx_insert_vtime.attach() {
81            links.push(link);
82        }
83        if let Ok(link) = skel.progs.scx_insert.attach() {
84            links.push(link);
85        }
86        if let Ok(link) = skel.progs.scx_dsq_move.attach() {
87            links.push(link);
88        }
89        if let Ok(link) = skel.progs.scx_dsq_move_set_vtime.attach() {
90            links.push(link);
91        }
92        if let Ok(link) = skel.progs.scx_dsq_move_set_slice.attach() {
93            links.push(link);
94        }
95    } else {
96        if let Ok(link) = skel.progs.scx_dispatch.attach() {
97            links.push(link);
98        }
99        if let Ok(link) = skel.progs.scx_dispatch_vtime.attach() {
100            links.push(link);
101        }
102        if let Ok(link) = skel.progs.scx_dispatch_from_dsq_set_vtime.attach() {
103            links.push(link);
104        }
105        if let Ok(link) = skel.progs.scx_dispatch_from_dsq_set_slice.attach() {
106            links.push(link);
107        }
108        if let Ok(link) = skel.progs.scx_dispatch_from_dsq.attach() {
109            links.push(link);
110        }
111    }
112    if let Ok(link) = skel.progs.on_cpuhp_enter.attach() {
113        links.push(link);
114    }
115    if let Ok(link) = skel.progs.on_cpuhp_exit.attach() {
116        links.push(link);
117    }
118    if let Ok(link) = skel.progs.on_pstate_sample.attach() {
119        links.push(link);
120    }
121
122    Ok(links)
123}
124
125fn run_trace(trace_args: &TraceArgs) -> Result<()> {
126    TermLogger::init(
127        match trace_args.verbose {
128            0 => simplelog::LevelFilter::Info,
129            1 => simplelog::LevelFilter::Debug,
130            _ => simplelog::LevelFilter::Trace,
131        },
132        SimplelogConfig::default(),
133        TerminalMode::Mixed,
134        ColorChoice::Auto,
135    )?;
136
137    let config = Config::default_config();
138    let worker_threads = config.worker_threads() as usize;
139    tokio::runtime::Builder::new_multi_thread()
140        .enable_all()
141        .worker_threads(if worker_threads > 2 {
142            worker_threads
143        } else {
144            4
145        })
146        .build()
147        .unwrap()
148        .block_on(async {
149            let (action_tx, mut action_rx) = mpsc::unbounded_channel();
150
151            let mut open_object = MaybeUninit::uninit();
152            let mut builder = BpfSkelBuilder::default();
153            if trace_args.verbose > 2 {
154                builder.obj_builder.debug(true);
155            }
156
157            let skel = builder.open(&mut open_object)?;
158            // set sample rate to 1 here to populate the BPF tctxs
159            skel.maps.data_data.sample_rate = 1;
160            compat::cond_kprobe_enable("gpu_memory_total", &skel.progs.on_gpu_memory_total)?;
161            compat::cond_kprobe_enable("hw_pressure_update", &skel.progs.on_hw_pressure_update)?;
162
163            let mut skel = skel.load()?;
164            let mut links = attach_progs(&mut skel)?;
165            links.push(skel.progs.on_sched_fork.attach()?);
166            links.push(skel.progs.on_sched_exec.attach()?);
167            links.push(skel.progs.on_sched_exit.attach()?);
168
169            let trace_dur = std::time::Duration::from_millis(trace_args.trace_ms);
170            let bpf_publisher = BpfEventActionPublisher::new(action_tx.clone());
171
172            let mut event_rbb = RingBufferBuilder::new();
173            let mut edm = EventDispatchManager::new(None, None);
174            let warmup_done = Arc::new(AtomicBool::new(false));
175            edm.register_bpf_handler(Box::new(bpf_publisher));
176            let event_handler = move |data: &[u8]| {
177                let mut event = bpf_event::default();
178                plain::copy_from_bytes(&mut event, data).expect("Event data buffer was too short");
179                let _ = edm.on_event(&event);
180                0
181            };
182
183            event_rbb.add(&skel.maps.events, event_handler)?;
184            let event_rb = event_rbb.build()?;
185            let shutdown = Arc::new(AtomicBool::new(false));
186            let stop_poll = shutdown.clone();
187            let stop_main = shutdown.clone();
188
189            tokio::spawn(async move {
190                loop {
191                    let _ = event_rb.poll(Duration::from_millis(1));
192                    if stop_poll.load(Ordering::Relaxed) {
193                        debug!("polling stopped");
194                        break;
195                    }
196                }
197            });
198
199            let (complete_tx, mut complete_rx) = mpsc::channel(1);
200            let trace_file_prefix = config.trace_file_prefix().to_string();
201            let trace_file = trace_args.output_file.clone();
202            let mut trace_manager = PerfettoTraceManager::new(trace_file_prefix, None);
203            info!("warming up for {}ms", trace_args.warmup_ms);
204            tokio::time::sleep(Duration::from_millis(trace_args.warmup_ms)).await;
205            debug!("starting trace");
206            let thread_warmup = warmup_done.clone();
207            trace_manager.start()?;
208            thread_warmup.store(true, Ordering::Relaxed);
209            tokio::spawn(async move {
210                let mut count = 0;
211                loop {
212                    let action = action_rx.recv().await;
213                    if let Some(a) = action {
214                        if thread_warmup.load(Ordering::Relaxed) {
215                            count += 1;
216                            trace_manager.on_action(&a).unwrap();
217                        }
218                    }
219                    if stop_main.load(Ordering::Relaxed) {
220                        trace_manager.stop(trace_file, None).unwrap();
221                        info!("trace complete, collected {} events", count);
222                        let _ = complete_tx.send(1).await;
223                        break;
224                    }
225                }
226            });
227
228            let mut tracer = Tracer::new(skel);
229            tracer.trace_async(trace_dur).await?;
230
231            // The order is important here:
232            // 1) first drop the links to detach the attached BPF programs
233            // 2) set the shutdown variable to stop background tokio threads
234            // 3) wait for the completion of the trace file generation to complete
235            drop(links);
236            shutdown.store(true, Ordering::Relaxed);
237            let _ = complete_rx.recv().await;
238
239            let stats = tracer.stats()?;
240            info!("{:?}", stats);
241
242            Ok(())
243        })
244}
245
246fn run_tui(tui_args: &TuiArgs) -> Result<()> {
247    if let Ok(log_path) = std::env::var("RUST_LOG_PATH") {
248        let log_level = match std::env::var("RUST_LOG") {
249            Ok(v) => LevelFilter::from_str(&v)?,
250            Err(_) => LevelFilter::Info,
251        };
252
253        WriteLogger::init(
254            log_level,
255            simplelog::Config::default(),
256            File::create(log_path)?,
257        )?;
258
259        log_panics::Config::new()
260            .backtrace_mode(log_panics::BacktraceMode::Resolved)
261            .install_panic_hook();
262    };
263
264    let config = Config::merge([
265        Config::from(tui_args.clone()),
266        Config::load().unwrap_or(Config::default_config()),
267    ]);
268    let keymap = config.active_keymap.clone();
269
270    tokio::runtime::Builder::new_multi_thread()
271        .enable_all()
272        .worker_threads(config.worker_threads() as usize)
273        .build()
274        .unwrap()
275        .block_on(async {
276            let (action_tx, mut action_rx) = mpsc::unbounded_channel();
277
278            let mut open_object = MaybeUninit::uninit();
279            let mut builder = BpfSkelBuilder::default();
280            if config.debug() {
281                builder.obj_builder.debug(true);
282            }
283            let bpf_publisher = BpfEventActionPublisher::new(action_tx.clone());
284            let mut edm = EventDispatchManager::new(None, None);
285            edm.register_bpf_handler(Box::new(bpf_publisher));
286
287            let skel = builder.open(&mut open_object)?;
288            skel.maps.rodata_data.long_tail_tracing_min_latency_ns =
289                tui_args.experimental_long_tail_tracing_min_latency_ns;
290
291            compat::cond_kprobe_enable("gpu_memory_total", &skel.progs.on_gpu_memory_total)?;
292            compat::cond_kprobe_enable("hw_pressure_update", &skel.progs.on_hw_pressure_update)?;
293            let mut skel = skel.load()?;
294            let mut links = attach_progs(&mut skel)?;
295            skel.progs.scxtop_init.test_run(ProgramInput::default())?;
296
297            if tui_args.experimental_long_tail_tracing {
298                skel.maps.data_data.trace_duration_ns = config.trace_duration_ns();
299                skel.maps.data_data.trace_warmup_ns = config.trace_warmup_ns();
300
301                let binary = tui_args
302                    .experimental_long_tail_tracing_binary
303                    .clone()
304                    .unwrap();
305                let symbol = tui_args
306                    .experimental_long_tail_tracing_symbol
307                    .clone()
308                    .unwrap();
309
310                links.extend([
311                    skel.progs.long_tail_tracker_exit.attach_uprobe_with_opts(
312                        -1, /* pid, -1 == all */
313                        binary.clone(),
314                        0,
315                        UprobeOpts {
316                            retprobe: true,
317                            func_name: symbol.clone(),
318                            ..Default::default()
319                        },
320                    )?,
321                    skel.progs.long_tail_tracker_entry.attach_uprobe_with_opts(
322                        -1, /* pid, -1 == all */
323                        binary.clone(),
324                        0,
325                        UprobeOpts {
326                            retprobe: false,
327                            func_name: symbol.clone(),
328                            ..Default::default()
329                        },
330                    )?,
331                ]);
332            };
333
334            let mut tui = Tui::new(keymap.clone(), config.tick_rate_ms())?;
335            let mut event_rbb = RingBufferBuilder::new();
336            let event_handler = move |data: &[u8]| {
337                let mut event = bpf_event::default();
338                plain::copy_from_bytes(&mut event, data).expect("Event data buffer was too short");
339                let _ = edm.on_event(&event);
340                0
341            };
342            event_rbb.add(&skel.maps.events, event_handler)?;
343            let event_rb = event_rbb.build()?;
344            let scheduler = read_file_string(SCHED_NAME_PATH).unwrap_or("".to_string());
345
346            let mut app = App::new(
347                config,
348                scheduler,
349                100,
350                tui_args.process_id,
351                action_tx.clone(),
352                skel,
353            )?;
354
355            tui.enter()?;
356
357            let shutdown = app.should_quit.clone();
358            tokio::spawn(async move {
359                loop {
360                    let _ = event_rb.poll(Duration::from_millis(1));
361                    if shutdown.load(Ordering::Relaxed) {
362                        break;
363                    }
364                }
365            });
366
367            if tui_args.mangoapp_tracing {
368                let stop_mangoapp = app.should_quit.clone();
369                let mangoapp_path = CString::new(tui_args.mangoapp_path.clone()).unwrap();
370                let poll_intvl_ms = tui_args.mangoapp_poll_intvl_ms;
371                let tx = action_tx.clone();
372                tokio::spawn(async move {
373                    poll_mangoapp(
374                        mangoapp_path,
375                        poll_intvl_ms,
376                        tx,
377                        stop_mangoapp,
378                    )
379                    .await
380                });
381            }
382
383
384            loop {
385                tokio::select! {
386                    ev = tui.next() => {
387                        let ev = ev?;
388                        match ev {
389                            Event::Quit => { action_tx.send(Action::Quit)?; },
390                            Event::Tick => action_tx.send(Action::Tick)?,
391                            Event::TickRateChange(tick_rate_ms) => action_tx.send(
392                                Action::TickRateChange(std::time::Duration::from_millis(tick_rate_ms)),
393                            )?,
394                            Event::Render => {
395                                if app.should_quit.load(Ordering::Relaxed) {
396                                    break;
397                                }
398                                tui.draw(|f| app.render(f).expect("Failed to render application"))?;
399                            }
400                            Event::Key(_) => {
401                                let action = get_action(&app, &keymap, ev);
402                                action_tx.send(action)?;
403                            }
404                            _ => {}
405                    }}
406
407                    ac = action_rx.recv() => {
408                        let ac = ac.ok_or(anyhow!("actions channel closed"))?;
409                        app.handle_action(&ac)?;
410                    }
411                }
412            }
413            tui.exit()?;
414            drop(links);
415
416            Ok(())
417        })
418}
419
420fn main() -> Result<()> {
421    let args = Cli::parse();
422
423    match &args.command.unwrap_or(Commands::Tui(args.tui)) {
424        Commands::Tui(tui_args) => {
425            run_tui(tui_args)?;
426        }
427        Commands::Trace(trace_args) => {
428            run_trace(trace_args)?;
429        }
430        Commands::GenerateCompletions { shell, output } => {
431            generate_completions(Cli::command(), *shell, output.clone())
432                .unwrap_or_else(|_| panic!("Failed to generate completions for {}", shell));
433        }
434    }
435    Ok(())
436}
scxtop/main.rs

scxtop/
main.rs
