From 7c32acece06f3760e6046196e10eb7bce5a2821f Mon Sep 17 00:00:00 2001
From: Jordan Rome <linux@jordanrome.com>
Date: Mon, 19 Feb 2024 16:50:32 -0800
Subject: [PATCH] Add libbpf logging to the rust schedulers

This is to get better logs when failing to load, attach, etc.
---
 rust/scx_utils/src/lib.rs           |  3 +++
 rust/scx_utils/src/libbpf_logger.rs | 20 ++++++++++++++++++++
 scheds/rust/scx_layered/src/main.rs |  2 ++
 scheds/rust/scx_rustland/src/bpf.rs |  2 ++
 scheds/rust/scx_rusty/src/main.rs   |  2 ++
 5 files changed, 29 insertions(+)
 create mode 100644 rust/scx_utils/src/libbpf_logger.rs

diff --git a/rust/scx_utils/src/lib.rs b/rust/scx_utils/src/lib.rs
index 53644cf..c98d90b 100644
--- a/rust/scx_utils/src/lib.rs
+++ b/rust/scx_utils/src/lib.rs
@@ -37,6 +37,9 @@ pub use bpf_builder::BpfBuilder;
 
 pub mod ravg;
 
+mod libbpf_logger;
+pub use libbpf_logger::init_libbpf_logging;
+
 mod user_exit_info;
 pub use user_exit_info::UserExitInfo;
 pub use user_exit_info::ScxExitKind;
diff --git a/rust/scx_utils/src/libbpf_logger.rs b/rust/scx_utils/src/libbpf_logger.rs
new file mode 100644
index 0000000..e40de1b
--- /dev/null
+++ b/rust/scx_utils/src/libbpf_logger.rs
@@ -0,0 +1,20 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This software may be used and distributed according to the terms of the
+// GNU General Public License version 2.
+
+use libbpf_rs::{PrintLevel, set_print};
+
+fn print_to_log(level: PrintLevel, msg: String) {
+    match level {
+        PrintLevel::Debug => log::debug!("{}", msg),
+        PrintLevel::Info => log::info!("{}", msg),
+        PrintLevel::Warn => log::warn!("{}", msg),
+    }
+}
+
+pub fn init_libbpf_logging(
+    level: Option<PrintLevel>,
+) {
+    set_print(Some((level.unwrap_or(PrintLevel::Debug), print_to_log)));
+}
\ No newline at end of file
diff --git a/scheds/rust/scx_layered/src/main.rs b/scheds/rust/scx_layered/src/main.rs
index a861d34..1612724 100644
--- a/scheds/rust/scx_layered/src/main.rs
+++ b/scheds/rust/scx_layered/src/main.rs
@@ -38,6 +38,7 @@ use prometheus_client::encoding::text::encode;
 use prometheus_client::metrics::family::Family;
 use prometheus_client::metrics::gauge::Gauge;
 use prometheus_client::registry::Registry;
+use scx_utils::init_libbpf_logging;
 use scx_utils::ravg::ravg_read;
 use scx_utils::uei_exited;
 use scx_utils::uei_report;
@@ -1225,6 +1226,7 @@ impl<'a> Scheduler<'a> {
         // Open the BPF prog first for verification.
         let mut skel_builder = BpfSkelBuilder::default();
         skel_builder.obj_builder.debug(opts.verbose > 1);
+        init_libbpf_logging(None);
         let mut skel = skel_builder.open().context("Failed to open BPF program")?;
 
         // Initialize skel according to @opts.
diff --git a/scheds/rust/scx_rustland/src/bpf.rs b/scheds/rust/scx_rustland/src/bpf.rs
index 73909f5..9e11b3d 100644
--- a/scheds/rust/scx_rustland/src/bpf.rs
+++ b/scheds/rust/scx_rustland/src/bpf.rs
@@ -18,6 +18,7 @@ use libc::{sched_param, sched_setscheduler};
 mod alloc;
 use alloc::*;
 
+use scx_utils::init_libbpf_logging;
 use scx_utils::uei_exited;
 use scx_utils::uei_report;
 
@@ -226,6 +227,7 @@ impl<'cb> BpfScheduler<'cb> {
     pub fn init(slice_us: u64, nr_cpus_online: i32, partial: bool, debug: bool) -> Result<Self> {
         // Open the BPF prog first for verification.
         let skel_builder = BpfSkelBuilder::default();
+        init_libbpf_logging(None);
         let mut skel = skel_builder.open().context("Failed to open BPF program")?;
 
         // Lock all the memory to prevent page faults that could trigger potential deadlocks during
diff --git a/scheds/rust/scx_rusty/src/main.rs b/scheds/rust/scx_rusty/src/main.rs
index 19d0b91..05621bb 100644
--- a/scheds/rust/scx_rusty/src/main.rs
+++ b/scheds/rust/scx_rusty/src/main.rs
@@ -33,6 +33,7 @@ use log::info;
 use log::trace;
 use log::warn;
 use ordered_float::OrderedFloat;
+use scx_utils::init_libbpf_logging;
 use scx_utils::ravg::ravg_read;
 use scx_utils::uei_exited;
 use scx_utils::uei_report;
@@ -982,6 +983,7 @@ impl<'a> Scheduler<'a> {
         // Open the BPF prog first for verification.
         let mut skel_builder = BpfSkelBuilder::default();
         skel_builder.obj_builder.debug(opts.verbose > 0);
+        init_libbpf_logging(None);
         let mut skel = skel_builder.open().context("Failed to open BPF program")?;
 
         // Initialize skel according to @opts.