From caf91a7aabca44de924e306ab458928bc7c60f04 Mon Sep 17 00:00:00 2001
From: Ruben Vargas Palma <ruben.vp8510@gmail.com>
Date: Wed, 6 Mar 2019 23:00:18 -0600
Subject: [PATCH] Process FOLLOWS_FROM spans in TraceView (#335)

* Process FOLLOWS_FROM spans in TraceView

Signed-off-by: Ruben Vargas <ruben.vp8510@gmail.com>

* Add test for FOLLOWS_FROM span relation for SpanTreeOffset

Signed-off-by: Ruben Vargas <ruben.vp8510@gmail.com>
---
 .../TraceTimelineViewer/SpanTreeOffset.js     | 23 ++++++++------
 .../SpanTreeOffset.test.js                    | 30 +++++++++++++++++++
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.js b/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.js
index 9bd0971886..fbe762cbaf 100644
--- a/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.js
+++ b/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.js
@@ -16,7 +16,6 @@
 
 import cx from 'classnames';
 import _get from 'lodash/get';
-import _find from 'lodash/find';
 import React from 'react';
 import IoChevronRight from 'react-icons/lib/io/chevron-right';
 import IoIosArrowDown from 'react-icons/lib/io/ios-arrow-down';
@@ -54,14 +53,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');
@@ -106,6 +98,19 @@ export class UnconnectedSpanTreeOffset extends React.PureComponent<SpanTreeOffse
     }
   };
 
+  computeAncestors(rootSpan: Span) {
+    let currentSpan: Span = rootSpan;
+    while (currentSpan && Array.isArray(currentSpan.references) && currentSpan.references.length) {
+      const { refType, span } = currentSpan.references[0] || {};
+      if (span && (refType === 'FOLLOWS_FROM' || refType === 'CHILD_OF')) {
+        this.ancestorIds.push(span.spanID);
+        currentSpan = span;
+      } else {
+        break;
+      }
+    }
+  }
+
   render() {
     const { showChildrenIcon, childrenVisible, onClick, span } = this.props;
     const { hasChildren, spanID } = span;
diff --git a/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.test.js b/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.test.js
index 5cefbdc4fd..0cc16f1b3c 100644
--- a/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.test.js
+++ b/packages/jaeger-ui/src/components/TracePage/TraceTimelineViewer/SpanTreeOffset.test.js
@@ -43,6 +43,7 @@ describe('SpanTreeOffset', () => {
     ],
     spanID: ownSpanID,
   };
+
   const specialRootID = 'root';
   let props;
   let wrapper;
@@ -77,6 +78,35 @@ describe('SpanTreeOffset', () => {
       expect(indentGuides.prop('data-ancestor-id')).toBe(specialRootID);
     });
 
+    it('renders one .SpanTreeOffset--indentGuide per ancestor span, plus one for entire trace, including FOLLOWS_FROM', () => {
+      props.span = {
+        hasChildren: false,
+        references: [
+          {
+            refType: 'FOLLOWS_FROM',
+            span: {
+              spanID: parentSpanID,
+              references: [
+                {
+                  refType: 'CHILD_OF',
+                  span: {
+                    spanID: rootSpanID,
+                  },
+                },
+              ],
+            },
+          },
+        ],
+        spanID: ownSpanID,
+      };
+      wrapper = shallow(<UnconnectedSpanTreeOffset {...props} />);
+      const indentGuides = wrapper.find('.SpanTreeOffset--indentGuide');
+      expect(indentGuides.length).toBe(3);
+      expect(indentGuides.at(0).prop('data-ancestor-id')).toBe(specialRootID);
+      expect(indentGuides.at(1).prop('data-ancestor-id')).toBe(rootSpanID);
+      expect(indentGuides.at(2).prop('data-ancestor-id')).toBe(parentSpanID);
+    });
+
     it('renders one .SpanTreeOffset--indentGuide per ancestor span, plus one for entire trace', () => {
       const indentGuides = wrapper.find('.SpanTreeOffset--indentGuide');
       expect(indentGuides.length).toBe(3);