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