Introduce tests for load-balanced channel.

Author: dkostyrev

Cargo.lock

@@ -137,6 +137,7 @@ rust_test_suite(
         "@crates//:mock_instant",
         "@crates//:opentelemetry",
         "@crates//:opentelemetry-http",
+        "@crates//:opentelemetry_sdk",
         "@crates//:parking_lot",
         "@crates//:pretty_assertions",
         "@crates//:rand",
@@ -163,10 +164,13 @@ rust_test(
     ],
     deps = [
         "@crates//:http-body-util",
+        "@crates//:opentelemetry-proto",
+        "@crates//:opentelemetry_sdk",
         "@crates//:pretty_assertions",
         "@crates//:rand",
         "@crates//:serde_json",
         "@crates//:tempfile",
+        "@crates//:tokio-stream",
         "@crates//:tracing-test",
     ],
 )

nativelink-util/BUILD.bazel

@@ -98,6 +98,13 @@ walkdir = { version = "2.5.0", default-features = false }
 
 [dev-dependencies]
 nativelink-macro = { path = "../nativelink-macro" }
+opentelemetry-proto = { version = "0.29.0", features = [
+  "gen-tonic",
+  "metrics",
+] }
+opentelemetry_sdk = { version = "0.29.0", features = [
+  "testing"
+] }
 
 axum = { version = "0.8.3", default-features = false }
 http-body-util = { version = "0.1.3", default-features = false }

nativelink-util/Cargo.toml

@@ -383,3 +383,217 @@ impl<S> tower::Layer<S> for OtlpLayer {
         OtlpMiddleware::new(service, self.identity_required)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use core::convert::Infallible;
+    use core::future::Future;
+    use core::pin::Pin;
+    use core::task::{Context, Poll};
+    use std::sync::Arc;
+
+    use nativelink_macro::nativelink_test;
+    use opentelemetry::metrics::MeterProvider as _;
+    use opentelemetry_proto::tonic::collector::metrics::v1::{
+        ExportMetricsServiceRequest, ExportMetricsServiceResponse,
+    };
+    use opentelemetry_sdk::metrics::SdkMeterProvider;
+    use parking_lot::Mutex;
+
+    use super::*;
+
+    // Serialize env-var mutations across tests: env vars are process-global,
+    // so concurrent writes would be a data race on Unix.
+    //
+    // Serialize env-var mutations across tests: env vars are process-global,
+    // so concurrent writes would be a data race on Unix.  The guard is always
+    // scoped so it drops before any `.await`, which satisfies
+    // `clippy::await_holding_lock`.
+    static ENV_LOCK: Mutex<()> = Mutex::new(());
+
+    #[nativelink_test("crate")]
+    async fn channel_absent_when_env_not_set() {
+        {
+            let _guard = ENV_LOCK.lock();
+            // SAFETY: ENV_LOCK serializes all env-var writes in this test module.
+            unsafe { env::remove_var(NL_OTEL_ENDPOINT) };
+        }
+        assert!(
+            maybe_load_balanced_channel().await.is_none(),
+            "Expected None when {NL_OTEL_ENDPOINT} is unset"
+        );
+    }
+
+    #[nativelink_test("crate")]
+    async fn channel_present_when_valid_endpoint_set() {
+        {
+            let _guard = ENV_LOCK.lock();
+            // SAFETY: ENV_LOCK serializes all env-var writes in this test module.
+            unsafe { env::set_var(NL_OTEL_ENDPOINT, "http://localhost:4317") };
+        }
+        let result = maybe_load_balanced_channel().await;
+        {
+            let _guard = ENV_LOCK.lock();
+            // SAFETY: ENV_LOCK serializes all env-var writes in this test module.
+            unsafe { env::remove_var(NL_OTEL_ENDPOINT) };
+        }
+        assert!(
+            result.is_some(),
+            "Expected Some(channel) when {NL_OTEL_ENDPOINT} points to a valid URL"
+        );
+    }
+
+    // A minimal in-process OTLP metrics collector backed by tonic 0.13.
+    //
+    // `opentelemetry-proto`'s generated `MetricsServiceServer` depends on tonic
+    // 0.12, which is incompatible with the tonic 0.13 used by the rest of the
+    // crate.  To avoid the version conflict we implement the gRPC service
+    // manually: re-use only the prost message types from `opentelemetry-proto`
+    // and wire up the routing with tonic 0.13 primitives directly.
+    #[derive(Clone)]
+    struct TestMetricsService {
+        received: Arc<Mutex<Vec<ExportMetricsServiceRequest>>>,
+    }
+
+    const METRICS_EXPORT_PATH: &str =
+        "/opentelemetry.proto.collector.metrics.v1.MetricsService/Export";
+
+    impl tower::Service<hyper::http::Request<tonic::body::Body>> for TestMetricsService {
+        type Response = Response<tonic::body::Body>;
+        type Error = Infallible;
+        type Future =
+            Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
+
+        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn call(&mut self, req: hyper::http::Request<tonic::body::Body>) -> Self::Future {
+            let received = self.received.clone();
+            Box::pin(async move {
+                if req.uri().path() == METRICS_EXPORT_PATH {
+                    let export_svc = tower::service_fn(
+                        move |request: tonic::Request<ExportMetricsServiceRequest>| {
+                            let received = received.clone();
+                            async move {
+                                received.lock().push(request.into_inner());
+                                Ok::<tonic::Response<ExportMetricsServiceResponse>, tonic::Status>(
+                                    tonic::Response::new(ExportMetricsServiceResponse::default()),
+                                )
+                            }
+                        },
+                    );
+                    // ProstCodec<T, U>: Encode = T (sent to client),
+                    //                   Decode = U (received from client).
+                    let mut grpc = tonic::server::Grpc::new(tonic::codec::ProstCodec::<
+                        ExportMetricsServiceResponse,
+                        ExportMetricsServiceRequest,
+                    >::default());
+                    Ok(grpc.unary(export_svc, req).await)
+                } else {
+                    let mut resp = Response::new(tonic::body::Body::empty());
+                    resp.headers_mut().insert(
+                        tonic::Status::GRPC_STATUS,
+                        // `Code` in the outer scope is nativelink_error::Code,
+                        // so the `tonic::` prefix is necessary here.
+                        #[expect(
+                            unused_qualifications,
+                            reason = "tonic::Code differs from the nativelink_error::Code in outer scope"
+                        )]
+                        (tonic::Code::Unimplemented as i32).into(),
+                    );
+                    Ok(resp)
+                }
+            })
+        }
+    }
+
+    impl tonic::server::NamedService for TestMetricsService {
+        const NAME: &'static str = "opentelemetry.proto.collector.metrics.v1.MetricsService";
+    }
+
+    #[expect(
+        clippy::disallowed_methods,
+        reason = "`tokio::test` uses `tokio::runtime::Builder::new_multi_thread` and \
+                  `tokio::runtime::Runtime::block_on` internally"
+    )]
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn metrics_pushed_via_load_balanced_channel() {
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind test listener");
+        let port = listener.local_addr().expect("local_addr").port();
+        let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
+
+        let received: Arc<Mutex<Vec<ExportMetricsServiceRequest>>> = Arc::new(Mutex::new(vec![]));
+        let svc = TestMetricsService {
+            received: received.clone(),
+        };
+
+        crate::background_spawn!("otlp_test_server", async move {
+            let _ = tonic::transport::Server::builder()
+                .add_service(svc)
+                .serve_with_incoming(incoming)
+                .await;
+        });
+
+        {
+            let _guard = ENV_LOCK.lock();
+            // SAFETY: ENV_LOCK serializes all env-var writes in this test module.
+            unsafe { env::set_var(NL_OTEL_ENDPOINT, format!("http://127.0.0.1:{port}")) };
+        }
+        let channel = maybe_load_balanced_channel()
+            .await
+            .expect("channel should be Some when NL_OTEL_ENDPOINT is set");
+        {
+            let _guard = ENV_LOCK.lock();
+            // SAFETY: ENV_LOCK serializes all env-var writes in this test module.
+            unsafe { env::remove_var(NL_OTEL_ENDPOINT) };
+        }
+
+        let exporter = MetricExporter::builder()
+            .with_tonic()
+            .with_channel(channel.into())
+            .with_protocol(Protocol::Grpc)
+            .build()
+            .expect("MetricExporter");
+
+        let meter_provider = SdkMeterProvider::builder()
+            .with_periodic_exporter(exporter)
+            .build();
+
+        let meter = meter_provider.meter("nativelink_test");
+        let counter = meter
+            .u64_counter("test.pushes")
+            .with_description("Counter pushed through the load-balanced channel")
+            .build();
+        counter.add(42, &[KeyValue::new("operation", "cache_read")]);
+
+        meter_provider
+            .force_flush()
+            .expect("force_flush should succeed");
+
+        {
+            let batches = received.lock();
+            assert!(
+                !batches.is_empty(),
+                "Expected at least one OTLP ExportMetricsServiceRequest at the test server"
+            );
+
+            let metric_names: Vec<&str> = batches
+                .iter()
+                .flat_map(|req| &req.resource_metrics)
+                .flat_map(|rm| &rm.scope_metrics)
+                .flat_map(|sm| &sm.metrics)
+                .map(|m| m.name.as_str())
+                .collect();
+
+            assert!(
+                metric_names.contains(&"test.pushes"),
+                "Counter 'test.pushes' should have been delivered via gRPC, got: {metric_names:?}"
+            );
+        }
+
+        meter_provider.shutdown().expect("shutdown");
+    }
+}

nativelink-util/src/telemetry.rs

@@ -8,8 +8,10 @@ use hyper::{Request, Response, StatusCode, Uri};
 use nativelink_macro::nativelink_test;
 use nativelink_util::telemetry::ClientHeaders;
 use opentelemetry::baggage::BaggageExt;
+use opentelemetry::metrics::MeterProvider as _;
 use opentelemetry::{Context, KeyValue, global};
 use opentelemetry_http::HeaderExtractor as OTELHeaderExtractor;
+use opentelemetry_sdk::metrics::{InMemoryMetricExporter, PeriodicReader, SdkMeterProvider};
 use tonic::body::Body;
 use tower::{Service, ServiceBuilder, ServiceExt};
 use tracing::{debug, error, warn};
@@ -153,3 +155,43 @@ async fn oltp_logs_with_headers() -> Result<(), Box<dyn core::error::Error>> {
 
     Ok(())
 }
+
+#[nativelink_test]
+async fn metrics_are_tracked() -> Result<(), Box<dyn core::error::Error>> {
+    let exporter = InMemoryMetricExporter::default();
+    let meter_provider = SdkMeterProvider::builder()
+        .with_reader(PeriodicReader::builder(exporter.clone()).build())
+        .build();
+
+    let meter = meter_provider.meter("nativelink_test");
+    let counter = meter
+        .u64_counter("test.operations")
+        .with_description("Total test operations")
+        .build();
+
+    counter.add(3, &[KeyValue::new("operation", "cache_read")]);
+    counter.add(7, &[KeyValue::new("operation", "cache_write")]);
+
+    meter_provider.force_flush()?;
+
+    let finished = exporter.get_finished_metrics()?;
+    assert!(
+        !finished.is_empty(),
+        "Expected at least one ResourceMetrics batch after force_flush"
+    );
+
+    let recorded_names: Vec<&str> = finished
+        .iter()
+        .flat_map(|rm| &rm.scope_metrics)
+        .flat_map(|sm| &sm.metrics)
+        .map(|m| m.name.as_ref())
+        .collect();
+
+    assert!(
+        recorded_names.contains(&"test.operations"),
+        "Counter 'test.operations' should appear in exported metrics, got: {recorded_names:?}"
+    );
+
+    meter_provider.shutdown()?;
+    Ok(())
+}