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,14 +24,18 @@ 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 std:: env;
38+ use std:: sync:: { OnceLock } ;
3739use tracing:: debug;
3840use tracing:: metadata:: LevelFilter ;
3941use tracing_opentelemetry:: { MetricsLayer , layer} ;
@@ -103,7 +105,7 @@ fn tracing_stdout_layer() -> impl Layer<Registry> {
103105///
104106/// Returns `Err` if logging was already initialized or if the exporters can't
105107/// be initialized.
106- pub fn init_tracing ( ) -> Result < ( ) , nativelink_error:: Error > {
108+ pub async fn init_tracing ( ) -> Result < ( ) , nativelink_error:: Error > {
107109 static INITIALIZED : OnceLock < ( ) > = OnceLock :: new ( ) ;
108110
109111 if INITIALIZED . get ( ) . is_some ( ) {
@@ -128,13 +130,18 @@ pub fn init_tracing() -> Result<(), nativelink_error::Error> {
128130 ] ) ;
129131 global:: set_text_map_propagator ( propagator) ;
130132
133+ let maybe_channel = maybe_load_balanced_channel ( ) . await ;
134+
131135 // Logs
136+ let mut log_exporter_builder = LogExporter :: builder ( ) . with_tonic ( ) ;
137+ if let Some ( channel) = maybe_channel. clone ( ) {
138+ log_exporter_builder = log_exporter_builder. with_channel ( channel. into ( ) ) ;
139+ }
132140 let otlp_log_layer = OpenTelemetryTracingBridge :: new (
133141 & SdkLoggerProvider :: builder ( )
134142 . with_resource ( resource. clone ( ) )
135143 . with_batch_exporter (
136- LogExporter :: builder ( )
137- . with_tonic ( )
144+ log_exporter_builder
138145 . with_protocol ( Protocol :: Grpc )
139146 . build ( )
140147 . map_err ( |e| make_err ! ( Code :: Internal , "{e}" ) )
@@ -145,13 +152,16 @@ pub fn init_tracing() -> Result<(), nativelink_error::Error> {
145152 . with_filter ( otlp_filter ( ) ) ;
146153
147154 // Traces
155+ let mut span_exporter_builder = SpanExporter :: builder ( ) . with_tonic ( ) ;
156+ if let Some ( channel) = maybe_channel. clone ( ) {
157+ span_exporter_builder = span_exporter_builder. with_channel ( channel. into ( ) ) ;
158+ }
148159 let otlp_trace_layer = layer ( )
149160 . with_tracer (
150161 SdkTracerProvider :: builder ( )
151162 . with_resource ( resource. clone ( ) )
152163 . with_batch_exporter (
153- SpanExporter :: builder ( )
154- . with_tonic ( )
164+ span_exporter_builder
155165 . with_protocol ( Protocol :: Grpc )
156166 . build ( )
157167 . map_err ( |e| make_err ! ( Code :: Internal , "{e}" ) )
@@ -163,11 +173,14 @@ pub fn init_tracing() -> Result<(), nativelink_error::Error> {
163173 . with_filter ( otlp_filter ( ) ) ;
164174
165175 // Metrics
176+ let mut metric_exporter_builder = MetricExporter :: builder ( ) . with_tonic ( ) ;
177+ if let Some ( channel) = maybe_channel {
178+ metric_exporter_builder = metric_exporter_builder. with_channel ( channel. into ( ) ) ;
179+ }
166180 let meter_provider = SdkMeterProvider :: builder ( )
167181 . with_resource ( resource)
168182 . with_periodic_exporter (
169- MetricExporter :: builder ( )
170- . with_tonic ( )
183+ metric_exporter_builder
171184 . with_protocol ( Protocol :: Grpc )
172185 . build ( )
173186 . map_err ( |e| make_err ! ( Code :: Internal , "{e}" ) )
@@ -191,6 +204,36 @@ pub fn init_tracing() -> Result<(), nativelink_error::Error> {
191204 Ok ( ( ) )
192205}
193206
207+ const NL_OTEL_ENDPOINT : & str = "NL_OTEL_ENDPOINT" ;
208+
209+ async fn maybe_load_balanced_channel ( ) -> Option < LoadBalancedChannel > {
210+ match env:: var ( NL_OTEL_ENDPOINT ) {
211+ Ok ( endpoint) => {
212+ let url = Url :: parse ( endpoint. as_str ( ) ) . map_err ( |e| {
213+ make_err ! ( Code :: Internal , "Unable to parse endpoint {endpoint}: {e:?}" )
214+ } ) . unwrap ( ) ;
215+
216+ let host = url
217+ . host ( )
218+ . err_tip ( || format ! ( "Unable to get host from endpoint {endpoint}" ) )
219+ . unwrap ( ) ;
220+ let port = url
221+ . port ( )
222+ . err_tip ( || format ! ( "Unable to get port from endpoint {endpoint}" ) )
223+ . unwrap ( ) ;
224+
225+ Some (
226+ LoadBalancedChannel :: builder ( ( host. to_string ( ) , port) )
227+ . channel ( )
228+ . await
229+ . map_err ( |e| make_err ! ( Code :: Internal , "Invalid hostname '{endpoint}': {e}" ) )
230+ . unwrap ( ) ,
231+ )
232+ }
233+ Err ( _) => None ,
234+ }
235+ }
236+
194237/// Custom metadata key field for Bazel metadata.
195238const BAZEL_METADATA_KEY : & str = "bazel.metadata" ;
196239
@@ -201,6 +244,7 @@ const BAZEL_REQUESTMETADATA_HEADER: &str = "build.bazel.remote.execution.v2.requ
201244
202245use opentelemetry:: baggage:: BaggageExt ;
203246use opentelemetry:: context:: FutureExt ;
247+ use url:: Url ;
204248
205249#[ derive( Debug , Clone ) ]
206250pub struct OtlpMiddleware < S > {
0 commit comments