From 577e1a0e504c057b8dc572ee8182261e4b941d70 Mon Sep 17 00:00:00 2001
From: Ruben Vargas <ruben.vp8510@gmail.com>
Date: Fri, 22 Feb 2019 00:36:44 -0500
Subject: [PATCH] Process FOLLOW_UP spans in TraceView

Signed-off-by: Ruben Vargas <ruben.vp8510@gmail.com>
---
 .../TraceTimelineViewer/SpanTreeOffset.js     | 35 ++++++++++++++-----
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.js b/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.js
index 133c08b4a7..4688c2f4eb 100644
--- a/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.js
+++ b/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.js
@@ -52,14 +52,7 @@ export class UnconnectedSpanTreeOffset extends React.PureComponent<SpanTreeOffse
     super(props);
 
     this.ancestorIds = [];
-    let currentSpan: Span = props.span;
-    while (currentSpan) {
-      currentSpan = _get(_find(currentSpan.references, { refType: 'CHILD_OF' }), 'span');
-      if (currentSpan) {
-        this.ancestorIds.push(currentSpan.spanID);
-      }
-    }
-
+    this.computeAncestors(props.span);
     // Some traces have multiple root-level spans, this connects them all under one guideline and adds the
     // necessary padding for the collapse icon on root-level spans.
     this.ancestorIds.push('root');
@@ -104,6 +97,32 @@ export class UnconnectedSpanTreeOffset extends React.PureComponent<SpanTreeOffse
     }
   };
 
+  checkForFollows(span: Span) {
+    // Could we have more than one follows_from ?
+    const follows = _get(_find(span.references, { refType: 'FOLLOWS_FROM' }), 'span');
+    if (follows) {
+      this.ancestorIds.push(follows.spanID);
+      this.computeAncestors(follows);
+      return true;
+    }
+    return false;
+  }
+
+  computeAncestors(span: Span) {
+    let currentSpan: Span = span;
+    const isFollow = this.checkForFollows(currentSpan);
+    // Assuming that follows and child_of are excludes.
+    if (!isFollow) {
+      while (currentSpan) {
+        currentSpan = _get(_find(currentSpan.references, { refType: 'CHILD_OF' }), 'span');
+        if (currentSpan) {
+          this.ancestorIds.push(currentSpan.spanID);
+          this.checkForFollows(currentSpan);
+        }
+      }
+    }
+  }
+
   render() {
     const { childrenVisible, onClick, span } = this.props;
     const { hasChildren, spanID } = span;