pomerium/internal/telemetry/trace/buffer_test.go
2025-01-13 17:34:35 +00:00

229 lines
6.7 KiB
Go

package trace_test
import (
"encoding/binary"
"fmt"
"testing"
"unique"
"github.com/pomerium/pomerium/internal/telemetry/trace"
"github.com/pomerium/pomerium/internal/testutil"
"github.com/stretchr/testify/assert"
oteltrace "go.opentelemetry.io/otel/trace"
commonv1 "go.opentelemetry.io/proto/otlp/common/v1"
resourcev1 "go.opentelemetry.io/proto/otlp/resource/v1"
tracev1 "go.opentelemetry.io/proto/otlp/trace/v1"
)
type (
Trace uint32
Span uint32
Scope uint32
Schema uint32
Resource uint32
)
func (n Trace) String() string { return fmt.Sprintf("Trace %d", n) }
func (n Span) String() string { return fmt.Sprintf("Span %d", n) }
func (n Scope) String() string { return fmt.Sprintf("Scope %d", n) }
func (n Schema) String() string { return fmt.Sprintf("Schema %d", n) }
func (n Resource) String() string { return fmt.Sprintf("Resource %d", n) }
func (n Trace) ID() unique.Handle[oteltrace.TraceID] {
id, _ := trace.ToTraceID(n.B())
return id
}
func (n Trace) B() []byte {
var id oteltrace.TraceID
binary.BigEndian.PutUint32(id[12:], uint32(n))
return id[:]
}
func (n Span) ID() oteltrace.SpanID {
id, _ := trace.ToSpanID(n.B())
return id
}
func (n Span) B() []byte {
var id oteltrace.SpanID
binary.BigEndian.PutUint32(id[4:], uint32(n))
return id[:]
}
func (n Scope) Make(s ...Schema) *trace.ScopeInfo {
if len(s) == 0 {
s = append(s, Schema(0))
}
return trace.NewScopeInfo(&commonv1.InstrumentationScope{
Name: n.String(),
Version: "v1",
Attributes: []*commonv1.KeyValue{
{
Key: "id",
Value: &commonv1.AnyValue{
Value: &commonv1.AnyValue_IntValue{
IntValue: int64(n),
},
},
},
},
}, s[0].String())
}
func (n Resource) Make(s ...Schema) *trace.ResourceInfo {
if len(s) == 0 {
s = append(s, Schema(0))
}
return trace.NewResourceInfo(&resourcev1.Resource{
Attributes: []*commonv1.KeyValue{
{
Key: "name",
Value: &commonv1.AnyValue{
Value: &commonv1.AnyValue_StringValue{
StringValue: n.String(),
},
},
},
{
Key: "id",
Value: &commonv1.AnyValue{
Value: &commonv1.AnyValue_IntValue{
IntValue: int64(n),
},
},
},
},
}, s[0].String())
}
func Traceparent(trace Trace, span Span, sampled bool) string {
sampledStr := "00"
if sampled {
sampledStr = "01"
}
return fmt.Sprintf("00-%s-%s-%s", trace.ID().Value(), span.ID(), sampledStr)
}
func TestBuffer(t *testing.T) {
t.Parallel()
// start time determines sort order of spans within a resource+scope group
s := []*tracev1.Span{
{TraceId: Trace(1).B(), SpanId: Span(1).B(), StartTimeUnixNano: 1},
{TraceId: Trace(1).B(), SpanId: Span(2).B(), StartTimeUnixNano: 2},
{TraceId: Trace(2).B(), SpanId: Span(3).B(), StartTimeUnixNano: 3},
{TraceId: Trace(2).B(), SpanId: Span(4).B(), StartTimeUnixNano: 4},
{TraceId: Trace(1).B(), SpanId: Span(5).B(), StartTimeUnixNano: 5},
{TraceId: Trace(1).B(), SpanId: Span(6).B(), StartTimeUnixNano: 6},
{TraceId: Trace(2).B(), SpanId: Span(7).B(), StartTimeUnixNano: 7},
{TraceId: Trace(2).B(), SpanId: Span(8).B(), StartTimeUnixNano: 8},
{TraceId: Trace(1).B(), SpanId: Span(9).B(), StartTimeUnixNano: 9},
{TraceId: Trace(1).B(), SpanId: Span(10).B(), StartTimeUnixNano: 10},
{TraceId: Trace(2).B(), SpanId: Span(11).B(), StartTimeUnixNano: 11},
{TraceId: Trace(2).B(), SpanId: Span(12).B(), StartTimeUnixNano: 12},
{TraceId: Trace(1).B(), SpanId: Span(13).B(), StartTimeUnixNano: 13},
{TraceId: Trace(1).B(), SpanId: Span(14).B(), StartTimeUnixNano: 14},
{TraceId: Trace(2).B(), SpanId: Span(15).B(), StartTimeUnixNano: 15},
{TraceId: Trace(2).B(), SpanId: Span(16).B(), StartTimeUnixNano: 16},
}
newTestBuffer := func() *trace.Buffer {
b := trace.NewBuffer()
b.Insert(Resource(1).Make(), Scope(1).Make(), s[0])
b.Insert(Resource(1).Make(), Scope(1).Make(), s[1])
b.Insert(Resource(1).Make(), Scope(1).Make(), s[2])
b.Insert(Resource(1).Make(), Scope(1).Make(), s[3])
b.Insert(Resource(1).Make(), Scope(2).Make(), s[4])
b.Insert(Resource(1).Make(), Scope(2).Make(), s[5])
b.Insert(Resource(1).Make(), Scope(2).Make(), s[6])
b.Insert(Resource(1).Make(), Scope(2).Make(), s[7])
b.Insert(Resource(2).Make(), Scope(1).Make(), s[8])
b.Insert(Resource(2).Make(), Scope(1).Make(), s[9])
b.Insert(Resource(2).Make(), Scope(1).Make(), s[10])
b.Insert(Resource(2).Make(), Scope(1).Make(), s[11])
b.Insert(Resource(2).Make(), Scope(2).Make(), s[12])
b.Insert(Resource(2).Make(), Scope(2).Make(), s[13])
b.Insert(Resource(2).Make(), Scope(2).Make(), s[14])
b.Insert(Resource(2).Make(), Scope(2).Make(), s[15])
return b
}
newExpectedSpans := func() []*tracev1.ResourceSpans {
return []*tracev1.ResourceSpans{
{
Resource: Resource(1).Make().Resource,
ScopeSpans: []*tracev1.ScopeSpans{
{
Scope: Scope(1).Make().Scope,
Spans: []*tracev1.Span{s[0], s[1], s[2], s[3]},
SchemaUrl: Schema(0).String(),
},
{
Scope: Scope(2).Make().Scope,
Spans: []*tracev1.Span{s[4], s[5], s[6], s[7]},
SchemaUrl: Schema(0).String(),
},
},
SchemaUrl: Schema(0).String(),
},
{
Resource: Resource(2).Make().Resource,
ScopeSpans: []*tracev1.ScopeSpans{
{
Scope: Scope(1).Make().Scope,
Spans: []*tracev1.Span{s[8], s[9], s[10], s[11]},
SchemaUrl: Schema(0).String(),
},
{
Scope: Scope(2).Make().Scope,
Spans: []*tracev1.Span{s[12], s[13], s[14], s[15]},
SchemaUrl: Schema(0).String(),
},
},
SchemaUrl: Schema(0).String(),
},
}
}
t.Run("Flush", func(t *testing.T) {
b := newTestBuffer()
actual := b.Flush()
assert.True(t, b.IsEmpty())
testutil.AssertProtoEqual(t, newExpectedSpans(), actual)
})
t.Run("FlushAs", func(t *testing.T) {
b := newTestBuffer()
actual := b.FlushAs(Trace(100).ID())
assert.True(t, b.IsEmpty())
expected := newExpectedSpans()
for _, resourceSpans := range expected {
for _, scopeSpans := range resourceSpans.ScopeSpans {
for _, span := range scopeSpans.Spans {
span.TraceId = Trace(100).B()
}
}
}
testutil.AssertProtoEqual(t, expected, actual)
})
t.Run("Default scope", func(t *testing.T) {
b := trace.NewBuffer()
b.Insert(Resource(1).Make(Schema(2)), trace.NewScopeInfo(nil, ""), s[0])
b.Insert(Resource(1).Make(Schema(2)), trace.NewScopeInfo(nil, ""), s[1])
b.Insert(Resource(1).Make(Schema(2)), trace.NewScopeInfo(nil, ""), s[2])
actual := b.Flush()
testutil.AssertProtoEqual(t, []*tracev1.ResourceSpans{
{
Resource: Resource(1).Make(Schema(2)).Resource,
ScopeSpans: []*tracev1.ScopeSpans{
{
Scope: nil,
Spans: []*tracev1.Span{s[0], s[1], s[2]},
SchemaUrl: "",
},
},
SchemaUrl: Schema(2).String(),
},
}, actual)
})
}