diff --git a/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java b/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java index 4ed6f02..3be6e08 100644 --- a/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java +++ b/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java @@ -20,6 +20,24 @@ static void enableClassInjection() { ClassInjector.enableClassInjection(TestAgent.getInstrumentation()); } + @Test + void repeatedEnablementPreservesClassInjection() throws Exception { + ClassLoader loader = newCL(); + String name = "datadog.injected.RepeatedEnablementSample"; + byte[] bytecode = sampleClass(name); + Class injected = ClassInjector.injectClass(name, bytecode, loader); + + ClassInjector.enableClassInjection(TestAgent.getInstrumentation()); + + assertThat(ClassInjector.injectClass(name, bytecode, loader)).isSameAs(injected); + assertThat(injected.getMethod("value").invoke(null)).isEqualTo(42); + ClassLoader anotherLoader = newCL(); + Class another = ClassInjector.injectClass(name, bytecode, anotherLoader); + assertThat(another).isNotSameAs(injected); + assertThat(another.getClassLoader()).isSameAs(anotherLoader); + assertThat(another.getMethod("value").invoke(null)).isEqualTo(42); + } + @Test void injectsBootClass() throws Exception { String name = "datadog.injected.BootSample"; diff --git a/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreTest.java b/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreTest.java index bdfd17a..216aeda 100644 --- a/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreTest.java +++ b/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreTest.java @@ -249,6 +249,43 @@ void differentKeysHaveIndependentEntries() { assertEquals("v2", store.get(key2)); } + @Test + void equalKeysAreStoredByIdentity() { + char[] characters = {'k', 'e', 'y'}; + String first = new String(characters); + String second = new String(characters); + assertThat(first).isEqualTo(second).isNotSameAs(second); + + store.put(first, "first"); + assertNull(store.get(second)); + assertEquals("second", store.getOrPut(second, "second")); + assertEquals("first", store.get(first)); + assertEquals("second", store.get(second)); + + assertEquals("first", store.remove(first)); + assertNull(store.get(first)); + assertEquals("second", store.get(second)); + assertEquals("computed", store.getOrCompute(first, key -> "computed")); + assertEquals("computed", store.get(first)); + assertEquals("second", store.get(second)); + } + + @Test + void sameKeyWithDifferentValueTypesIsIsolated() { + Object key = new Object(); + ObjectStore strings = ObjectStore.of(Object.class, String.class); + ObjectStore integers = ObjectStore.of(Object.class, Integer.class); + + strings.put(key, "string"); + assertNull(integers.get(key)); + assertEquals(42, integers.getOrCompute(key, ignored -> 42)); + assertEquals("string", strings.get(key)); + assertEquals(42, integers.remove(key)); + assertNull(integers.get(key)); + assertEquals("string", strings.get(key)); + assertEquals("string", strings.remove(key)); + } + @Test void sameKeyInDifferentStoresIsIsolated() { ObjectStore storeA = ObjectStore.of("test.IsolA.Key", "test.IsolA.Value"); diff --git a/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java b/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java index afc02ed..31ea02f 100644 --- a/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java +++ b/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java @@ -1,5 +1,6 @@ package datadog.instrument.utils; +import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Collections.singletonList; import static java.util.function.Function.identity; import static java.util.stream.Collectors.joining; @@ -26,7 +27,9 @@ import java.net.URLClassLoader; import java.nio.file.Files; import java.nio.file.Paths; +import java.util.ArrayList; import java.util.Arrays; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.UUID; @@ -179,6 +182,106 @@ void generateEmptyJavaSource(@TempDir File classesDir) throws Exception { } } + @Test + void generateLargeJavaSource(@TempDir File directory) throws Exception { + File trieDir = new File(directory, "input"); + File javaDir = new File(directory, "output"); + Files.createDirectories(new File(trieDir, "example").toPath()); + Files.createDirectories(javaDir.toPath()); + Map mappings = largeTrieMappings(); + List lines = new ArrayList<>(); + mappings.forEach((key, value) -> lines.add(value + " " + key)); + lines.add("default.Value"); // omitted values default to one + Files.write(new File(trieDir, "example/large_names.trie").toPath(), lines, UTF_8); + + ClassNameTrie.JavaGenerator.main( + new String[] {trieDir.toString(), javaDir.toString(), "example/large_names.trie"}); + + assertEquals( + 0, + getSystemJavaCompiler() + .run( + null, + null, + null, + "-d", + javaDir.toString(), + new File(javaDir, "example/LargeNamesTrie.java").toString())); + + try (URLClassLoader cl = URLClassLoader.newInstance(new URL[] {javaDir.toURI().toURL()})) { + Class trieClass = cl.loadClass("example.LargeNamesTrie"); + Method apply = trieClass.getMethod("apply", String.class); + Method applyFromIndex = trieClass.getMethod("apply", String.class, int.class); + for (Map.Entry entry : mappings.entrySet()) { + assertEquals(entry.getValue(), apply.invoke(null, entry.getKey())); + assertEquals(entry.getValue(), applyFromIndex.invoke(null, "prefix/" + entry.getKey(), 7)); + } + assertEquals(1, apply.invoke(null, "default.Value")); + assertEquals(-1, apply.invoke(null, "missing.Type")); + assertEquals(-1, applyFromIndex.invoke(null, "prefix/missing.Type", 7)); + } + } + + @Test + void roundTripWithLongJumps() throws IOException { + Map mappings = largeTrieMappings(); + ClassNameTrie.Builder builder = new ClassNameTrie.Builder(); + mappings.forEach(builder::put); + + ClassNameTrie restored = + ClassNameTrie.readFrom(new DataInputStream(new ByteArrayInputStream(serialize(builder)))); + + for (Map.Entry entry : mappings.entrySet()) { + assertEquals(entry.getValue(), restored.apply(entry.getKey())); + assertEquals(entry.getValue(), restored.apply("prefix/" + entry.getKey(), 7)); + } + assertEquals(-1, restored.apply("missing.Type")); + assertEquals(-1, restored.apply("aaa")); + } + + private static Map largeTrieMappings() { + Map mappings = new LinkedHashMap<>(); + // The first subtree exceeds the short-jump range and the Java string chunk size. + String padding = padding(7700); + for (char branch = 'a'; branch <= 'h'; branch++) { + mappings.put("aa" + branch + padding, branch - 'a' + 1); + } + for (char branch = 'b'; branch <= 'q'; branch++) { + mappings.put("a" + branch, branch - 'a' + 10); + } + return mappings; + } + + @TableTest({ + "Scenario | First | Second | Third", + "shorten a leaf segment | abc | ab | a ", + "prefix at a branch boundary | abc | abd | ab ", + "diverge left at the last char | abc | abb | a ", + "extend a left branch | abc | abbx | ab ", + "extend then shorten a segment | abc | abcd | ab ", + "prefix before a short leaf | ab | a | ac ", + "diverge right at last char | abc | abx | aby ", + "prefix before longer branches | abxyz | abqrs | ab " + }) + void insertingPrefixesAndDivergingKeysPreservesMappings(String first, String second, String third) + throws IOException { + ClassNameTrie.Builder builder = new ClassNameTrie.Builder(); + String[] keys = {first, second, third}; + for (int i = 0; i < keys.length; i++) { + builder.put(keys[i], i + 1); + ClassNameTrie trie = builder.buildTrie(); + ClassNameTrie restored = + ClassNameTrie.readFrom(new DataInputStream(new ByteArrayInputStream(serialize(builder)))); + for (int j = 0; j <= i; j++) { + assertEquals(j + 1, builder.apply(keys[j])); + assertEquals(j + 1, trie.apply(keys[j])); + assertEquals(j + 1, restored.apply(keys[j])); + } + assertEquals(-1, trie.apply("missing")); + assertEquals(-1, trie.apply(keys[i] + "missing")); + } + } + @Test @SuppressWarnings("UnnecessaryUnicodeEscape") void manual() throws IOException {