1212// See the License for the specific language governing permissions and
1313// limitations under the License.
1414
15- use core:: default:: Default ;
16- use std:: env;
17- use std:: sync:: OnceLock ;
18-
1915use base64:: Engine ;
2016use base64:: prelude:: BASE64_STANDARD_NO_PAD ;
17+ use core:: default:: Default ;
18+ use ginepro:: LoadBalancedChannel ;
2119use hyper:: http:: Response ;
2220use nativelink_error:: { Code , ResultExt , make_err} ;
2321use nativelink_proto:: build:: bazel:: remote:: execution:: v2:: RequestMetadata ;
@@ -26,15 +24,19 @@ use opentelemetry::trace::{TraceContextExt, Tracer, TracerProvider};
2624use opentelemetry:: { KeyValue , global} ;
2725use opentelemetry_appender_tracing:: layer:: OpenTelemetryTracingBridge ;
2826use opentelemetry_http:: HeaderExtractor ;
29- use opentelemetry_otlp:: { LogExporter , MetricExporter , Protocol , SpanExporter , WithExportConfig } ;
27+ use opentelemetry_otlp:: {
28+ LogExporter , MetricExporter , Protocol , SpanExporter , WithExportConfig , WithTonicConfig ,
29+ } ;
3030use opentelemetry_sdk:: Resource ;
3131use opentelemetry_sdk:: logs:: SdkLoggerProvider ;
3232use opentelemetry_sdk:: metrics:: SdkMeterProvider ;
3333use opentelemetry_sdk:: propagation:: { BaggagePropagator , TraceContextPropagator } ;
3434use opentelemetry_sdk:: trace:: SdkTracerProvider ;
3535use opentelemetry_semantic_conventions:: attribute:: ENDUSER_ID ;
3636use prost:: Message ;
37- use tracing:: debug;
37+ use std:: env;
38+ use std:: sync:: { OnceLock } ;
39+ use tracing:: { debug, info} ;
3840use tracing:: metadata:: LevelFilter ;
3941use tracing_opentelemetry:: { MetricsLayer , layer} ;
4042use tracing_subscriber:: filter:: Directive ;
@@ -104,7 +106,7 @@ fn tracing_stdout_layer() -> impl Layer<Registry> {
104106///
105107/// Returns `Err` if logging was already initialized or if the exporters can't
106108/// be initialized.
107- pub fn init_tracing ( ) -> Result < ( ) , nativelink_error:: Error > {
109+ pub async fn init_tracing ( ) -> Result < ( ) , nativelink_error:: Error > {
108110 static INITIALIZED : OnceLock < ( ) > = OnceLock :: new ( ) ;
109111
110112 if INITIALIZED . get ( ) . is_some ( ) {
@@ -129,13 +131,18 @@ pub fn init_tracing() -> Result<(), nativelink_error::Error> {
129131 ] ) ;
130132 global:: set_text_map_propagator ( propagator) ;
131133
134+ let maybe_channel = maybe_load_balanced_channel ( ) . await ;
135+
132136 // Logs
137+ let mut log_exporter_builder = LogExporter :: builder ( ) . with_tonic ( ) ;
138+ if let Some ( channel) = maybe_channel. clone ( ) {
139+ log_exporter_builder = log_exporter_builder. with_channel ( channel. into ( ) ) ;
140+ }
133141 let otlp_log_layer = OpenTelemetryTracingBridge :: new (
134142 & SdkLoggerProvider :: builder ( )
135143 . with_resource ( resource. clone ( ) )
136144 . with_batch_exporter (
137- LogExporter :: builder ( )
138- . with_tonic ( )
145+ log_exporter_builder
139146 . with_protocol ( Protocol :: Grpc )
140147 . build ( )
141148 . map_err ( |e| make_err ! ( Code :: Internal , "{e}" ) )
@@ -146,13 +153,16 @@ pub fn init_tracing() -> Result<(), nativelink_error::Error> {
146153 . with_filter ( otlp_filter ( ) ) ;
147154
148155 // Traces
156+ let mut span_exporter_builder = SpanExporter :: builder ( ) . with_tonic ( ) ;
157+ if let Some ( channel) = maybe_channel. clone ( ) {
158+ span_exporter_builder = span_exporter_builder. with_channel ( channel. into ( ) ) ;
159+ }
149160 let otlp_trace_layer = layer ( )
150161 . with_tracer (
151162 SdkTracerProvider :: builder ( )
152163 . with_resource ( resource. clone ( ) )
153164 . with_batch_exporter (
154- SpanExporter :: builder ( )
155- . with_tonic ( )
165+ span_exporter_builder
156166 . with_protocol ( Protocol :: Grpc )
157167 . build ( )
158168 . map_err ( |e| make_err ! ( Code :: Internal , "{e}" ) )
@@ -164,11 +174,14 @@ pub fn init_tracing() -> Result<(), nativelink_error::Error> {
164174 . with_filter ( otlp_filter ( ) ) ;
165175
166176 // Metrics
177+ let mut metric_exporter_builder = MetricExporter :: builder ( ) . with_tonic ( ) ;
178+ if let Some ( channel) = maybe_channel {
179+ metric_exporter_builder = metric_exporter_builder. with_channel ( channel. into ( ) ) ;
180+ }
167181 let meter_provider = SdkMeterProvider :: builder ( )
168182 . with_resource ( resource)
169183 . with_periodic_exporter (
170- MetricExporter :: builder ( )
171- . with_tonic ( )
184+ metric_exporter_builder
172185 . with_protocol ( Protocol :: Grpc )
173186 . build ( )
174187 . map_err ( |e| make_err ! ( Code :: Internal , "{e}" ) )
@@ -192,6 +205,36 @@ pub fn init_tracing() -> Result<(), nativelink_error::Error> {
192205 Ok ( ( ) )
193206}
194207
208+ const NL_OTEL_ENDPOINT : & str = "NL_OTEL_ENDPOINT" ;
209+
210+ async fn maybe_load_balanced_channel ( ) -> Option < LoadBalancedChannel > {
211+ match env:: var ( NL_OTEL_ENDPOINT ) {
212+ Ok ( endpoint) => {
213+ let url = Url :: parse ( endpoint. as_str ( ) ) . map_err ( |e| {
214+ make_err ! ( Code :: Internal , "Unable to parse endpoint {endpoint}: {e:?}" )
215+ } ) . unwrap ( ) ;
216+
217+ let host = url
218+ . host ( )
219+ . err_tip ( || format ! ( "Unable to get host from endpoint {endpoint}" ) )
220+ . unwrap ( ) ;
221+ let port = url
222+ . port ( )
223+ . err_tip ( || format ! ( "Unable to get port from endpoint {endpoint}" ) )
224+ . unwrap ( ) ;
225+
226+ Some (
227+ LoadBalancedChannel :: builder ( ( host. to_string ( ) , port) )
228+ . channel ( )
229+ . await
230+ . map_err ( |e| make_err ! ( Code :: Internal , "Invalid hostname '{endpoint}': {e}" ) )
231+ . unwrap ( ) ,
232+ )
233+ }
234+ Err ( _) => None ,
235+ }
236+ }
237+
195238/// Custom metadata key field for Bazel metadata.
196239const BAZEL_METADATA_KEY : & str = "bazel.metadata" ;
197240
@@ -202,6 +245,7 @@ const BAZEL_REQUESTMETADATA_HEADER: &str = "build.bazel.remote.execution.v2.requ
202245
203246use opentelemetry:: baggage:: BaggageExt ;
204247use opentelemetry:: context:: FutureExt ;
248+ use url:: Url ;
205249
206250#[ derive( Debug , Clone ) ]
207251pub struct OtlpMiddleware < S > {
0 commit comments