From cb646ebc92c9ff652e22c8aa66cf44251d533e5a Mon Sep 17 00:00:00 2001
From: Subhra264 <chakrabortysubhradeep556@gmail.com>
Date: Fri, 12 Jul 2024 16:48:28 +0530
Subject: [PATCH] feat: Introduce otel grpc endpoint to export traces (#3963)

Fixes #3940

## New ENVs:
- `ZO_TRACING_GRPC_ORGANIZATION` - The name of the organization used to
ingest traces
- `ZO_TRACING_GRPC_STREAM_NAME` - The stream name to use for trace
export.
- `OTEL_OTLP_GRPC_ENDPOINT` - The grpc endpoint to use for trace export.
E.g. - `grpc://localhost:5081`. If this is specified,
`OTEL_OTLP_HTTP_ENDPOINT` is ignored.

`ZO_TRACING_HEADER_VALUE` can be used for authorization token.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added configuration options for GRPC tracing, including URL, type, and
header settings.
- **Improvements**
- Enhanced tracer exporter logic for better configurability based on
tracing type.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 src/config/src/config.rs | 14 ++++++++++++
 src/main.rs              | 48 +++++++++++++++++++++++++++++++---------
 2 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/src/config/src/config.rs b/src/config/src/config.rs
index db749a7159..e4d43e44f3 100644
--- a/src/config/src/config.rs
+++ b/src/config/src/config.rs
@@ -574,6 +574,20 @@ pub struct Common {
     pub tracing_search_enabled: bool,
     #[env_config(name = "OTEL_OTLP_HTTP_ENDPOINT", default = "")]
     pub otel_otlp_url: String,
+    #[env_config(name = "OTEL_OTLP_GRPC_ENDPOINT", default = "")]
+    pub otel_otlp_grpc_url: String,
+    #[env_config(
+        name = "ZO_TRACING_GRPC_ORGANIZATION",
+        default = "",
+        help = "Used in metadata when exporting traces to grpc endpoint."
+    )]
+    pub tracing_grpc_header_org: String,
+    #[env_config(
+        name = "ZO_TRACING_GRPC_STREAM_NAME",
+        default = "",
+        help = "Used in metadata when exporting traces to grpc endpoint."
+    )]
+    pub tracing_grpc_header_stream_name: String,
     #[env_config(name = "ZO_TRACING_HEADER_KEY", default = "Authorization")]
     pub tracing_header_key: String,
     #[env_config(
diff --git a/src/main.rs b/src/main.rs
index 89c5d90f49..cde369efd8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -64,7 +64,10 @@ use pyroscope::PyroscopeAgent;
 #[cfg(feature = "profiling")]
 use pyroscope_pprofrs::{pprof_backend, PprofConfig};
 use tokio::sync::oneshot;
-use tonic::codec::CompressionEncoding;
+use tonic::{
+    codec::CompressionEncoding,
+    metadata::{MetadataKey, MetadataMap, MetadataValue},
+};
 use tracing_appender::non_blocking::WorkerGuard;
 use tracing_subscriber::Registry;
 
@@ -630,14 +633,14 @@ pub(crate) fn setup_logs() -> tracing_appender::non_blocking::WorkerGuard {
 fn enable_tracing() -> Result<(), anyhow::Error> {
     let cfg = get_config();
     opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new());
-    let mut headers = HashMap::new();
-    headers.insert(
-        cfg.common.tracing_header_key.clone(),
-        cfg.common.tracing_header_value.clone(),
-    );
-    let tracer = opentelemetry_otlp::new_pipeline()
-        .tracing()
-        .with_exporter(
+    let tracer = opentelemetry_otlp::new_pipeline().tracing();
+    let tracer = if cfg.common.otel_otlp_grpc_url.is_empty() {
+        tracer.with_exporter({
+            let mut headers = HashMap::new();
+            headers.insert(
+                cfg.common.tracing_header_key.clone(),
+                cfg.common.tracing_header_value.clone(),
+            );
             opentelemetry_otlp::new_exporter()
                 .http()
                 .with_http_client(
@@ -646,8 +649,31 @@ fn enable_tracing() -> Result<(), anyhow::Error> {
                         .build()?,
                 )
                 .with_endpoint(&cfg.common.otel_otlp_url)
-                .with_headers(headers),
-        )
+                .with_headers(headers)
+        })
+    } else {
+        tracer.with_exporter({
+            let mut metadata = MetadataMap::new();
+            metadata.insert(
+                MetadataKey::from_str(&cfg.common.tracing_header_key).unwrap(),
+                MetadataValue::from_str(&cfg.common.tracing_header_value).unwrap(),
+            );
+            metadata.insert(
+                MetadataKey::from_str(&cfg.grpc.org_header_key).unwrap(),
+                MetadataValue::from_str(&cfg.common.tracing_grpc_header_org).unwrap(),
+            );
+            metadata.insert(
+                MetadataKey::from_str(&cfg.grpc.stream_header_key).unwrap(),
+                MetadataValue::from_str(&cfg.common.tracing_grpc_header_stream_name).unwrap(),
+            );
+            opentelemetry_otlp::new_exporter()
+                .tonic()
+                .with_endpoint(&cfg.common.otel_otlp_grpc_url)
+                .with_metadata(metadata)
+                .with_protocol(opentelemetry_otlp::Protocol::Grpc)
+        })
+    };
+    let tracer = tracer
         .with_trace_config(sdktrace::config().with_resource(Resource::new(vec![
             KeyValue::new("service.name", cfg.common.node_role.to_string()),
             KeyValue::new("service.instance", cfg.common.instance_name.to_string()),