From 18534280113a79172a8ed314f3169f4f2881ae2a Mon Sep 17 00:00:00 2001
From: Toby <me@tobyan.com>
Date: Tue, 9 May 2023 04:58:29 +0800
Subject: [PATCH] feat: add otlptracegrpc otlptracehttp headers support for
 Uptrace (#3219)

Signed-off-by: Toby Yan <me@tobyan.com>
Co-authored-by: cong <zhangcong1992@gmail.com>
---
 core/trace/agent.go      | 20 ++++++++++++++++----
 core/trace/agent_test.go |  8 +++++++-
 core/trace/config.go     |  9 +++++----
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/core/trace/agent.go b/core/trace/agent.go
index 939c1d9e..9b293e45 100644
--- a/core/trace/agent.go
+++ b/core/trace/agent.go
@@ -71,17 +71,29 @@ func createExporter(c Config) (sdktrace.SpanExporter, error) {
 		// endpoint can not reach.
 		// If the connection not dial success, the global otel ErrorHandler will catch error
 		// when reporting data like other exporters.
-		return otlptracegrpc.New(
-			context.Background(),
+		opts := []otlptracegrpc.Option{
 			otlptracegrpc.WithInsecure(),
 			otlptracegrpc.WithEndpoint(c.Endpoint),
+		}
+		if len(c.OtlpHeaders) > 0 {
+			opts = append(opts, otlptracegrpc.WithHeaders(c.OtlpHeaders))
+		}
+		return otlptracegrpc.New(
+			context.Background(),
+			opts...,
 		)
 	case kindOtlpHttp:
 		// Not support flexible configuration now.
-		return otlptracehttp.New(
-			context.Background(),
+		opts := []otlptracehttp.Option{
 			otlptracehttp.WithInsecure(),
 			otlptracehttp.WithEndpoint(c.Endpoint),
+		}
+		if len(c.OtlpHeaders) > 0 {
+			opts = append(opts, otlptracehttp.WithHeaders(c.OtlpHeaders))
+		}
+		return otlptracehttp.New(
+			context.Background(),
+			opts...,
 		)
 	default:
 		return nil, fmt.Errorf("unknown exporter: %s", c.Batcher)
diff --git a/core/trace/agent_test.go b/core/trace/agent_test.go
index 853fa6ba..a3871a42 100644
--- a/core/trace/agent_test.go
+++ b/core/trace/agent_test.go
@@ -36,14 +36,20 @@ func TestStartAgent(t *testing.T) {
 		Batcher:  "otlp",
 	}
 	c5 := Config{
-		Name:     "grpc",
+		Name:     "otlpgrpc",
 		Endpoint: endpoint3,
 		Batcher:  kindOtlpGrpc,
+		OtlpHeaders: map[string]string{
+			"uptrace-dsn": "http://project2_secret_token@localhost:14317/2",
+		},
 	}
 	c6 := Config{
 		Name:     "otlphttp",
 		Endpoint: endpoint4,
 		Batcher:  kindOtlpHttp,
+		OtlpHeaders: map[string]string{
+			"uptrace-dsn": "http://project2_secret_token@localhost:14318/2",
+		},
 	}
 	c7 := Config{
 		Name:     "UDP",
diff --git a/core/trace/config.go b/core/trace/config.go
index 83c62c3c..9e6f7fb0 100644
--- a/core/trace/config.go
+++ b/core/trace/config.go
@@ -5,8 +5,9 @@ const TraceName = "go-zero"
 
 // A Config is an opentelemetry config.
 type Config struct {
-	Name     string  `json:",optional"`
-	Endpoint string  `json:",optional"`
-	Sampler  float64 `json:",default=1.0"`
-	Batcher  string  `json:",default=jaeger,options=jaeger|zipkin|otlpgrpc|otlphttp"`
+	Name        string            `json:",optional"`
+	Endpoint    string            `json:",optional"`
+	Sampler     float64           `json:",default=1.0"`
+	Batcher     string            `json:",default=jaeger,options=jaeger|zipkin|otlpgrpc|otlphttp"`
+	OtlpHeaders map[string]string `json:",optional"` // uptrace-dsn: 'http://project2_secret_token@localhost:14317/2'
 }