From 42fe1485cf17dec8416bad157ef0422648e32726 Mon Sep 17 00:00:00 2001
From: bnaecker <ben@oxide.computer>
Date: Thu, 13 Jun 2024 17:06:51 -0700
Subject: [PATCH] Add an xtask for listing all USDT probes in Omicron (#5895)

---
 Cargo.lock                  |  2 +
 dev-tools/xtask/Cargo.toml  |  6 ++-
 dev-tools/xtask/src/main.rs | 11 +++++
 dev-tools/xtask/src/usdt.rs | 84 +++++++++++++++++++++++++++++++++++++
 4 files changed, 101 insertions(+), 2 deletions(-)
 create mode 100644 dev-tools/xtask/src/usdt.rs

diff --git a/Cargo.lock b/Cargo.lock
index af3738e0a0..5949144ab5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -11384,9 +11384,11 @@ dependencies = [
  "slog-term",
  "strum",
  "swrite",
+ "tabled",
  "tar",
  "tokio",
  "toml 0.8.13",
+ "usdt",
 ]
 
 [[package]]
diff --git a/dev-tools/xtask/Cargo.toml b/dev-tools/xtask/Cargo.toml
index 745e16dea6..9268eafc3e 100644
--- a/dev-tools/xtask/Cargo.toml
+++ b/dev-tools/xtask/Cargo.toml
@@ -15,6 +15,7 @@ cargo_metadata.workspace = true
 clap.workspace = true
 flate2.workspace = true
 futures.workspace = true
+fs-err.workspace = true
 macaddr.workspace = true
 md5 = "0.7.0"
 reqwest = { workspace = true, features = [ "default-tls" ] }
@@ -25,8 +26,9 @@ slog-async.workspace = true
 slog-bunyan.workspace = true
 slog-term.workspace = true
 strum.workspace = true
+swrite.workspace = true
+tabled.workspace = true
 tar.workspace = true
 tokio = { workspace = true, features = ["full"] }
 toml.workspace = true
-fs-err.workspace = true
-swrite.workspace = true
+usdt.workspace = true
diff --git a/dev-tools/xtask/src/main.rs b/dev-tools/xtask/src/main.rs
index 96afdac17d..3d8acceb3d 100644
--- a/dev-tools/xtask/src/main.rs
+++ b/dev-tools/xtask/src/main.rs
@@ -15,6 +15,7 @@ mod clippy;
 mod download;
 #[cfg_attr(not(target_os = "illumos"), allow(dead_code))]
 mod external;
+mod usdt;
 
 #[cfg(target_os = "illumos")]
 mod verify_libraries;
@@ -66,6 +67,15 @@ enum Cmds {
     /// (this command is only available on illumos)
     #[cfg(not(target_os = "illumos"))]
     VirtualHardware,
+
+    /// Print USDT probes in Omicron binaries.
+    Probes {
+        /// An optional filter applied to binary names.
+        ///
+        /// This is a simple substring match. Any binary with the filter as a
+        /// substring of its name will be examined for probes.
+        filter: Option<String>,
+    },
 }
 
 #[tokio::main]
@@ -91,6 +101,7 @@ async fn main() -> Result<()> {
         Cmds::Releng | Cmds::VerifyLibraries | Cmds::VirtualHardware => {
             anyhow::bail!("this command is only available on illumos");
         }
+        Cmds::Probes { filter } => usdt::print_probes(filter),
     }
 }
 
diff --git a/dev-tools/xtask/src/usdt.rs b/dev-tools/xtask/src/usdt.rs
new file mode 100644
index 0000000000..ad183656a4
--- /dev/null
+++ b/dev-tools/xtask/src/usdt.rs
@@ -0,0 +1,84 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+//! Xtask for printing USDT probes.
+
+use crate::load_workspace;
+use tabled::settings::Style;
+use tabled::Table;
+use tabled::Tabled;
+
+#[derive(Tabled)]
+#[tabled(rename_all = "SCREAMING_SNAKE_CASE")]
+struct Probe {
+    binary: String,
+    provider: String,
+    probe: String,
+    arguments: String,
+}
+
+pub(crate) fn print_probes(filter: Option<String>) -> anyhow::Result<()> {
+    const SKIP_ME: &[&str] = &["bootstrap", "xtask"];
+    const PATHS: &[&str] = &["release", "debug"];
+    let workspace = load_workspace()?;
+
+    // Find all local packages and any binaries they contain, and attemp to find
+    // contained DTrace probes in the object files.
+    let mut entries = Vec::new();
+    for bin_target in
+        workspace.workspace_packages().iter().flat_map(|package| {
+            package.targets.iter().filter(|target| {
+                target.is_bin()
+                    && !SKIP_ME.contains(&target.name.as_str())
+                    && filter
+                        .as_ref()
+                        .map(|filt| target.name.contains(filt))
+                        .unwrap_or(true)
+            })
+        })
+    {
+        let maybe_path = PATHS
+            .iter()
+            .filter_map(|p| {
+                let path =
+                    workspace.target_directory.join(p).join(&bin_target.name);
+                if path.exists() {
+                    Some(path)
+                } else {
+                    None
+                }
+            })
+            .next();
+        let Some(path) = maybe_path else {
+            continue;
+        };
+        match usdt::probe_records(&path) {
+            Ok(sections) => {
+                let all_providers = sections
+                    .into_iter()
+                    .flat_map(|section| section.providers.into_values());
+                for provider in all_providers {
+                    for probe in provider.probes.into_values() {
+                        let arguments =
+                            format!("[{}]", probe.arguments.join(","));
+                        entries.push(Probe {
+                            binary: bin_target.name.clone(),
+                            provider: provider.name.clone(),
+                            probe: probe.name,
+                            arguments,
+                        });
+                    }
+                }
+            }
+            Err(e) => {
+                eprintln!(
+                    "Failed to extract DTrace probes from '{path}': \
+                    {e}, or it may have zero probes"
+                );
+            }
+        }
+    }
+    println!("{}", Table::new(entries).with(Style::empty()));
+    Ok(())
+}