diff --git a/server/model/traces.go b/server/model/traces.go index e95d824579..f27d34a736 100644 --- a/server/model/traces.go +++ b/server/model/traces.go @@ -23,12 +23,12 @@ func NewTrace(traceID string, spans []Span) Trace { spanMap[spanID] = &spanCopy } - var rootSpan *Span = nil + rootSpans := make([]*Span, 0) for _, span := range spanMap { parentID := span.Attributes["parent_id"] parentSpan, found := spanMap[parentID] if !found { - rootSpan = span + rootSpans = append(rootSpans, span) continue } @@ -36,10 +36,17 @@ func NewTrace(traceID string, spans []Span) Trace { span.Parent = parentSpan } + var rootSpan Span + if len(rootSpans) == 1 { + rootSpan = *rootSpans[0] + } else { + rootSpan = Span{ID: IDGen.SpanID(), Name: TemporaryRootSpanName, Attributes: make(Attributes), Children: rootSpans} + } + id, _ := trace.TraceIDFromHex(traceID) trace := Trace{ ID: id, - RootSpan: *rootSpan, + RootSpan: rootSpan, } trace = trace.Sort() @@ -89,6 +96,7 @@ func (t *Trace) Sort() Trace { } const TriggerSpanName = "Tracetest trigger" +const TemporaryRootSpanName = "Temporary Tracetest root span" func (t *Trace) HasRootSpan() bool { return t.RootSpan.Name == TriggerSpanName @@ -110,6 +118,12 @@ func (t *Trace) InsertRootSpan(newRoot Span) *Trace { } func replaceRoot(oldRoot, newRoot Span) Span { + if oldRoot.Name == TemporaryRootSpanName { + // Replace the temporary root with the actual root + newRoot.Children = oldRoot.Children + return newRoot + } + if oldRoot.Attributes == nil { oldRoot.Attributes = make(Attributes) } diff --git a/server/model/traces_test.go b/server/model/traces_test.go new file mode 100644 index 0000000000..cf29ba8ba7 --- /dev/null +++ b/server/model/traces_test.go @@ -0,0 +1,155 @@ +package model_test + +import ( + "testing" + "time" + + "github.com/kubeshop/tracetest/server/model" + "github.com/kubeshop/tracetest/server/traces" + "github.com/stretchr/testify/assert" + v1 "go.opentelemetry.io/proto/otlp/trace/v1" +) + +func TestTraces(t *testing.T) { + rootSpan := newSpan("Root", nil) + childSpan1 := newSpan("child 1", &rootSpan) + childSpan2 := newSpan("child 2", &rootSpan) + grandchildSpan := newSpan("grandchild", &childSpan2) + + spans := []model.Span{rootSpan, childSpan1, childSpan2, grandchildSpan} + trace := model.NewTrace("trace", spans) + + assert.Len(t, trace.Flat, 4) + assert.Equal(t, "Root", trace.RootSpan.Name) + assert.Equal(t, "child 1", trace.RootSpan.Children[0].Name) + assert.Equal(t, "child 2", trace.RootSpan.Children[1].Name) + assert.Equal(t, "grandchild", trace.RootSpan.Children[1].Children[0].Name) +} + +func TestTraceWithMultipleRoots(t *testing.T) { + root1 := newSpan("Root 1", nil) + root1Child := newSpan("Child from root 1", &root1) + root2 := newSpan("Root 2", nil) + root2Child := newSpan("Child from root 2", &root2) + root3 := newSpan("Root 3", nil) + root3Child := newSpan("Child from root 3", &root3) + + spans := []model.Span{root1, root1Child, root2, root2Child, root3, root3Child} + trace := model.NewTrace("trace", spans) + + // agreggate root + 3 roots + 3 child + assert.Len(t, trace.Flat, 7) + assert.Equal(t, model.TemporaryRootSpanName, trace.RootSpan.Name) + assert.Equal(t, "Root 1", trace.RootSpan.Children[0].Name) + assert.Equal(t, "Root 2", trace.RootSpan.Children[1].Name) + assert.Equal(t, "Root 3", trace.RootSpan.Children[2].Name) + assert.Equal(t, "Child from root 1", trace.RootSpan.Children[0].Children[0].Name) + assert.Equal(t, "Child from root 2", trace.RootSpan.Children[1].Children[0].Name) + assert.Equal(t, "Child from root 3", trace.RootSpan.Children[2].Children[0].Name) +} + +func TestTraceWithMultipleRootsFromOtel(t *testing.T) { + root1 := newOtelSpan("Root 1", nil) + root1Child := newOtelSpan("Child from root 1", root1) + root2 := newOtelSpan("Root 2", nil) + root2Child := newOtelSpan("Child from root 2", root2) + root3 := newOtelSpan("Root 3", nil) + root3Child := newOtelSpan("Child from root 3", root3) + + tracesData := &v1.TracesData{ + ResourceSpans: []*v1.ResourceSpans{ + { + ScopeSpans: []*v1.ScopeSpans{ + { + Spans: []*v1.Span{root1, root1Child, root2, root2Child, root3, root3Child}, + }, + }, + }, + }, + } + + trace := traces.FromOtel(tracesData) + + // agreggate root + 3 roots + 3 child + assert.Len(t, trace.Flat, 7) + assert.Equal(t, model.TemporaryRootSpanName, trace.RootSpan.Name) + assert.Equal(t, "Root 1", trace.RootSpan.Children[0].Name) + assert.Equal(t, "Root 2", trace.RootSpan.Children[1].Name) + assert.Equal(t, "Root 3", trace.RootSpan.Children[2].Name) + assert.Equal(t, "Child from root 1", trace.RootSpan.Children[0].Children[0].Name) + assert.Equal(t, "Child from root 2", trace.RootSpan.Children[1].Children[0].Name) + assert.Equal(t, "Child from root 3", trace.RootSpan.Children[2].Children[0].Name) +} + +func TestInjectingNewRootWhenSingleRoot(t *testing.T) { + rootSpan := newSpan("Root", nil) + childSpan1 := newSpan("child 1", &rootSpan) + childSpan2 := newSpan("child 2", &rootSpan) + grandchildSpan := newSpan("grandchild", &childSpan2) + + spans := []model.Span{rootSpan, childSpan1, childSpan2, grandchildSpan} + trace := model.NewTrace("trace", spans) + + newRoot := newSpan("new Root", nil) + newTrace := trace.InsertRootSpan(newRoot) + + assert.Len(t, newTrace.Flat, 5) + assert.Equal(t, "new Root", newTrace.RootSpan.Name) + assert.Len(t, newTrace.RootSpan.Children, 1) + assert.Equal(t, "Root", newTrace.RootSpan.Children[0].Name) +} + +func TestInjectingNewRootWhenMultipleRoots(t *testing.T) { + root1 := newSpan("Root 1", nil) + root1Child := newSpan("Child from root 1", &root1) + root2 := newSpan("Root 2", nil) + root2Child := newSpan("Child from root 2", &root2) + root3 := newSpan("Root 3", nil) + root3Child := newSpan("Child from root 3", &root3) + + spans := []model.Span{root1, root1Child, root2, root2Child, root3, root3Child} + trace := model.NewTrace("trace", spans) + + newRoot := newSpan("new Root", nil) + newTrace := trace.InsertRootSpan(newRoot) + + assert.Len(t, newTrace.Flat, 7) + assert.Equal(t, "new Root", newTrace.RootSpan.Name) + assert.Len(t, newTrace.RootSpan.Children, 3) + assert.Equal(t, "Root 1", newTrace.RootSpan.Children[0].Name) + assert.Equal(t, "Root 2", newTrace.RootSpan.Children[1].Name) + assert.Equal(t, "Root 3", newTrace.RootSpan.Children[2].Name) +} + +func newSpan(name string, parent *model.Span) model.Span { + span := model.Span{ + ID: model.IDGen.SpanID(), + Name: name, + Parent: parent, + Attributes: make(model.Attributes), + StartTime: time.Now(), + EndTime: time.Now().Add(1 * time.Second), + } + + if parent != nil { + span.Attributes["parent_id"] = parent.ID.String() + } + + return span +} + +func newOtelSpan(name string, parent *v1.Span) *v1.Span { + id := model.IDGen.SpanID() + var parentId []byte = nil + if parent != nil { + parentId = parent.SpanId + } + + return &v1.Span{ + SpanId: id[:], + Name: name, + ParentSpanId: parentId, + StartTimeUnixNano: uint64(time.Now().UnixNano()), + EndTimeUnixNano: uint64(time.Now().Add(1 * time.Second).UnixNano()), + } +}