From ad7ad996d37b032005f47f998ff99c382374cf81 Mon Sep 17 00:00:00 2001 From: Mik <29287456+Mikulasz12@users.noreply.github.com> Date: Sat, 26 Sep 2026 19:50:04 +0100 Subject: [PATCH] Fix Java sampler export mutating profiling data --- .../common/sampler/java/JavaNodeExporter.java | 10 +++- .../common/sampler/node/StackTraceNode.java | 6 +++ .../spark/common/sampler/node/NodeTest.java | 48 +++++++++++++++++++ 3 files changed, 63 insertions(+), 1 deletion(-) diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaNodeExporter.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaNodeExporter.java index c1100860b8d..04621e91593 100644 --- a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaNodeExporter.java +++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaNodeExporter.java @@ -27,6 +27,7 @@ import me.lucko.spark.proto.SparkSamplerProtos; import java.util.ArrayList; +import java.util.BitSet; import java.util.Collection; import java.util.List; @@ -81,11 +82,18 @@ protected Collection exportChildren(Collection c } List list = new ArrayList<>(children.size()); + BitSet detached = new BitSet(); outer: for (StackTraceNode child : children) { - for (StackTraceNode other : list) { + for (int i = 0; i < list.size(); i++) { + StackTraceNode other = list.get(i); if (this.mergeStrategy.shouldMerge(this.methodDisambiguator, other, child)) { + if (!detached.get(i)) { + other = other.copy(); + list.set(i, other); + detached.set(i); + } other.merge(child); continue outer; } 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 dd3539a6ed4..b10b0b38459 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 @@ -41,6 +41,12 @@ public StackTraceNode(Description description) { this.description = description; } + public StackTraceNode copy() { + StackTraceNode copy = new StackTraceNode(this.description); + copy.merge(this); + return copy; + } + public String getClassName() { return this.description.className(); } diff --git a/spark-common/src/test/java/me/lucko/spark/common/sampler/node/NodeTest.java b/spark-common/src/test/java/me/lucko/spark/common/sampler/node/NodeTest.java index 52477ce408e..2b30f6dcee9 100644 --- a/spark-common/src/test/java/me/lucko/spark/common/sampler/node/NodeTest.java +++ b/spark-common/src/test/java/me/lucko/spark/common/sampler/node/NodeTest.java @@ -25,7 +25,10 @@ import me.lucko.spark.common.sampler.SamplerMode; import me.lucko.spark.common.sampler.async.AsyncNodeExporter; import me.lucko.spark.common.sampler.async.AsyncStackTraceElement; +import me.lucko.spark.common.sampler.java.JavaNodeExporter; +import me.lucko.spark.common.sampler.java.MergeStrategy; import me.lucko.spark.common.sampler.window.ProtoTimeEncoder; +import me.lucko.spark.common.util.MethodDisambiguator; import me.lucko.spark.proto.SparkSamplerProtos; import org.junit.jupiter.api.Test; @@ -155,6 +158,51 @@ public void testExport() { assertEquals(expected, proto); } + @Test + public void testJavaExportDoesNotMutateSourceTree() { + StackTraceNode.Describer describer = (element, parent) -> new StackTraceNode.JavaDescription( + element.getClassName(), + element.getMethodName(), + element.getLineNumber(), + parent == null ? StackTraceNode.NULL_LINE_NUMBER : parent.getLineNumber() + ); + + StackTraceElement root = new StackTraceElement("java.lang.Thread", "run", "Thread.java", 1); + StackTraceElement caller = new StackTraceElement("test.Caller", "run", "Caller.java", 50); + StackTraceElement target1 = new StackTraceElement("test.Target", "work", "Target.java", 10); + StackTraceElement target2 = new StackTraceElement("test.Target", "work", "Target.java", 20); + + ThreadNode threadNode = new ThreadNode("Test Thread"); + threadNode.log(describer, new StackTraceElement[]{target1, caller, root}, TimeUnit.SECONDS.toMicros(1), WINDOW); + threadNode.log(describer, new StackTraceElement[]{target2, caller, root}, TimeUnit.SECONDS.toMicros(1), WINDOW); + + ProtoTimeEncoder timeEncoder = new ProtoTimeEncoder(SamplerMode.EXECUTION.valueTransformer(), ImmutableList.of(threadNode)); + JavaNodeExporter exporter = new JavaNodeExporter( + timeEncoder, + MergeStrategy.SAME_METHOD, + new MethodDisambiguator(className -> null) + ); + + SparkSamplerProtos.ThreadNode first = exporter.export(threadNode); + SparkSamplerProtos.ThreadNode second = exporter.export(threadNode); + + assertEquals(first, second); + assertEquals(3, first.getChildrenCount()); + + SparkSamplerProtos.StackTraceNode exportedTarget = first.getChildrenList().stream() + .filter(node -> node.getClassName().equals("test.Target")) + .findFirst() + .orElseThrow(AssertionError::new); + assertEquals(2000, exportedTarget.getTimes(0)); + + StackTraceNode threadRun = threadNode.getChildren().iterator().next(); + StackTraceNode callerNode = threadRun.getChildren().iterator().next(); + assertEquals(2, callerNode.getChildren().size()); + for (StackTraceNode target : callerNode.getChildren()) { + assertEquals(TimeUnit.SECONDS.toMicros(1), target.getTimeAccumulator(WINDOW).longValue()); + } + } + @Test public void testRemoveTimeWindows() { ThreadNode threadNode = new ThreadNode("Test Thread");