Add supertd_events Scuba logger

Summary:
Add a small Scuba wrapper module which logs to the supertd_events table. If the environment variable `SUPERTD_SCUBA_LOGFILE` is set, the samples are logged to that file instead of Scuba.

Includes a convenience macro `scuba!`. Instead of writing:

```
td_utils::supertd_events::sample_builder()
    .add("event", "COMMAND_SUCCESS")
    .add("count", foos_run + bars_launched)
    .add("data", serde_json::to_string(&serde_json::json!({
        "foos_run": foos_run,
        "bars_launched": bars_launched,
    })).unwrap())
    .log();
```

the macro allows us to write:

```
td_util::scuba!(
    event: "COMMAND_SUCCESS",
    count: foos_run + bars_launched,
    data: json!({
        "foos_run": foos_run,
        "bars_launched": bars_launched,
    })
);
```

The `scuba!` macro brings the `serde_json::json!` macro into scope for the `data` column, and passes the provided value through `serde_json::to_string(...)`. This only kicks in for the `data` column, which is intended to be event-specific properties (so that we don't bloat up the table with columns that are only populated by one event). This data can be used in Scuba queries or derived columns using `JSON_EXTRACT`.

Each sample automatically includes [several useful server and build data fields](https://www.internalfb.com/code/fbsource/[ec6f25826e22a953d23af7d59165e8002ff2ab55]/fbcode/common/rust/scuba/src/builder.rs?lines=188%2C191%2C194%2C197%2C202%2C206%2C210%2C214%2C218%2C222%2C231%2C238-239) via [`ScubaSampleBuilder::add_common_server_data`](https://www.internalfb.com/code/fbsource/[ec6f25826e22a953d23af7d59165e8002ff2ab55]/fbcode/common/rust/scuba/src/builder.rs?lines=245) and several Sandcastle variables (see `add_sandcastle_columns` in this diff).

See {D53788472} for more examples of how this might be used.

Reviewed By: Acesine

Differential Revision: D53689274

fbshipit-source-id: 2de6fec92d7e834c13e1ea7995a44e9900022ebd

btd/Cargo.toml

@@ -12,6 +12,7 @@ anyhow = "1.0"
 clap = {version = "4.1.4", features = ["derive"]}
 rayon = "1.7.0"
 derive_more = "0.99.3"
+fbinit = { workspace = true }
 glob = "0.3.0"
 itertools = "0.10.5"
 serde = {version = "1.0", features = ["derive"]}

btd/bin/main.rs

@@ -8,9 +8,9 @@
  */
 
 use td_util::cli::parse_args;
-use td_util::tracing::init_tracing;
 
-pub fn main() -> anyhow::Result<()> {
-    init_tracing();
+#[fbinit::main]
+pub fn main(fb: fbinit::FacebookInit) -> anyhow::Result<()> {
+    let _guard = td_util::init(fb);
     btd::main(parse_args()?)
 }

supertd/Cargo.toml

@@ -10,7 +10,7 @@ path = "bin/main.rs"
 [dependencies]
 anyhow = "1.0"
 clap = {version = "4.1.4", features = ["derive"]}
-fbinit = "0.1"
+fbinit = { workspace = true }
 
 td_util = {path = "../td_util"}
 btd = {path = "../btd"}

supertd/bin/main.rs

@@ -12,7 +12,6 @@ use clap::FromArgMatches;
 use clap::Parser;
 use fbinit::FacebookInit;
 use td_util::cli::get_args;
-use td_util::tracing::init_tracing;
 
 /// Generic binary for the pieces of the new target-determinator framework.
 #[allow(clippy::large_enum_variant)] // Only one instance, so not a big deal
@@ -34,10 +33,7 @@ enum Args {
 
 #[fbinit::main]
 pub fn main(fb: FacebookInit) -> anyhow::Result<()> {
-    // Avoid a warning in open source code
-    let _unused_in_oss = fb;
-
-    init_tracing();
+    let _guard = td_util::init(fb);
 
     let mut command = Args::command();
     if std::env::var_os("SUPERTD_IGNORE_EXTRA_ARGUMENTS") == Some("1".into()) {

targets/Cargo.toml

@@ -10,5 +10,6 @@ path = "bin/main.rs"
 [dependencies]
 anyhow = "1.0"
 clap = {version = "4.1.4", features = ["derive"]}
+fbinit = { workspace = true }
 
 td_util = {path = "../td_util"}

targets/bin/main.rs

@@ -8,9 +8,9 @@
  */
 
 use td_util::cli::parse_args;
-use td_util::tracing::init_tracing;
 
-pub fn main() -> anyhow::Result<()> {
-    init_tracing();
+#[fbinit::main]
+pub fn main(fb: fbinit::FacebookInit) -> anyhow::Result<()> {
+    let _guard = td_util::init(fb);
     targets::main(parse_args()?)
 }

td_util/Cargo.toml

@@ -9,10 +9,12 @@ argfile = "0.1.5"
 clap = {version = "4.1.4"}
 derive_more = "0.99.3"
 equivalent = "1.0.0"
+fbinit = { workspace = true }
 lazy_static = "1.4.0"
 static_interner.version = "0.1"
 # @oss-disable: static_interner.path = "../../buck2/shed/static_interner"
 static_interner.default-features = false
+scuba = { workspace = true }
 parse-display = "0.8.2"
 rayon = "1.6.1"
 serde = {version = "1.0", features = ["derive"]}

td_util/src/lib.rs

@@ -14,4 +14,18 @@ pub mod prelude;
 pub mod project;
 pub mod schedules;
 pub mod string;
+pub mod supertd_events;
 pub mod tracing;
+
+/// Initialize `tracing` and `supertd_events` Scuba client.
+///
+/// Returns a guard that flushes the Scuba client when dropped.
+///
+/// # Panics
+///
+/// Panics if environment variable `SUPERTD_SCUBA_LOGFILE` is set and the log
+/// file cannot be opened for writing.
+pub fn init(fb: fbinit::FacebookInit) -> supertd_events::ScubaClientGuard {
+    tracing::init_tracing();
+    supertd_events::init(fb)
+}

td_util/src/supertd_events.rs

@@ -0,0 +1,158 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under both the MIT license found in the
+ * LICENSE-MIT file in the root directory of this source tree and the Apache
+ * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+ * of this source tree.
+ */
+
+//! Simple interface for logging to the `supertd_events` dataset.
+
+use std::sync::OnceLock;
+
+use scuba::ScubaSampleBuilder;
+pub use serde_json;
+
+const SCUBA_DATASET: &str = "supertd_events";
+
+static BUILDER: OnceLock<ScubaSampleBuilder> = OnceLock::new();
+
+/// Initialize the Scuba client for the `supertd_events` dataset.
+///
+/// Returns a guard that flushes the Scuba client when dropped.
+///
+/// Expects `tracing` to be initialized.
+///
+/// If the environment variable `SUPERTD_SCUBA_LOGFILE` is set, then log to that
+/// filename instead of Scuba (useful for testing).
+///
+/// # Panics
+///
+/// Panics if `SUPERTD_SCUBA_LOGFILE` is set and the log file cannot be opened
+/// for writing.
+pub fn init(fb: fbinit::FacebookInit) -> ScubaClientGuard {
+    let mut builder = match std::env::var_os("SUPERTD_SCUBA_LOGFILE") {
+        None => ScubaSampleBuilder::new(fb, SCUBA_DATASET),
+        Some(path) => ScubaSampleBuilder::with_discard()
+            .with_log_file(path)
+            .unwrap(),
+    };
+    builder.add_common_server_data();
+    add_sandcastle_columns(&mut builder);
+    if BUILDER.set(builder).is_err() {
+        tracing::error!("supertd_events Scuba client initialized twice");
+    }
+    ScubaClientGuard(())
+}
+
+/// Log a sample to the `supertd_events` dataset.
+///
+/// The `event` column should be a distinct string for each source location
+/// logging an event.
+///
+/// The `data` column contains JSON-encoded data specific to that event (so that
+/// we do not inflate the number of columns in the Scuba table with properties
+/// populated by only one event). Use this data in derived columns or queries
+/// using `JSON_EXTRACT`.
+///
+/// If [`init`] has not been invoked, the sample will not be logged.
+///
+/// # Panics
+///
+/// Panics if [`serde_json::to_string`] fails for `data`.
+///
+/// # Examples
+///
+/// ```
+/// # use serde_json::json;
+/// # let foos_run = 10;
+/// # let bars_launched = 2;
+/// td_util::scuba!(
+///     event: "COMMAND_SUCCESS",
+///     count: foos_run + bars_launched,
+///     data: json!({
+///         "arbitrary": ["JSON", "object"],
+///         "foos_run": foos_run,
+///         "bars_launched": bars_launched,
+///     })
+/// );
+/// ```
+#[macro_export]
+macro_rules! scuba {
+    ( $($key:ident : $value:expr),* $(,)? ) => {
+        let mut builder = $crate::supertd_events::sample_builder();
+        $($crate::scuba! { @SET_FIELD(builder, $key, $value) })*
+        builder.log();
+    };
+    ( @SET_FIELD ( $builder:ident, data, $value:expr ) ) => {
+        // This unwrap should only fail for map keys with invalid UTF-8.
+        $builder.add("data", $crate::supertd_events::serde_json::to_string(&$value).unwrap());
+    };
+    ( @SET_FIELD ( $builder:ident, $key:ident, $value:expr ) ) => {
+        $builder.add(stringify!($key), $value);
+    };
+}
+
+/// Get the sample builder for the `supertd_events` dataset.
+///
+/// Most use cases should use the [`scuba!`] macro instead of this function, but
+/// this function can be used to access the underlying [`ScubaSampleBuilder`]
+/// (in order to do things like set the sampling rate).
+pub fn sample_builder() -> ScubaSampleBuilder {
+    BUILDER
+        .get()
+        .cloned()
+        .unwrap_or_else(ScubaSampleBuilder::with_discard)
+}
+
+fn add_sandcastle_columns(sample: &mut ScubaSampleBuilder) {
+    let Some(nexus_path) = std::env::var_os("SANDCASTLE_NEXUS") else {
+        return;
+    };
+    let nexus_path = std::path::Path::new(&nexus_path);
+    if !nexus_path.exists() {
+        return;
+    }
+    let variables_path = nexus_path.join("variables");
+    let variables = [
+        "SANDCASTLE_ALIAS_NAME",
+        "SANDCASTLE_ALIAS",
+        "SANDCASTLE_COMMAND_NAME",
+        "SANDCASTLE_INSTANCE_ID",
+        "SANDCASTLE_IS_DRY_RUN",
+        "SANDCASTLE_JOB_OWNER",
+        "SANDCASTLE_NONCE",
+        "SANDCASTLE_PHABRICATOR_DIFF_ID",
+        "SANDCASTLE_SCHEDULE_TYPE",
+        "SANDCASTLE_TYPE",
+        "SANDCASTLE_URL",
+        "SKYCASTLE_ACTION_ID",
+        "SKYCASTLE_JOB_ID",
+        "SKYCASTLE_WORKFLOW_RUN_ID",
+        "STEP_IDX",
+    ];
+    for var in variables {
+        let var_lowercase = var.to_ascii_lowercase();
+        if let Ok(value) = std::fs::read_to_string(variables_path.join(var)) {
+            sample.add(var_lowercase, value);
+        } else if let Ok(value) = std::env::var(var) {
+            sample.add(var_lowercase, value);
+        }
+    }
+}
+
+/// Flushes the `supertd_events` Scuba client when dropped.
+///
+/// Make sure this value is in scope for the duration of the program so that we
+/// flush the client upon program exit.
+#[must_use]
+pub struct ScubaClientGuard(());
+
+impl Drop for ScubaClientGuard {
+    fn drop(&mut self) {
+        if let Some(builder) = BUILDER.get() {
+            builder.flush(std::time::Duration::from_secs(5));
+        }
+    }
+}