Merge pull request #73 from openzipkin/refactor-node

Re-organizes Node so that it can be used on objects besides Span

zipkin/src/main/java/zipkin/internal/CorrectForClockSkew.java

@@ -13,6 +13,8 @@
  */
 package zipkin.internal;
 
+import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -39,9 +41,13 @@ public ClockSkew(Endpoint endpoint, long skew) {
   public static List<Span> apply(List<Span> spans) {
     for (Span s : spans) {
       if (s.parentId == null) {
-        SpanNode tree = SpanNode.create(s, spans);
+        Node<Span> tree = Node.constructTree(spans);
         adjust(tree, null);
-        return tree.toSpans();
+        List<Span> result = new ArrayList<>(spans.size());
+        for (Iterator<Node<Span>> i = tree.traverse(); i.hasNext();) {
+          result.add(i.next().value());
+        }
+        return result;
       }
     }
     return spans;
@@ -51,20 +57,20 @@ public static List<Span> apply(List<Span> spans) {
    * Recursively adjust the timestamps on the span tree. Root span is the reference point, all
    * children's timestamps gets adjusted based on that span's timestamps.
    */
-  private static void adjust(SpanNode node, @Nullable ClockSkew skewFromParent) {
+  private static void adjust(Node<Span> node, @Nullable ClockSkew skewFromParent) {
     // adjust skew for the endpoint brought over from the parent span
     if (skewFromParent != null) {
-      node.span = adjustTimestamps(node.span, skewFromParent);
+      node.value(adjustTimestamps(node.value(), skewFromParent));
     }
 
     // Is there any skew in the current span?
-    ClockSkew skew = getClockSkew(node.span);
+    ClockSkew skew = getClockSkew(node.value());
     if (skew != null) {
       // the current span's skew may be a different endpoint than skewFromParent, adjust again.
-      node.span = adjustTimestamps(node.span, skew);
+      node.value(adjustTimestamps(node.value(), skew));
 
       // propagate skew to any children
-      for (SpanNode child : node.children) {
+      for (Node<Span> child : node.children()) {
         adjust(child, skew);
       }
     }

zipkin/src/main/java/zipkin/internal/Node.java

@@ -0,0 +1,152 @@
+/**
+ * Copyright 2015-2016 The OpenZipkin 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 zipkin.internal;
+
+import java.util.ArrayDeque;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+import zipkin.Span;
+
+/**
+ * Convenience type representing a tree. This is here because multiple facets in zipkin require
+ * traversing the trace tree. For example, looking at network boundaries to correct clock skew, or
+ * counting requests imply visiting the tree.
+ *
+ * @param <V> the node's value. Ex a full span or a tuple like {@code (serviceName, isLocal)}
+ */
+public final class Node<V> {
+
+  /** Set via {@link #addChild(Node)} */
+  private Node<V> parent;
+  /** mutable as some transformations, such as clock skew, adjust this. */
+  private V value;
+  /** mutable to avoid allocating lists for childless nodes */
+  private List<Node<V>> children = Collections.emptyList();
+
+  /** Returns the parent, or null if root */
+  @Nullable
+  public Node<V> parent() {
+    return parent;
+  }
+
+  public V value() {
+    return value;
+  }
+
+  public Node<V> value(V newValue) {
+    this.value = newValue;
+    return this;
+  }
+
+  public Node<V> addChild(Node<V> child) {
+    child.parent = this;
+    if (children.equals(Collections.emptyList())) children = new LinkedList<>();
+    children.add(child);
+    return this;
+  }
+
+  /** Returns the children of this node. */
+  public Collection<Node<V>> children() {
+    return children;
+  }
+
+  /** Traverses the tree, breadth-first. */
+  public Iterator<Node<V>> traverse() {
+    return new BreadthFirstIterator<>(this);
+  }
+
+  static final class BreadthFirstIterator<V> implements Iterator<Node<V>> {
+    private final Queue<Node<V>> queue = new ArrayDeque<>();
+
+    BreadthFirstIterator(Node<V> root) {
+      queue.add(root);
+    }
+
+    @Override
+    public boolean hasNext() {
+      return !queue.isEmpty();
+    }
+
+    @Override
+    public Node<V> next() {
+      Node<V> result = queue.remove();
+      queue.addAll(result.children);
+      return result;
+    }
+
+    @Override
+    public void remove() {
+      throw new UnsupportedOperationException("remove");
+    }
+  }
+
+  /**
+   * @param trace spans that belong to the same {@link Span#traceId trace}, in any order.
+   */
+  static Node<Span> constructTree(List<Span> trace) {
+    TreeBuilder<Span> treeBuilder = new TreeBuilder<>();
+    for (Span s : trace) {
+      treeBuilder.addNode(s.parentId, s.id, s);
+    }
+    return treeBuilder.build();
+  }
+
+  /**
+   * Some operations do not require the entire span object. This creates a tree given (parent id,
+   * id) pairs.
+   *
+   * @param <V> same type as {@link Node#value}
+   */
+  public static final class TreeBuilder<V> {
+    Node<V> rootNode = null;
+
+    // Nodes representing the trace tree
+    Map<Long, Node<V>> idToNode = new LinkedHashMap<>();
+    // Collect the parent-child relationships between all spans.
+    Map<Long, Long> idToParent = new LinkedHashMap<>(idToNode.size());
+
+    public void addNode(Long parentId, long id, @Nullable V value) {
+      Node<V> node = new Node<V>().value(value);
+      if (parentId == null) { // special-case root
+        rootNode = node;
+      } else {
+        idToNode.put(id, node);
+        idToParent.put(id, parentId);
+      }
+    }
+
+    /** Builds a tree from calls to {@link #addNode}, or returns an empty tree. */
+    public Node<V> build() {
+      // Materialize the tree using parent - child relationships
+      for (Map.Entry<Long, Long> entry : idToParent.entrySet()) {
+        Node<V> node = idToNode.get(entry.getKey());
+        Node<V> parent = idToNode.get(entry.getValue());
+        if (parent == null && rootNode == null) { // handle headless trace
+          rootNode = node;
+        } else if (parent == null) { // attribute missing parents to root
+          rootNode.addChild(node);
+        } else {
+          parent.addChild(node);
+        }
+      }
+      return rootNode != null ? rootNode : new Node<V>();
+    }
+  }
+}

zipkin/src/test/java/zipkin/internal/NodeTest.java

@@ -0,0 +1,81 @@
+/**
+ * Copyright 2015-2016 The OpenZipkin 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 zipkin.internal;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.junit.Test;
+import zipkin.Span;
+import zipkin.TestObjects;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class NodeTest {
+
+  /**
+   * <p>The following tree should traverse in alphabetical order <pre>{@code
+   *
+   *          a
+   *        / | \
+   *       b  c  d
+   *      /|\     \
+   *     e f g     h
+   * }</pre>
+   */
+  @Test
+  public void traversesBreadthFirst() {
+    Node<Character> a = new Node<Character>().value('a');
+    Node<Character> b = new Node<Character>().value('b');
+    Node<Character> c = new Node<Character>().value('c');
+    Node<Character> d = new Node<Character>().value('d');
+    // root(a) has children b, c, d
+    a.addChild(b).addChild(c).addChild(d);
+    Node<Character> e = new Node<Character>().value('e');
+    Node<Character> f = new Node<Character>().value('f');
+    Node<Character> g = new Node<Character>().value('g');
+    // child(b) has children e, f, g
+    b.addChild(e).addChild(f).addChild(g);
+    Node<Character> h = new Node<Character>().value('h');
+    // f has no children
+    // child(g) has child h
+    g.addChild(h);
+
+    assertThat(a.traverse()).extracting(Node::value)
+        .containsExactly('a', 'b', 'c', 'd', 'e', 'f', 'g', 'h');
+  }
+
+  /**
+   * Makes sure that the trace tree is constructed based on parent-child, not by parameter order.
+   */
+  @Test
+  public void constructsTraceTree() {
+    // TRACE is sorted with root span first, lets shuffle them to make
+    // sure the trace is stitched together by id.
+    List<Span> copy = new ArrayList<>(TestObjects.TRACE);
+
+    Collections.shuffle(copy);
+
+    Node<Span> root = Node.constructTree(copy);
+    assertThat(root.value())
+        .isEqualTo(TestObjects.TRACE.get(0));
+
+    assertThat(root.children()).extracting(Node::value)
+        .containsExactly(TestObjects.TRACE.get(1));
+
+    Node<Span> child = root.children().iterator().next();
+    assertThat(child.children()).extracting(Node::value)
+        .containsExactly(TestObjects.TRACE.get(2));
+  }
+}