From 1808ae39ef00094b8ee6b29f9e8f213001ceca7b Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Sat, 17 Oct 2026 14:26:08 +0100 Subject: [PATCH] Preserve trie mappings when long-jump capacity is exhausted An insertion could modify trie data before discovering that the long-jump table was full, leaving existing mappings corrupted after rejection. Check local allocations and ancestor jump conversions before modifying the trie. Use the actual insertion size so valid insertions still work when the table is full. Skip the ancestor scan when enough slots remain for every possible conversion. Add regression tests for rejected and valid insertions at capacity, plus builder benchmarks for exclusion rules, class names, and long jumps. --- .../utils/ClassNameTrieBenchmark.java | 2 +- .../utils/ClassNameTrieBuilderBenchmark.java | 88 +++++++++++ .../instrument/utils/ClassNameTrie.java | 85 ++++++++-- .../instrument/utils/ClassNameTrieTest.java | 145 +++++++++++++++++- 4 files changed, 301 insertions(+), 19 deletions(-) create mode 100644 utils/src/jmh/java/datadog/instrument/utils/ClassNameTrieBuilderBenchmark.java diff --git a/utils/src/jmh/java/datadog/instrument/utils/ClassNameTrieBenchmark.java b/utils/src/jmh/java/datadog/instrument/utils/ClassNameTrieBenchmark.java index 22cad91..b72e705 100644 --- a/utils/src/jmh/java/datadog/instrument/utils/ClassNameTrieBenchmark.java +++ b/utils/src/jmh/java/datadog/instrument/utils/ClassNameTrieBenchmark.java @@ -200,7 +200,7 @@ private static char normalize(char c) { return c == '/' ? '.' : c; } - private static void buildTrieEntries(BiConsumer builder) { + static void buildTrieEntries(BiConsumer builder) { builder.accept("cinnamon.*", 1); builder.accept("clojure.*", 1); builder.accept("co.elastic.apm.*", 1); diff --git a/utils/src/jmh/java/datadog/instrument/utils/ClassNameTrieBuilderBenchmark.java b/utils/src/jmh/java/datadog/instrument/utils/ClassNameTrieBuilderBenchmark.java new file mode 100644 index 0000000..53fff23 --- /dev/null +++ b/utils/src/jmh/java/datadog/instrument/utils/ClassNameTrieBuilderBenchmark.java @@ -0,0 +1,88 @@ +package datadog.instrument.utils; + +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.openjdk.jmh.annotations.Mode.AverageTime; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; + +@State(Scope.Benchmark) +@BenchmarkMode(AverageTime) +@OutputTimeUnit(MILLISECONDS) +@Fork( + value = 3, + jvmArgsAppend = {"-Xms512m", "-Xmx512m"}) +@Warmup(iterations = 3, time = 1) +@Measurement(iterations = 5, time = 1) +@Threads(1) +public class ClassNameTrieBuilderBenchmark { + + @Param({"exclusions", "classNames", "longJumps"}) + public String dataset; + + private String[] keys; + private int[] values; + + @Setup + public void setup() throws IOException { + List names = new ArrayList<>(); + List numbers = new ArrayList<>(); + switch (dataset) { + case "exclusions": + ClassNameTrieBenchmark.buildTrieEntries( + (name, value) -> { + names.add(name); + numbers.add(value); + }); + break; + case "classNames": + names.addAll(Files.readAllLines(Paths.get("src/test/resources/sample_class_names.txt"))); + break; + case "longJumps": + char[] chars = new char[7700]; + Arrays.fill(chars, 'p'); + String padding = new String(chars); + for (int group = 1; group <= 8; group++) { + for (int branch = 1; branch <= 8; branch++) { + names.add("" + (char) group + '\u0001' + (char) branch + padding); + } + for (int branch = 2; branch <= 65; branch++) { + char smallChar = (char) (branch <= 41 ? branch : branch + 1); // skip '*' + names.add("" + (char) group + smallChar); + } + } + break; + default: + throw new IllegalArgumentException("Unknown dataset: " + dataset); + } + keys = names.toArray(new String[0]); + values = new int[keys.length]; + for (int i = 0; i < values.length; i++) { + values[i] = numbers.isEmpty() ? 1 : numbers.get(i); + } + } + + @Benchmark + public ClassNameTrie.Builder build() { + ClassNameTrie.Builder builder = new ClassNameTrie.Builder(); + for (int i = 0; i < keys.length; i++) { + builder.put(keys[i], values[i]); + } + return builder; + } +} diff --git a/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java b/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java index ebf1c50..aeb7b06 100644 --- a/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java +++ b/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java @@ -110,6 +110,9 @@ public final class ClassNameTrie { /** Boundary marker where a jump is replaced by an index into the long jump table. */ private static final char LONG_JUMP_MARKER = 0xF000; // 4096 long-jumps, each >= 61440 + /** Maximum number of entries in the long jump table. */ + private static final int MAX_LONG_JUMPS = (char) ~LONG_JUMP_MARKER + 1; + /** A branch has at most 3 control characters: key, value, and optional jump offset/id. */ private static final int BRANCH_CONTROL_CHARS = 3; @@ -503,20 +506,20 @@ private void insertMapping(String key, char valueToInsert) { dataIndex += (branchCount * BRANCH_CONTROL_CHARS) - FIRST_JUMP_OMITTED; if ((value & LEAF_MARKER) != 0) { - // change leaf branch to a bud and append our new leaf node below it - trieData[valueIndex] = (char) ((value & ~LEAF_MARKER) | BUD_MARKER); + // append our new leaf node below the leaf branch and change it to a bud jumpOffset = appendLeaf(dataIndex, key, keyIndex, valueToInsert); + trieData[valueIndex] = (char) ((value & ~LEAF_MARKER) | BUD_MARKER); break; } else if ((value & BUD_MARKER) != 0) { // ignore bud as we still have more of the key to match continue; } else if (keyIndex == keyLength) { - // branch originally led to a segment - record it's now a bud - trieData[valueIndex] = (char) (valueToInsert | BUD_MARKER); if (value > 0) { // add the usual node preamble before the old segment jumpOffset = prependNode(dataIndex, value - 1); } + // branch originally led to a segment - record it's now a bud + trieData[valueIndex] = (char) (valueToInsert | BUD_MARKER); break; } @@ -535,11 +538,10 @@ private void insertMapping(String key, char valueToInsert) { if (dataIndex < segmentEnd) { if (keyIndex == keyLength) { // key is shorter than segment; add bud at the point the key ends - trieData[valueIndex] -= (char) (segmentEnd - (dataIndex - 1)); jumpOffset = insertBud(dataIndex - 1, valueToInsert, segmentEnd); + trieData[valueIndex] -= (char) (segmentEnd - (dataIndex - 1)); } else { // key diverges from segment; add leaf on left/right to capture that - trieData[valueIndex] -= (char) (segmentEnd - dataIndex); if (c < trieData[dataIndex]) { jumpOffset = insertLeafLeft(dataIndex, key, keyIndex, valueToInsert, segmentEnd); } else { @@ -547,6 +549,7 @@ private void insertMapping(String key, char valueToInsert) { insertLeafRight( dataIndex, key, keyIndex, valueToInsert, segmentEnd, subTrieEnd); } + trieData[valueIndex] -= (char) (segmentEnd - dataIndex); } break; // nothing more to add } @@ -556,8 +559,8 @@ private void insertMapping(String key, char valueToInsert) { if ((value & LEAF_MARKER) != 0) { if (keyIndex < keyLength) { // key goes past leaf segment, extend old leaf with a bud node and add our new leaf - trieData[valueIndex]--; jumpOffset = appendLeaf(dataIndex, key, keyIndex, valueToInsert); + trieData[valueIndex]--; break; } else { // duplicate key: overwrite existing value in the tree with the new value @@ -567,8 +570,8 @@ private void insertMapping(String key, char valueToInsert) { } else /* segment is followed by a node */ { if (keyIndex == keyLength) { // key stops at boundary between segment and node, insert bud node just before this - trieData[valueIndex]--; jumpOffset = prependNode(dataIndex - 1, valueToInsert | BUD_MARKER); + trieData[valueIndex]--; break; } } @@ -603,6 +606,27 @@ private int insertBranch( insertedCharacters = 3; /* branch-key, value, jump */ } + int subTrieStart = dataIndex + (branchCount * BRANCH_CONTROL_CHARS) - FIRST_JUMP_OMITTED; + + // count new long-jumps needed for this node's jumps after our new branch + int newLongJumps = 0; + if (newBranch < branchCount) { + int jumpIndex = dataIndex + (branchCount << 1) + newBranch - FIRST_JUMP_OMITTED; + int nextJump = (newBranch > 0 ? getJump(trieData[jumpIndex]) : 0) + remainingKeyLength; + if (nextJump >= LONG_JUMP_MARKER) { + newLongJumps++; + } + for (int b = newBranch + 1; b < branchCount; b++) { + if (becomesLongJump(trieData[++jumpIndex], remainingKeyLength)) { + newLongJumps++; + } + } + } else if (subTrieEnd - subTrieStart >= LONG_JUMP_MARKER) { + newLongJumps++; + } + + checkLongJumpCapacity(insertedCharacters, newLongJumps); + int i = dataIndex + newBranch, j = i + insertedCharacters; makeHole(i, insertedCharacters); @@ -619,8 +643,6 @@ private int insertBranch( // insert our new branch value trieData[i++] = (char) (collapseRight ? value | LEAF_MARKER : remainingKeyLength - 1); - int subTrieStart = dataIndex + (branchCount * BRANCH_CONTROL_CHARS) - FIRST_JUMP_OMITTED; - int precedingJump; if (newBranch < branchCount) { // adding branch on left/centre @@ -667,6 +689,8 @@ private int prependNode(int dataIndex, int value) { insertedCharacters--; } + checkLongJumpCapacity(insertedCharacters, 0); + int i = dataIndex, j = i + insertedCharacters; makeHole(i, insertedCharacters); @@ -690,6 +714,8 @@ private int insertBud(int dataIndex, int value, int segmentEnd) { insertedCharacters = 3; } + checkLongJumpCapacity(insertedCharacters, 0); + int i = dataIndex, j = i + insertedCharacters; makeHole(i, insertedCharacters); @@ -724,6 +750,8 @@ private int insertLeafLeft(int dataIndex, String key, int keyIndex, int value, i pivot++; } + checkLongJumpCapacity(insertedCharacters, 0); + int i = dataIndex, j = i + insertedCharacters; makeHole(i, insertedCharacters); @@ -764,6 +792,8 @@ private int insertLeafRight( insertedCharacters--; } + checkLongJumpCapacity(insertedCharacters, subTrieEnd - pivot >= LONG_JUMP_MARKER ? 1 : 0); + int i = dataIndex, j = i + insertedCharacters; makeHole(i, insertedCharacters); @@ -800,6 +830,8 @@ private int appendLeaf(int dataIndex, String key, int keyIndex, int value) { insertedCharacters--; } + checkLongJumpCapacity(insertedCharacters, 0); + int i = dataIndex, j = i + insertedCharacters; makeHole(i, insertedCharacters); @@ -836,14 +868,34 @@ private void makeHole(int dataIndex, int hole) { trieLength = newLength; } + /** + * Rejects the insertion before the trie is modified if it would overflow the long-jump table. + * + * @param delta how much content will be added, which shifts recorded jumps on the right + * @param newLongJumps long-jumps needed for the node where the content is being added + */ + private void checkLongJumpCapacity(int delta, int newLongJumps) { + int remaining = MAX_LONG_JUMPS - longJumpCount; + // Skip the scan if all recorded jumps can fit: each needs at most one new table entry. + if (remaining >= newLongJumps + jumpIndexCount) { + return; + } + + for (int i = 0; i < jumpIndexCount; i++) { + if (becomesLongJump(trieData[jumpIndexes[i]], delta)) { + newLongJumps++; + } + } + if (longJumpCount + newLongJumps > MAX_LONG_JUMPS) { + throw new IllegalArgumentException("Class-name does not fit in the trie"); + } + } + /** Checks a new jump value; long jumps are replaced with an index into the long-jump table. */ private char newJump(int jump) { if (jump < LONG_JUMP_MARKER) { return (char) jump; // jump is small enough to fit in the trie } - if (longJumpCount > (char) ~LONG_JUMP_MARKER) { - throw new IllegalArgumentException("Class-name does not fit in the trie"); - } if (longJumps == null) { longJumps = new int[16]; // create table on first use } else if (longJumpCount == longJumps.length) { @@ -870,7 +922,12 @@ private char updateJump(char jump, int delta) { } } - private boolean isLongJump(char jump) { + /** Returns {@code true} if increasing a short jump by the given delta needs a long-jump. */ + private static boolean becomesLongJump(char jump, int delta) { + return !isLongJump(jump) && jump + delta >= LONG_JUMP_MARKER; + } + + private static boolean isLongJump(char jump) { return (jump & LONG_JUMP_MARKER) == LONG_JUMP_MARKER; } diff --git a/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java b/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java index 4c30ae1..3caa5ef 100644 --- a/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java +++ b/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java @@ -5,6 +5,7 @@ import static java.util.stream.Collectors.joining; import static java.util.stream.Collectors.toMap; import static javax.tools.ToolProvider.getSystemJavaCompiler; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -25,6 +26,7 @@ import java.net.URLClassLoader; import java.nio.file.Files; import java.nio.file.Paths; +import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.UUID; @@ -253,9 +255,37 @@ void overflow() { @Test void tooManyLongJumps() { + String padding = padding(7700); + ClassNameTrie.Builder builder = builderAtLongJumpLimit(padding); + // last put should exceed the number of allowed long-jumps + assertThrows( + IllegalArgumentException.class, () -> builder.put("" + ((char) 67) + ((char) 2), 1)); + } + + @Test + void failedInsertionPreservesExistingMappings() { + String padding = padding(7700); + ClassNameTrie.Builder builder = builderAtLongJumpLimit(padding); + for (int j = 1; j <= 8; j++) { + assertEquals(1, builder.apply("" + (char) 67 + '\u0001' + (char) j + padding)); + } + + // A nontrivial suffix forces insertion to move existing data before allocating a long jump. + assertThrows( + IllegalArgumentException.class, () -> builder.put("" + (char) 67 + (char) 2 + padding, 1)); + + // Catching a rejected insertion must leave previously registered mappings readable. + for (int j = 1; j <= 8; j++) { + assertEquals( + 1, + builder.apply("" + (char) 67 + '\u0001' + (char) j + padding), + "Existing mapping " + j + " changed after a rejected insertion"); + } + } + + private static ClassNameTrie.Builder builderAtLongJumpLimit(String padding) { // 65 groups x 64 long-jumps/group > 4096-entry limit; each group's sub-trie spans // ~61600 chars (8 padded entries) forcing 64 small leaf entries to use long-jumps - String padding = Stream.generate(() -> " ").limit(7700).collect(joining()); ClassNameTrie.Builder builder = new ClassNameTrie.Builder(); for (int g = 1; g <= 65; g++) { char groupChar = (char) (g <= 41 ? g : g + 1); // skip '*' = chr(42) @@ -270,9 +300,116 @@ void tooManyLongJumps() { for (int j = 1; j <= 8; j++) { builder.put("" + ((char) 67) + '\u0001' + (char) j + padding, 1); } - // last put should exceed the number of allowed long-jumps - assertThrows( - IllegalArgumentException.class, () -> builder.put("" + ((char) 67) + ((char) 2), 1)); + return builder; + } + + @ParameterizedTest + @MethodSource("insertionsGrowingAncestorJump") + void failedAncestorJumpConversionPreservesTrie(String key) throws IOException { + ClassNameTrie.Builder builder = builderWithShortJumpNearLimit(0); + ByteArrayOutputStream before = new ByteArrayOutputStream(); + builder.writeTo(new DataOutputStream(before)); + + // These insertions need no local long jump, but grow an ancestor's short jump past the limit. + for (int attempt = 0; attempt < 2; attempt++) { + assertThrows(IllegalArgumentException.class, () -> builder.put(key, 2)); + ByteArrayOutputStream after = new ByteArrayOutputStream(); + builder.writeTo(new DataOutputStream(after)); + assertArrayEquals(before.toByteArray(), after.toByteArray()); + } + assertEquals(1, builder.apply("aa" + padding(7675))); + assertEquals(1, builder.apply("z")); + } + + static Stream insertionsGrowingAncestorJump() { + return Stream.of( + "aa" + padding(7675) + "x", // extend a leaf segment + "aa" + padding(100), // insert a bud within a segment + "aa", // prepend a node to a segment + "aao", // split a segment on the left + "aaq", // split a segment on the right + "az"); // add a branch + } + + @Test + void fullLongJumpTableAllowsOverwritesAndExistingJumpUpdates() { + ClassNameTrie.Builder builder = builderWithShortJumpNearLimit(1); + String key = "aa" + padding(7675); + builder.put(key + "x", 2); // consumes the final slot by converting the ancestor jump + assertEquals(2, builder.apply(key + "x")); + assertEquals(1, builder.apply(key)); + + builder.put(key, 3); // overwrite a bud + builder.put(key + "x", 4); // overwrite a leaf + builder.put(key + "xy", 5); // grows the existing long jump without allocating another + assertEquals(3, builder.apply(key)); + assertEquals(4, builder.apply(key + "x")); + assertEquals(5, builder.apply(key + "xy")); + assertEquals(1, builder.apply("z")); + } + + private static String padding(int length) { + char[] chars = new char[length]; + Arrays.fill(chars, 'p'); + return new String(chars); + } + + @Test + void fullLongJumpTableAllowsSmallInsertionsWithoutConversion() { + ClassNameTrie.Builder builder = builderWithShortJumpNearLimit(0, 10); + String key = "aa" + padding(7666); + + builder.put(key + "x", 2); // grows the ancestor jump from 61430 to 61433 + + assertEquals(1, builder.apply(key)); + assertEquals(2, builder.apply(key + "x")); + assertEquals(1, builder.apply("z")); + } + + @Test + void fullLongJumpTableAllowsBranchBeforeExistingLongJump() { + ClassNameTrie.Builder builder = builderWithShortJumpNearLimit(1); + String key = "aa" + padding(7675); + builder.put(key + "x", 2); // consumes the final slot + + builder.put("0", 3); // a collapsed branch adds no content to the existing jumps + + assertEquals(3, builder.apply("0")); + assertEquals(1, builder.apply(key)); + assertEquals(2, builder.apply(key + "x")); + assertEquals(1, builder.apply("z")); + } + + private static ClassNameTrie.Builder builderWithShortJumpNearLimit(int spareLongJumps) { + return builderWithShortJumpNearLimit(spareLongJumps, 1); + } + + private static ClassNameTrie.Builder builderWithShortJumpNearLimit( + int spareLongJumps, int distanceToLimit) { + // Root branches 'a' and 'z'; the 'a' sub-trie is just below the long-jump boundary. + // Populate the long-jump table to model a builder whose other slots are occupied. + int shortJump = 61440 - distanceToLimit; + char[] data = new char[6 + shortJump]; + data[0] = 2; + data[1] = 'a'; + data[2] = 'z'; + data[3] = 0; + data[4] = 0x8001; + data[5] = (char) shortJump; + data[6] = 8; + int offset = 0; + for (int branch = 0; branch < 8; branch++) { + int length = branch == 0 ? 7676 - distanceToLimit : 7676; + data[7 + branch] = (char) ('a' + branch); + data[15 + branch] = (char) length; + if (branch > 0) { + data[22 + branch] = (char) offset; + } + Arrays.fill(data, 30 + offset, 30 + offset + length, 'p'); + data[30 + offset + length] = 0x8001; + offset += length + 1; + } + return new ClassNameTrie.Builder(new ClassNameTrie(data, new int[4096 - spareLongJumps])); } @Test