Remove recursion in protobuf data

2 files changed, 82 insertions, 5 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/StackTraceNode.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/node/StackTraceNode.java
index ed938d5..c0dcc5b 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/StackTraceNode.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/node/StackTraceNode.java
@@ -65,7 +65,7 @@ public final class StackTraceNode extends AbstractNode {
         return this.description.parentLineNumber;
     }
 
-    public SparkSamplerProtos.StackTraceNode toProto(MergeMode mergeMode, ProtoTimeEncoder timeEncoder) {
+    public SparkSamplerProtos.StackTraceNode toProto(MergeMode mergeMode, ProtoTimeEncoder timeEncoder, Iterable<Integer> childrenRefs) {
         SparkSamplerProtos.StackTraceNode.Builder proto = SparkSamplerProtos.StackTraceNode.newBuilder()
                 .setClassName(this.description.className)
                 .setMethodName(this.description.methodName);
@@ -91,9 +91,7 @@ public final class StackTraceNode extends AbstractNode {
                     .ifPresent(proto::setMethodDesc);
         }
 
-        for (StackTraceNode child : exportChildren(mergeMode)) {
-            proto.addChildren(child.toProto(mergeMode, timeEncoder));
-        }
+        proto.addAllChildrenRefs(childrenRefs);
 
         return proto.build();
     }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/ThreadNode.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/node/ThreadNode.java
index 1dce523..9faece6 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/ThreadNode.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/node/ThreadNode.java
@@ -21,8 +21,14 @@
 package me.lucko.spark.common.sampler.node;
 
 import me.lucko.spark.common.sampler.window.ProtoTimeEncoder;
+import me.lucko.spark.common.util.IndexedListBuilder;
 import me.lucko.spark.proto.SparkSamplerProtos;
 
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.LinkedList;
+import java.util.List;
+
 /**
  * The root of a sampling stack for a given thread / thread group.
  */
@@ -92,10 +98,83 @@ public final class ThreadNode extends AbstractNode {
             proto.addTimes(time);
         }
 
+        // When converting to a proto, we change the data structure from a recursive tree to an array.
+        // Effectively, instead of:
+        //
+        //   {
+        //     data: 'one',
+        //     children: [
+        //       {
+        //         data: 'two',
+        //         children: [{ data: 'four' }]
+        //       },
+        //       { data: 'three' }
+        //     ]
+        //   }
+        //
+        // we transmit:
+        //
+        //   [
+        //     { data: 'one', children: [1, 2] },
+        //     { data: 'two', children: [3] }
+        //     { data: 'three', children: [] }
+        //     { data: 'four', children: [] }
+        //   ]
+        //
+
+        // the flattened array of nodes
+        IndexedListBuilder<SparkSamplerProtos.StackTraceNode> nodesArray = new IndexedListBuilder<>();
+
+        // Perform a depth-first post order traversal of the tree
+        Deque<Node> stack = new ArrayDeque<>();
+
+        // push the thread node's children to the stack
+        List<Integer> childrenRefs = new LinkedList<>();
         for (StackTraceNode child : exportChildren(mergeMode)) {
-            proto.addChildren(child.toProto(mergeMode, timeEncoder));
+            stack.push(new Node(child, childrenRefs));
+        }
+
+        Node node;
+        while (!stack.isEmpty()) {
+            node = stack.peek();
+
+            // on the first visit, just push this node's children and leave it on the stack
+            if (node.firstVisit) {
+                for (StackTraceNode child : node.stackTraceNode.exportChildren(mergeMode)) {
+                    stack.push(new Node(child, node.childrenRefs));
+                }
+                node.firstVisit = false;
+                continue;
+            }
+
+            // convert StackTraceNode to a proto
+            // - at this stage, we have already visited this node's children
+            // - the refs for each child are stored in node.childrenRefs
+            SparkSamplerProtos.StackTraceNode childProto = node.stackTraceNode.toProto(mergeMode, timeEncoder, node.childrenRefs);
+
+            // add the child proto to the nodes array, and record the ref in the parent
+            int childIndex = nodesArray.add(childProto);
+            node.parentChildrenRefs.add(childIndex);
+
+            // pop from the stack
+            stack.pop();
         }
 
+        proto.addAllChildrenRefs(childrenRefs);
+        proto.addAllChildren(nodesArray.build());
+
         return proto.build();
     }
+
+    private static final class Node {
+        private final StackTraceNode stackTraceNode;
+        private boolean firstVisit = true;
+        private final List<Integer> childrenRefs = new LinkedList<>();
+        private final List<Integer> parentChildrenRefs;
+
+        private Node(StackTraceNode node, List<Integer> parentChildrenRefs) {
+            this.stackTraceNode = node;
+            this.parentChildrenRefs = parentChildrenRefs;
+        }
+    }
 }