From 3f09c0fcde59b2eaee692d349ae70bc1fe341c0b Mon Sep 17 00:00:00 2001
From: Jakub Rydz <niechzyjetequila@gmail.com>
Date: Fri, 14 Apr 2023 11:43:34 +0200
Subject: [PATCH] treat child-of ref as priority

Signed-off-by: Jakub Rydz <niechzyjetequila@gmail.com>
---
 model/span.go      | 14 ++++++++++++--
 model/span_test.go |  1 +
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/model/span.go b/model/span.go
index 658025563f6..749f21dbf2a 100644
--- a/model/span.go
+++ b/model/span.go
@@ -97,13 +97,23 @@ func (s *Span) NormalizeTimestamps() {
 }
 
 // ParentSpanID returns ID of a parent span if it exists.
-// It searches for the first child-of reference pointing to the same trace ID.
+// It searches for the first child-of or follows-from reference pointing to the same trace ID.
 func (s *Span) ParentSpanID() SpanID {
+	var followsFromRef *SpanRef
 	for i := range s.References {
 		ref := &s.References[i]
-		if ref.TraceID == s.TraceID && (ref.RefType == ChildOf || ref.RefType == FollowsFrom) {
+		if ref.TraceID != s.TraceID {
+			continue
+		}
+		if ref.RefType == ChildOf {
 			return ref.SpanID
 		}
+		if followsFromRef == nil && ref.RefType == FollowsFrom {
+			followsFromRef = ref
+		}
+	}
+	if followsFromRef != nil {
+		return followsFromRef.SpanID
 	}
 	return SpanID(0)
 }
diff --git a/model/span_test.go b/model/span_test.go
index d7681d990a4..efd9b5b9c08 100644
--- a/model/span_test.go
+++ b/model/span_test.go
@@ -283,6 +283,7 @@ func TestParentSpanID(t *testing.T) {
 	assert.Equal(t, model.NewSpanID(777), span.ParentSpanID())
 
 	span.References = []model.SpanRef{
+		model.NewFollowsFromRef(span.TraceID, model.NewSpanID(777)),
 		model.NewChildOfRef(span.TraceID, model.NewSpanID(888)),
 	}
 	assert.Equal(t, model.NewSpanID(888), span.ParentSpanID())