custom event metadata from buckconfig

dont modify pool.rs

cleanup

Author: KlotzAndrew

app/buck2_common/src/legacy_configs.rs

@@ -12,6 +12,7 @@
 //! .buckconfig files as configuration)
 
 mod access;
+pub use access::parse_buckconfig_metadata;
 mod aggregator;
 pub mod args;
 pub mod cells;

app/buck2_common/src/legacy_configs/access.rs

@@ -12,6 +12,8 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use buck2_error::BuckErrorContext;
+use buck2_hash::StdBuckHashMap;
+use buck2_util::env_vars::substitute_env_vars;
 use gazebo::eq_chain;
 
 use crate::legacy_configs::configs::ConfigValue;
@@ -21,6 +23,26 @@ use crate::legacy_configs::configs::LegacyBuckConfigValue;
 use crate::legacy_configs::key::BuckconfigKeyRef;
 use crate::legacy_configs::view::LegacyBuckConfigView;
 
+/// Read the `[buck2_metadata]` section from a `LegacyBuckConfig` and resolve any `$VAR`
+/// references. Entries whose env vars are not set are skipped with a warning.
+pub fn parse_buckconfig_metadata(config: &LegacyBuckConfig) -> StdBuckHashMap<String, String> {
+    let mut map = StdBuckHashMap::default();
+    let Some(section) = config.get_section("buck2_metadata") else {
+        return map;
+    };
+    for (key, value) in section.iter() {
+        match substitute_env_vars(value.as_str()) {
+            Ok(resolved) => {
+                map.insert(key.to_owned(), resolved);
+            }
+            Err(e) => {
+                tracing::warn!("Skipping [buck2_metadata] key `{}`: {:#}", key, e);
+            }
+        }
+    }
+    map
+}
+
 impl LegacyBuckConfigView for &LegacyBuckConfig {
     fn get(&mut self, key: BuckconfigKeyRef) -> buck2_error::Result<Option<Arc<str>>> {
         Ok(LegacyBuckConfig::get(self, key).map(|v| v.to_owned().into()))

app/buck2_events/src/metadata.rs

@@ -19,6 +19,16 @@ use buck2_wrapper_common::BUCK2_WRAPPER_ENV_VAR;
 
 use crate::daemon_id::DaemonId;
 
+/// Collects metadata from the current binary and environment, merged with any extras, suitable for telemetry purposes.
+pub fn collect_with_extras(
+    daemon: &DaemonId,
+    extras: &StdBuckHashMap<String, String>,
+) -> StdBuckHashMap<String, String> {
+    let mut map = collect(daemon);
+    map.extend(extras.iter().map(|(k, v)| (k.clone(), v.clone())));
+    map
+}
+
 /// Collects metadata from the current binary and environment and writes it as map, suitable for telemetry purposes.
 pub fn collect(daemon: &DaemonId) -> StdBuckHashMap<String, String> {
     facebook_only();

app/buck2_server/src/ctx.rs

@@ -1123,7 +1123,10 @@ impl ServerCommandContextTrait for ServerCommandContext<'_> {
         // Facebook only: metadata collection for Scribe writes
         facebook_only();
 
-        let mut metadata = metadata::collect(&self.base_context.daemon.daemon_id);
+        let mut metadata = metadata::collect_with_extras(
+            &self.base_context.daemon.daemon_id,
+            &self.base_context.daemon.buckconfig_metadata,
+        );
 
         metadata.insert(
             "io_provider".to_owned(),

app/buck2_server/src/daemon/state.rs

@@ -27,6 +27,7 @@ use buck2_common::init::Timeout;
 use buck2_common::invocation_paths::InvocationPaths;
 use buck2_common::io::IoProvider;
 use buck2_common::legacy_configs::cells::BuckConfigBasedCells;
+use buck2_common::legacy_configs::parse_buckconfig_metadata;
 use buck2_common::legacy_configs::key::BuckconfigKeyRef;
 use buck2_common::sqlite::sqlite_db::SqliteDb;
 use buck2_common::sqlite::sqlite_db::SqliteIdentity;
@@ -208,6 +209,8 @@ pub struct DaemonStateData {
     /// Semaphores for running actions locally. These need to be shared across commands.
     #[allocative(skip)]
     pub named_semaphores_for_run_actions: Arc<NamedSemaphores>,
+
+    pub buckconfig_metadata: StdBuckHashMap<String, String>,
 }
 
 impl DaemonStateData {
@@ -741,6 +744,7 @@ impl DaemonState {
                 incremental_db_state,
                 daemon_id: daemon_id.dupe(),
                 named_semaphores_for_run_actions: Arc::new(NamedSemaphores::new()),
+                buckconfig_metadata: parse_buckconfig_metadata(root_config),
             }))
         };
         let daemon_listener_span = tracing::Span::current();

app/buck2_util/Cargo.toml

@@ -12,6 +12,8 @@ version = "0.1.0"
 [dependencies]
 allocative = { workspace = true }
 blake3 = { workspace = true }
+once_cell = { workspace = true }
+regex = { workspace = true }
 dupe = { workspace = true }
 futures = { workspace = true }
 pagable = { workspace = true }

app/buck2_util/src/env_vars.rs

@@ -0,0 +1,84 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is dual-licensed under either the MIT license found in the
+ * LICENSE-MIT file in the root directory of this source tree or the Apache
+ * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+ * of this source tree. You may select, at your option, one of the
+ * above-listed licenses.
+ */
+
+use std::env::VarError;
+
+use once_cell::sync::Lazy;
+use regex::Regex;
+
+/// Replace occurrences of `$FOO` in a string with the value of the env var `$FOO`.
+/// Returns an error if any referenced variable is not set.
+pub fn substitute_env_vars(s: &str) -> buck2_error::Result<String> {
+    substitute_env_vars_impl(s, |v| std::env::var(v))
+}
+
+pub fn substitute_env_vars_impl(
+    s: &str,
+    getter: impl Fn(&str) -> Result<String, VarError>,
+) -> buck2_error::Result<String> {
+    static ENV_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new("\\$[a-zA-Z_][a-zA-Z_0-9]*").unwrap());
+
+    let mut out = String::with_capacity(s.len());
+    let mut last_idx = 0;
+
+    for mat in ENV_REGEX.find_iter(s) {
+        out.push_str(&s[last_idx..mat.start()]);
+        let var = &mat.as_str()[1..];
+        let val = getter(var).map_err(|e| {
+            buck2_error::buck2_error!(
+                buck2_error::ErrorTag::Environment,
+                "Error substituting `{}`: {}",
+                mat.as_str(),
+                e
+            )
+        })?;
+        out.push_str(&val);
+        last_idx = mat.end();
+    }
+
+    if last_idx < s.len() {
+        out.push_str(&s[last_idx..s.len()]);
+    }
+
+    Ok(out)
+}
+
+#[cfg(test)]
+mod tests {
+    use std::env::VarError;
+
+    use super::*;
+
+    #[test]
+    fn test_substitute_env_vars() {
+        let getter = |s: &str| match s {
+            "FOO" => Ok("foo_value".to_owned()),
+            "BAR" => Ok("bar_value".to_owned()),
+            "BAZ" => Err(VarError::NotPresent),
+            _ => panic!("Unexpected"),
+        };
+
+        assert_eq!(
+            substitute_env_vars_impl("$FOO", getter).unwrap(),
+            "foo_value"
+        );
+        assert_eq!(
+            substitute_env_vars_impl("$FOO$BAR", getter).unwrap(),
+            "foo_valuebar_value"
+        );
+        assert_eq!(
+            substitute_env_vars_impl("some$FOO.bar", getter).unwrap(),
+            "somefoo_value.bar"
+        );
+        assert_eq!(substitute_env_vars_impl("foo", getter).unwrap(), "foo");
+        assert_eq!(substitute_env_vars_impl("FOO", getter).unwrap(), "FOO");
+        assert!(substitute_env_vars_impl("$FOO$BAZ", getter).is_err());
+    }
+}

app/buck2_util/src/lib.rs

@@ -14,6 +14,7 @@
 #![feature(used_with_arg)]
 
 pub mod arc_str;
+pub mod env_vars;
 pub mod async_move_clone;
 pub mod commas;
 pub mod cycle_detector;