diff --git a/context.go b/context.go index 90045f4f..43553655 100644 --- a/context.go +++ b/context.go @@ -22,8 +22,9 @@ import ( ) const ( - flagSampled = byte(1) - flagDebug = byte(2) + flagSampled = byte(1) + flagDebug = byte(2) + flagFirehose = byte(8) ) var ( @@ -88,6 +89,11 @@ func (c SpanContext) IsDebug() bool { return (c.flags & flagDebug) == flagDebug } +// IsFirehose indicates whether the firehose flag was set +func (c SpanContext) IsFirehose() bool { + return (c.flags & flagFirehose) == flagFirehose +} + // IsValid indicates whether this context actually represents a valid trace. func (c SpanContext) IsValid() bool { return c.traceID.IsValid() && c.spanID != 0 diff --git a/context_test.go b/context_test.go index 34dfc745..97b7e697 100644 --- a/context_test.go +++ b/context_test.go @@ -78,21 +78,79 @@ func TestSpanContext_WithBaggageItem(t *testing.T) { assert.Equal(t, map[string]string{"some-KEY": "Some-Other-Value"}, ctx.baggage) } -func TestSpanContext_SampledDebug(t *testing.T) { - ctx, err := ContextFromString("1:1:1:1") - require.NoError(t, err) - assert.True(t, ctx.IsSampled()) - assert.False(t, ctx.IsDebug()) +func TestSpanContext_Flags(t *testing.T) { - ctx, err = ContextFromString("1:1:1:3") - require.NoError(t, err) - assert.True(t, ctx.IsSampled()) - assert.True(t, ctx.IsDebug()) + var tests = map[string]struct { + in string + sampledFlag bool + debugFlag bool + firehoseFlag bool + }{ + "None": { + in: "1:1:1:0", + sampledFlag: false, + debugFlag: false, + firehoseFlag: false, + }, + "Sampled Only": { + in: "1:1:1:1", + sampledFlag: true, + debugFlag: false, + firehoseFlag: false, + }, - ctx, err = ContextFromString("1:1:1:0") - require.NoError(t, err) - assert.False(t, ctx.IsSampled()) - assert.False(t, ctx.IsDebug()) + "Debug Only": { + in: "1:1:1:2", + sampledFlag: false, + debugFlag: true, + firehoseFlag: false, + }, + + "IsFirehose Only": { + in: "1:1:1:8", + sampledFlag: false, + debugFlag: false, + firehoseFlag: true, + }, + + "Sampled And Debug": { + in: "1:1:1:3", + sampledFlag: true, + debugFlag: true, + firehoseFlag: false, + }, + + "Sampled And Firehose": { + in: "1:1:1:9", + sampledFlag: true, + debugFlag: false, + firehoseFlag: true, + }, + + "Debug And Firehose": { + in: "1:1:1:10", + sampledFlag: false, + debugFlag: true, + firehoseFlag: true, + }, + + "Sampled And Debug And Firehose": { + in: "1:1:1:11", + sampledFlag: true, + debugFlag: true, + firehoseFlag: true, + }, + } + + for name, tc := range tests { + t.Run(name, func(t *testing.T) { + ctx, err := ContextFromString(tc.in) + require.NoError(t, err) + assert.Equal(t, tc.sampledFlag, ctx.IsSampled()) + assert.Equal(t, tc.debugFlag, ctx.IsDebug()) + assert.Equal(t, tc.firehoseFlag, ctx.IsFirehose()) + }) + } } func TestSpanContext_CopyFrom(t *testing.T) { diff --git a/span.go b/span.go index 2b5cd37e..d4669fc3 100644 --- a/span.go +++ b/span.go @@ -318,3 +318,8 @@ func setSamplingPriority(s *Span, value interface{}) bool { } return false } + +// EnableFirehose enables firehose flag on the span context +func EnableFirehose(s *Span) { + s.context.flags |= flagFirehose +} diff --git a/span_test.go b/span_test.go index 1dd781d8..d7bf9322 100644 --- a/span_test.go +++ b/span_test.go @@ -126,6 +126,18 @@ func TestSetTag_SamplingPriority(t *testing.T) { assert.False(t, sp1.context.IsDebug(), "debug should not be allowed by the throttler") } +func TestSetFirehoseMode(t *testing.T) { + tracer, closer := NewTracer("DOOP", NewConstSampler(true), NewNullReporter()) + defer closer.Close() + + sp1 := tracer.StartSpan("s1").(*Span) + assert.False(t, sp1.context.IsFirehose()) + + EnableFirehose(sp1) + + assert.True(t, sp1.context.IsFirehose()) +} + type testThrottler struct { allowAll bool }