Introduce ParentReference adjuster (#3786)

* Introduce ParentReference adjuster

As demonstrated in https://github.com/jaegertracing/jaeger-ui/issues/966,
jaeger-ui expects the first reference to be of `CHILD_OF` type to properly
render the tree. This new adjuster achieves exactly that.

Signed-off-by: Ivan Babrou <[email protected]>

* Simplify algorithm, clean-up tests

Signed-off-by: Yuri Shkuro <[email protected]>

Co-authored-by: Yuri Shkuro <[email protected]>

cmd/query/app/querysvc/adjusters.go

@@ -30,5 +30,6 @@ func StandardAdjusters(maxClockSkewAdjust time.Duration) []adjuster.Adjuster {
 		adjuster.IPTagAdjuster(),
 		adjuster.SortLogFields(),
 		adjuster.SpanReferences(),
+		adjuster.ParentReference(),
 	}
 }

model/adjuster/parent_reference.go

@@ -0,0 +1,50 @@
+// Copyright (c) 2022 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package adjuster
+
+import (
+	"github.com/jaegertracing/jaeger/model"
+)
+
+// ParentReference returns an Adjuster that puts CHILD_OF references first.
+// This is necessary to match jaeger-ui expectations:
+// * https://github.com/jaegertracing/jaeger-ui/issues/966
+func ParentReference() Adjuster {
+	return Func(func(trace *model.Trace) (*model.Trace, error) {
+		for _, span := range trace.Spans {
+			firstChildOfRef := -1
+			firstOtherRef := -1
+			for i := 1; i < len(span.References); i++ {
+				if span.References[i].TraceID == span.TraceID {
+					if span.References[i].RefType == model.SpanRefType_CHILD_OF {
+						firstChildOfRef = i
+						break
+					} else if firstOtherRef == -1 {
+						firstOtherRef = i
+					}
+				}
+			}
+			swap := firstChildOfRef
+			if swap == -1 {
+				swap = firstOtherRef
+			}
+			if swap != -1 {
+				span.References[swap], span.References[0] = span.References[0], span.References[swap]
+			}
+		}
+
+		return trace, nil
+	})
+}

model/adjuster/parent_reference_test.go

@@ -0,0 +1,109 @@
+// Copyright (c) 2022 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package adjuster
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/jaegertracing/jaeger/model"
+)
+
+func TestParentReference(t *testing.T) {
+	a := model.NewTraceID(0, 1)
+	b := model.NewTraceID(0, 2)
+	childOf := func(id model.TraceID) model.SpanRef {
+		return model.NewChildOfRef(id, 1)
+	}
+	followsFrom := func(id model.TraceID) model.SpanRef {
+		return model.NewFollowsFromRef(id, 1)
+	}
+	followsFrom2 := func(id model.TraceID) model.SpanRef {
+		return model.NewFollowsFromRef(id, 2)
+	}
+
+	testCases := []struct {
+		name     string
+		incoming []model.SpanRef
+		expected []model.SpanRef
+	}{
+		{
+			name:     "empty",
+			incoming: []model.SpanRef{},
+			expected: []model.SpanRef{},
+		},
+		{
+			name:     "single child",
+			incoming: []model.SpanRef{childOf(a)},
+			expected: []model.SpanRef{childOf(a)},
+		},
+		{
+			name:     "single remote child",
+			incoming: []model.SpanRef{childOf(b)},
+			expected: []model.SpanRef{childOf(b)},
+		},
+		{
+			name:     "local and remote child in order",
+			incoming: []model.SpanRef{childOf(a), childOf(b)},
+			expected: []model.SpanRef{childOf(a), childOf(b)},
+		},
+		{
+			name:     "local and remote child out of order",
+			incoming: []model.SpanRef{childOf(b), childOf(a)},
+			expected: []model.SpanRef{childOf(a), childOf(b)},
+		},
+		{
+			name:     "local child, remote follows",
+			incoming: []model.SpanRef{followsFrom(b), childOf(a)},
+			expected: []model.SpanRef{childOf(a), followsFrom(b)},
+		},
+		{
+			name:     "remote, local, local follows - keep order",
+			incoming: []model.SpanRef{followsFrom(b), followsFrom2(a), followsFrom(a)},
+			expected: []model.SpanRef{followsFrom2(a), followsFrom(b), followsFrom(a)},
+		},
+		{
+			name:     "remote child, local follows",
+			incoming: []model.SpanRef{childOf(b), followsFrom(a)},
+			expected: []model.SpanRef{followsFrom(a), childOf(b)},
+		},
+		{
+			name:     "remote child, local follows, local child",
+			incoming: []model.SpanRef{childOf(b), followsFrom(a), childOf(a)},
+			expected: []model.SpanRef{childOf(a), followsFrom(a), childOf(b)},
+		},
+		{
+			name:     "remote follows, local follows, local child",
+			incoming: []model.SpanRef{followsFrom(b), followsFrom(a), childOf(a)},
+			expected: []model.SpanRef{childOf(a), followsFrom(a), followsFrom(b)},
+		},
+	}
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			trace := &model.Trace{
+				Spans: []*model.Span{
+					{
+						TraceID:    a,
+						References: testCase.incoming,
+					},
+				},
+			}
+			trace, err := ParentReference().Adjust(trace)
+			assert.NoError(t, err)
+			assert.Equal(t, testCase.expected, trace.Spans[0].References)
+		})
+	}
+}