From 0ee275a53eca8774268e432d2c912d6f4ef563e8 Mon Sep 17 00:00:00 2001
From: kubarydz <niechzyjetequila@gmail.com>
Date: Sun, 16 Apr 2023 20:07:04 +0200
Subject: [PATCH] Allow follows-from reference as a parent span id (#4382)

## Which problem is this PR solving?
Resolves #4380

## Short description of the changes
handle FOLLOWS_FROM as a parent span ID

---------

Signed-off-by: Jakub Rydz <niechzyjetequila@gmail.com>
Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
---
 model/adjuster/clockskew.go |  1 -
 model/span.go               | 14 ++++++++++++--
 model/span_test.go          |  5 +++++
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/model/adjuster/clockskew.go b/model/adjuster/clockskew.go
index 167590d862d..dccc5e8d580 100644
--- a/model/adjuster/clockskew.go
+++ b/model/adjuster/clockskew.go
@@ -119,7 +119,6 @@ func (a *clockSkewAdjuster) buildNodesMap() {
 func (a *clockSkewAdjuster) buildSubGraphs() {
 	a.roots = make(map[model.SpanID]*node)
 	for _, n := range a.spans {
-		// TODO handle FOLLOWS_FROM references
 		if n.span.ParentSpanID() == 0 {
 			a.roots[n.span.SpanID] = n
 			continue
diff --git a/model/span.go b/model/span.go
index 5136e06b3ea..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 {
+		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 9aa50389365..efd9b5b9c08 100644
--- a/model/span_test.go
+++ b/model/span_test.go
@@ -277,6 +277,11 @@ func TestParentSpanID(t *testing.T) {
 	span := makeSpan(model.String("k", "v"))
 	assert.Equal(t, model.NewSpanID(123), span.ParentSpanID())
 
+	span.References = []model.SpanRef{
+		model.NewFollowsFromRef(span.TraceID, model.NewSpanID(777)),
+	}
+	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)),