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

We need to create the trace tree in dependency store code. Particularly,
we don't want to retrieve the entire span just to count calls. This
refactors the Node object so that it can be used when all span details
aren't present.

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));
+  }
+}