feat: add trace headers flag (#253)

Author: nrwiersma

cmd.go

@@ -20,3 +20,15 @@ func Split(slice []string, sep string) ([][2]string, error) {
 
 	return res, nil
 }
+
+func sliceToMap(s []string) (map[string]string, error) {
+	m := make(map[string]string, len(s))
+	kvs, err := Split(s, "=")
+	if err != nil {
+		return nil, err
+	}
+	for _, kv := range kvs {
+		m[kv[0]] = kv[1]
+	}
+	return m, nil
+}

profile.go

@@ -90,15 +90,12 @@ func NewProfiler(c *cli.Context, svc string, log *logger.Logger) (*pyroscope.Pro
 		Path:   u.Path,
 	}
 
-	tags := map[string]string{}
+	var tags map[string]string
 	if pairs := c.StringSlice(FlagProfilingTags); len(pairs) > 0 {
-		strTags, err := Split(pairs, "=")
+		tags, err = sliceToMap(pairs)
 		if err != nil {
 			return nil, err
 		}
-		for _, kv := range strTags {
-			tags[kv[0]] = kv[1]
-		}
 	}
 
 	types := allProfilingTypes

trace.go

@@ -22,6 +22,7 @@ const (
 	FlagTracingEndpoint         = "tracing.endpoint"
 	FlagTracingEndpointInsecure = "tracing.endpoint-insecure"
 	FlagTracingTags             = "tracing.tags"
+	FlagTracingHeaders          = "tracing.headers"
 	FlagTracingRatio            = "tracing.ratio"
 )
 
@@ -47,6 +48,11 @@ var TracingFlags = Flags{
 		Usage:   "A list of tags appended to every trace. Format: key=value.",
 		EnvVars: []string{"TRACING_TAGS"},
 	},
+	&cli.StringSliceFlag{
+		Name:    FlagTracingHeaders,
+		Usage:   "A list of headers appended to every trace when supported by the exporter. Format: key=value.",
+		EnvVars: []string{"TRACING_HEADERS"},
+	},
 	&cli.Float64Flag{
 		Name:    FlagTracingRatio,
 		Usage:   "The ratio between 0 and 1 of sample traces to take.",
@@ -94,19 +100,30 @@ func createExporter(c *cli.Context) (trace.SpanExporter, error) {
 	backend := c.String(FlagTracingExporter)
 	endpoint := c.String(FlagTracingEndpoint)
 
+	var (
+		headers map[string]string
+		err     error
+	)
+	if h := c.StringSlice(FlagTracingHeaders); len(h) > 0 {
+		headers, err = sliceToMap(h)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	switch backend {
 	case "":
 		return nil, nil
 	case "zipkin":
 		return zipkin.New(endpoint)
 	case "otlphttp":
-		opts := []otlptracehttp.Option{otlptracehttp.WithEndpoint(endpoint)}
+		opts := []otlptracehttp.Option{otlptracehttp.WithEndpoint(endpoint), otlptracehttp.WithHeaders(headers)}
 		if c.Bool(FlagTracingEndpointInsecure) {
 			opts = append(opts, otlptracehttp.WithInsecure())
 		}
 		return otlptracehttp.New(c.Context, opts...)
 	case "otlpgrpc":
-		opts := []otlptracegrpc.Option{otlptracegrpc.WithEndpoint(endpoint)}
+		opts := []otlptracegrpc.Option{otlptracegrpc.WithEndpoint(endpoint), otlptracegrpc.WithHeaders(headers)}
 		if c.Bool(FlagTracingEndpointInsecure) {
 			opts = append(opts, otlptracegrpc.WithInsecure())
 		}

trace_test.go

@@ -17,6 +17,7 @@ func TestNewTracer(t *testing.T) {
 		exporter string
 		endpoint string
 		tags     []string
+		headers  []string
 		ratio    float64
 		wantErr  require.ErrorAssertionFunc
 	}{
@@ -42,6 +43,14 @@ func TestNewTracer(t *testing.T) {
 			ratio:    1.0,
 			wantErr:  require.NoError,
 		},
+		{
+			name:     "with headers",
+			exporter: "otlphttp",
+			endpoint: "localhost:1234",
+			headers:  []string{"cluster=test", "namespace=num"},
+			ratio:    1.0,
+			wantErr:  require.NoError,
+		},
 		{
 			name:     "unknown exporter",
 			exporter: "some-exporter",
@@ -74,6 +83,7 @@ func TestNewTracer(t *testing.T) {
 			fs.String(cmd.FlagTracingExporter, test.exporter, "doc")
 			fs.String(cmd.FlagTracingEndpoint, test.endpoint, "doc")
 			fs.Var(cli.NewStringSlice(test.tags...), cmd.FlagTracingTags, "doc")
+			fs.Var(cli.NewStringSlice(test.headers...), cmd.FlagTracingHeaders, "doc")
 			fs.Float64(cmd.FlagTracingRatio, test.ratio, "doc")
 
 			log := logger.New(io.Discard, logger.LogfmtFormat(), logger.Error)