Skip to content

Commit

Permalink
feat: add trace headers flag (#253)
Browse files Browse the repository at this point in the history
nrwiersma authored Apr 4, 2024
1 parent e155169 commit 7ef5ca0
Showing 4 changed files with 43 additions and 7 deletions.
12 changes: 12 additions & 0 deletions cmd.go
Original file line number Diff line number Diff line change
@@ -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
}
7 changes: 2 additions & 5 deletions profile.go
Original file line number Diff line number Diff line change
@@ -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
21 changes: 19 additions & 2 deletions trace.go
Original file line number Diff line number Diff line change
@@ -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())
}
10 changes: 10 additions & 0 deletions trace_test.go
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 7ef5ca0

Please sign in to comment.