diff --git a/buildSrc/src/main/kotlin/instrument-glue.gradle.kts b/buildSrc/src/main/kotlin/instrument-glue.gradle.kts index e68647b..8ed87e5 100644 --- a/buildSrc/src/main/kotlin/instrument-glue.gradle.kts +++ b/buildSrc/src/main/kotlin/instrument-glue.gradle.kts @@ -10,6 +10,7 @@ sourceSets { create("glue") { output.dir(generatedGlueResources) output.dir(generatedGlueJava) + java.include("**/glue/*") } main { resources.srcDir(generatedGlueResources) diff --git a/field-inject/build.gradle.kts b/field-inject/build.gradle.kts index ed5f91a..d0d3ce1 100644 --- a/field-inject/build.gradle.kts +++ b/field-inject/build.gradle.kts @@ -1,7 +1,29 @@ plugins { id("java-common") + id("instrument-glue") } dependencies { + + testImplementation(sourceSets["glue"].output) + testImplementation(project(":utils")) + testImplementation(libs.asm.commons) + testImplementation(libs.asm.util) + jmh("com.blogspot.mydailyjava:weak-lock-free:0.18") } + +// include GlobalObjectStore in glue generation, so we can repackage it for injection +sourceSets["glue"].java { + srcDir("src/main/java") + include("**/GlobalObjectStore.java") +} + +// field-inject generates dispatch + access glue for injection into the bootstrap classpath +extra["glue"] = listOf("ObjectStoreGlue") + +tasks.jar { + // ObjectStoreGlue only contains large string constants that get inlined into FieldInjector + // - the inlining means we can safely drop ObjectStoreGlue from the final jar to save space + excludes.add("datadog/instrument/glue/ObjectStoreGlue.class") +} diff --git a/field-inject/src/glue/java/datadog/instrument/glue/ObjectStoreGlueGenerator.java b/field-inject/src/glue/java/datadog/instrument/glue/ObjectStoreGlueGenerator.java new file mode 100644 index 0000000..115c174 --- /dev/null +++ b/field-inject/src/glue/java/datadog/instrument/glue/ObjectStoreGlueGenerator.java @@ -0,0 +1,510 @@ +/* + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). + * Copyright 2025-Present Datadog, Inc. + */ + +package datadog.instrument.glue; + +import static datadog.instrument.glue.GlueGenerator.classHeader; +import static datadog.instrument.glue.GlueGenerator.packBytecode; +import static org.objectweb.asm.ClassWriter.COMPUTE_FRAMES; +import static org.objectweb.asm.Opcodes.*; + +import datadog.instrument.utils.JVM; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.Consumer; +import org.objectweb.asm.AnnotationVisitor; +import org.objectweb.asm.ClassReader; +import org.objectweb.asm.ClassVisitor; +import org.objectweb.asm.ClassWriter; +import org.objectweb.asm.FieldVisitor; +import org.objectweb.asm.Label; +import org.objectweb.asm.MethodVisitor; +import org.objectweb.asm.Type; +import org.objectweb.asm.TypePath; +import org.objectweb.asm.commons.ClassRemapper; +import org.objectweb.asm.commons.Remapper; +import org.objectweb.asm.commons.SimpleRemapper; + +/** + * Generates {@code KeyWithValue}/{@code ObjectStoreDispatch} bytecode, plus java.lang-relocated + * copies of them and {@code GlobalObjectStore} packed as glue for the bootstrap classpath. + */ +final class ObjectStoreGlueGenerator { + + // generated dispatch glue + private static final String KEYWITHVALUE_CLASS = "datadog/instrument/fieldinject/KeyWithValue"; + private static final String OBJECTSTOREDISPATCH_CLASS = + "datadog/instrument/fieldinject/ObjectStoreDispatch"; + + // repackaged support glue + private static final String GLOBALOBJECTSTORE_CLASS = + "datadog/instrument/fieldinject/GlobalObjectStore"; + private static final String GLOBALOBJECTSTORE_STOREKEY_CLASS = + GLOBALOBJECTSTORE_CLASS + "$StoreKey"; + private static final String GLOBALOBJECTSTORE_LOOKUPKEY_CLASS = + GLOBALOBJECTSTORE_CLASS + "$LookupKey"; + private static final String GLOBALOBJECTSTORE_LOOKUPKEYCACHE_CLASS = + GLOBALOBJECTSTORE_CLASS + "$LookupKeyCache"; + private static final String GLOBALOBJECTSTORE_VALUECOMPUTER_CLASS = + GLOBALOBJECTSTORE_CLASS + "$ValueComputer"; + + private static final String OBJECT_CLASS = "java/lang/Object"; + private static final String FUNCTION_CLASS = "java/util/function/Function"; + + // bootstrap glue must be located inside the java.lang namespace for accessibility reasons + private static final String JAVA_LANG_PREFIX = "java/lang/$Datadog$"; + + private static final String INJECTED_STORE_MARKER = "__dd_instrument$"; + + private static final String GET_ACCESSOR = "$get$" + INJECTED_STORE_MARKER; + private static final String PUT_ACCESSOR = "$put$" + INJECTED_STORE_MARKER; + private static final String GET_ACCESSOR_DESCRIPTOR = "(I)L" + OBJECT_CLASS + ";"; + private static final String PUT_ACCESSOR_DESCRIPTOR = "(IL" + OBJECT_CLASS + ";)V"; + + private static final String GET_DESCRIPTOR = "(L" + OBJECT_CLASS + ";I)L" + OBJECT_CLASS + ";"; + private static final String PUT_DESCRIPTOR = "(L" + OBJECT_CLASS + ";IL" + OBJECT_CLASS + ";)V"; + private static final String GETORPUT_DESCRIPTOR = + "(L" + OBJECT_CLASS + ";IL" + OBJECT_CLASS + ";)L" + OBJECT_CLASS + ";"; + private static final String GETORCOMPUTE_DESCRIPTOR = + "(L" + OBJECT_CLASS + ";IL" + FUNCTION_CLASS + ";)L" + OBJECT_CLASS + ";"; + private static final String REMOVE_DESCRIPTOR = "(L" + OBJECT_CLASS + ";I)L" + OBJECT_CLASS + ";"; + + private static final String FUNCTION_APPLY_DESCRIPTOR = + "(L" + OBJECT_CLASS + ";)L" + OBJECT_CLASS + ";"; + + private ObjectStoreGlueGenerator() {} + + /** + * Generates glue to dispatch field-injected values, and writes it to the given location. + * + * @param resourcePath where to write resource files + * @param javaPath where to write Java files + * @throws IOException if the files cannot be written + * @see GlueGenerator#main + */ + public static void generateGlue(Path resourcePath, Path javaPath) throws IOException { + byte[] keyWithValueBytecode = generateKeyWithValueBytecode(); + byte[] objectStoreDispatchBytecode = generateObjectStoreDispatchBytecode(); + + // first generate dispatch glue for the field-inject package + Path fieldInjectPackagePath = resourcePath.resolveSibling("fieldinject"); + Files.createDirectories(fieldInjectPackagePath); + writeClassResource(fieldInjectPackagePath, KEYWITHVALUE_CLASS, keyWithValueBytecode); + writeClassResource( + fieldInjectPackagePath, OBJECTSTOREDISPATCH_CLASS, objectStoreDispatchBytecode); + + // next collect all the glue we need to repackage for the bootstrap. + List relocatedClasses = new ArrayList<>(); + relocatedClasses.add( + new RelocatedClass("KEYWITHVALUE", KEYWITHVALUE_CLASS, keyWithValueBytecode)); + relocatedClasses.add( + new RelocatedClass( + "OBJECTSTOREDISPATCH", OBJECTSTOREDISPATCH_CLASS, objectStoreDispatchBytecode)); + relocatedClasses.add( + new RelocatedClass( + "GLOBALOBJECTSTORE", + GLOBALOBJECTSTORE_CLASS, + readClassResource(GLOBALOBJECTSTORE_CLASS))); + relocatedClasses.add( + new RelocatedClass( + "GLOBALOBJECTSTORE_STOREKEY", + GLOBALOBJECTSTORE_STOREKEY_CLASS, + readClassResource(GLOBALOBJECTSTORE_STOREKEY_CLASS))); + relocatedClasses.add( + new RelocatedClass( + "GLOBALOBJECTSTORE_LOOKUPKEY", + GLOBALOBJECTSTORE_LOOKUPKEY_CLASS, + readClassResource(GLOBALOBJECTSTORE_LOOKUPKEY_CLASS))); + relocatedClasses.add( + new RelocatedClass( + "GLOBALOBJECTSTORE_LOOKUPKEYCACHE", + GLOBALOBJECTSTORE_LOOKUPKEYCACHE_CLASS, + readClassResource(GLOBALOBJECTSTORE_LOOKUPKEYCACHE_CLASS))); + relocatedClasses.add( + new RelocatedClass( + "GLOBALOBJECTSTORE_VALUECOMPUTER", + GLOBALOBJECTSTORE_VALUECOMPUTER_CLASS, + readClassResource(GLOBALOBJECTSTORE_VALUECOMPUTER_CLASS))); + + Map renames = new HashMap<>(); + for (RelocatedClass relocated : relocatedClasses) { + renames.put(relocated.originalClass, relocated.relocatedClass); + } + Remapper remapper = new SimpleRemapper(ASM9, renames); + + // make the bootstrap glue available for class injection + List lines = new ArrayList<>(); + classHeader(lines, "ObjectStoreGlue"); + lines.add(" /** Shared prefix of bootstrap glue */"); + lines.add(" String PREFIX = \"" + JAVA_LANG_PREFIX.replace('/', '.') + "\";"); + for (RelocatedClass relocated : relocatedClasses) { + packRelocatedClass(lines, remapper, relocated); + } + lines.add("}"); + Files.write(javaPath.resolve("ObjectStoreGlue.java"), lines, StandardCharsets.UTF_8); + } + + /** A class to be relocated into the {@code java.lang} namespace and packed into the glue. */ + private static final class RelocatedClass { + final String constantName; + final String originalClass; + final String relocatedClass; + final byte[] originalBytecode; + + RelocatedClass(String constantName, String originalClass, byte[] originalBytecode) { + String simpleName = originalClass.substring(originalClass.lastIndexOf('/') + 1); + this.constantName = constantName; + this.originalClass = originalClass; + this.relocatedClass = JAVA_LANG_PREFIX + simpleName; + this.originalBytecode = originalBytecode; + } + } + + /** Generates bytecode for the {@code KeyWithValue} marker interface. */ + static byte[] generateKeyWithValueBytecode() { + ClassWriter cw = new ClassWriter(0); + cw.visit( + V1_8, + ACC_PUBLIC | ACC_INTERFACE | ACC_ABSTRACT | ACC_SYNTHETIC, + KEYWITHVALUE_CLASS, + null, + OBJECT_CLASS, + null); + cw.visitMethod(ACC_PUBLIC | ACC_ABSTRACT, GET_ACCESSOR, GET_ACCESSOR_DESCRIPTOR, null, null) + .visitEnd(); + cw.visitMethod(ACC_PUBLIC | ACC_ABSTRACT, PUT_ACCESSOR, PUT_ACCESSOR_DESCRIPTOR, null, null) + .visitEnd(); + cw.visitEnd(); + return cw.toByteArray(); + } + + /** Generates bytecode for {@code ObjectStoreDispatch}. */ + static byte[] generateObjectStoreDispatchBytecode() { + ClassWriter cw = new ClassWriter(COMPUTE_FRAMES); + cw.visit( + V1_8, + ACC_PUBLIC | ACC_FINAL | ACC_SYNTHETIC, + OBJECTSTOREDISPATCH_CLASS, + null, + OBJECT_CLASS, + null); + + // no constructor: this class is only ever used via its static dispatch methods + + // simple two-way dispatch: fast path calls the accessor directly, otherwise fall back + generateCheckedDispatch(cw, "get", GET_DESCRIPTOR); + generateCheckedDispatch(cw, "put", PUT_DESCRIPTOR); + + // double-checked get-then-put dispatch, shared by getOrPut/getOrCompute/remove + generateDoubleCheckedDispatch( + cw, "getOrPut", GETORPUT_DESCRIPTOR, mv -> mv.visitVarInsn(ALOAD, 2)); + generateDoubleCheckedDispatch( + cw, + "getOrCompute", + GETORCOMPUTE_DESCRIPTOR, + mv -> { + mv.visitVarInsn(ALOAD, 2); + mv.visitVarInsn(ALOAD, 0); + mv.visitMethodInsn( + INVOKEINTERFACE, FUNCTION_CLASS, "apply", FUNCTION_APPLY_DESCRIPTOR, true); + }); + generateDoubleCheckedDispatch(cw, "remove", REMOVE_DESCRIPTOR, mv -> mv.visitInsn(ACONST_NULL)); + + // weakGet/weakPut skip the fast-path accessor and always go straight to the global store + generateForwardDispatch(cw, "weakGet", "get", GET_DESCRIPTOR); + generateForwardDispatch(cw, "weakPut", "put", PUT_DESCRIPTOR); + + cw.visitEnd(); + return cw.toByteArray(); + } + + /** Whether the descriptor's 3rd argument is the value to store, by convention. */ + private static boolean hasValueArg(String descriptor) { + return Type.getArgumentTypes(descriptor).length > 2; + } + + /** Generates a method that unconditionally delegates to {@code GlobalObjectStore}. */ + private static void generateForwardDispatch( + ClassWriter cw, String method, String forwardMethod, String descriptor) { + final boolean hasValueArg = hasValueArg(descriptor); + final int returnOpcode = Type.getReturnType(descriptor).getOpcode(IRETURN); + + final MethodVisitor mv = + cw.visitMethod(ACC_PUBLIC | ACC_STATIC, method, descriptor, null, null); + + mv.visitCode(); + mv.visitVarInsn(ALOAD, 0); + emitFallback(mv, forwardMethod, descriptor, hasValueArg, returnOpcode); + mv.visitMaxs(-1, -1); + mv.visitEnd(); + } + + /** Generates a getter or putter that dispatches to the fast-path accessor or the fallback. */ + private static void generateCheckedDispatch(ClassWriter cw, String method, String descriptor) { + final boolean hasValueArg = hasValueArg(descriptor); + final int returnOpcode = Type.getReturnType(descriptor).getOpcode(IRETURN); + final String accessorMethod = hasValueArg ? PUT_ACCESSOR : GET_ACCESSOR; + final String accessorDescriptor = + hasValueArg ? PUT_ACCESSOR_DESCRIPTOR : GET_ACCESSOR_DESCRIPTOR; + + final MethodVisitor mv = + cw.visitMethod(ACC_PUBLIC | ACC_STATIC, method, descriptor, null, null); + + final Label fallback = new Label(); + + mv.visitCode(); + emitInstanceOfCheck(mv, fallback); + + // fast path: ((KeyWithValue) key).accessorMethod(storeId[, value]) + mv.visitTypeInsn(CHECKCAST, KEYWITHVALUE_CLASS); + mv.visitVarInsn(ILOAD, 1); + if (hasValueArg) { + mv.visitVarInsn(ALOAD, 2); + } + mv.visitMethodInsn( + INVOKEINTERFACE, KEYWITHVALUE_CLASS, accessorMethod, accessorDescriptor, true); + mv.visitInsn(returnOpcode); + + mv.visitLabel(fallback); + emitFallback(mv, method, descriptor, hasValueArg, returnOpcode); + + mv.visitMaxs(-1, -1); + mv.visitEnd(); + } + + /** + * Generates the double-checked get/put dispatch shared by {@code getOrPut}, {@code getOrCompute} + * and {@code remove}; {@code valueToStore} emits the value to put on the slow path. + */ + private static void generateDoubleCheckedDispatch( + ClassWriter cw, String method, String descriptor, Consumer valueToStore) { + + final boolean hasValueArg = hasValueArg(descriptor); + final int accessorLocal = hasValueArg ? 3 : 2; + final int existingLocal = accessorLocal + 1; + final int storedLocal = hasValueArg ? existingLocal : existingLocal + 1; + final int skipToReturnOpcode = hasValueArg ? IFNONNULL : IFNULL; + + final MethodVisitor mv = + cw.visitMethod(ACC_PUBLIC | ACC_STATIC, method, descriptor, null, null); + + final Label fallback = new Label(); + final Label returnExisting = new Label(); + final Label skipPut = new Label(); + final Label monitorEntered = new Label(); + final Label monitorExited = new Label(); + final Label unlockAndRethrow = new Label(); + + mv.visitCode(); + mv.visitTryCatchBlock(monitorEntered, monitorExited, unlockAndRethrow, null); + + emitInstanceOfCheck(mv, fallback); + + // KeyWithValue accessor = (KeyWithValue) key; + mv.visitTypeInsn(CHECKCAST, KEYWITHVALUE_CLASS); + mv.visitVarInsn(ASTORE, accessorLocal); + + // Object existing = accessor.$dd_instrument_$get$(storeId); + emitAccessorGet(mv, accessorLocal, existingLocal); + + // if (existing [is/isn't] null) return existing; + mv.visitVarInsn(ALOAD, existingLocal); + mv.visitJumpInsn(skipToReturnOpcode, returnExisting); + + // synchronized (accessor) { + mv.visitVarInsn(ALOAD, accessorLocal); + mv.visitInsn(MONITORENTER); + mv.visitLabel(monitorEntered); + + // existing = accessor.$dd_instrument_$get$(storeId); (re-check now we hold the lock) + emitAccessorGet(mv, accessorLocal, existingLocal); + + mv.visitVarInsn(ALOAD, existingLocal); + mv.visitJumpInsn(skipToReturnOpcode, skipPut); + + // stored = ; accessor.$dd_instrument_$put$(storeId, stored); + valueToStore.accept(mv); + mv.visitVarInsn(ASTORE, storedLocal); + mv.visitVarInsn(ALOAD, accessorLocal); + mv.visitVarInsn(ILOAD, 1); + mv.visitVarInsn(ALOAD, storedLocal); + mv.visitMethodInsn( + INVOKEINTERFACE, KEYWITHVALUE_CLASS, PUT_ACCESSOR, PUT_ACCESSOR_DESCRIPTOR, true); + + mv.visitLabel(skipPut); + mv.visitVarInsn(ALOAD, accessorLocal); + mv.visitInsn(MONITOREXIT); + mv.visitLabel(monitorExited); + mv.visitJumpInsn(GOTO, returnExisting); + + // } unlock and rethrow if anything above failed + mv.visitLabel(unlockAndRethrow); + mv.visitVarInsn(ALOAD, accessorLocal); + mv.visitInsn(MONITOREXIT); + mv.visitInsn(ATHROW); + + mv.visitLabel(returnExisting); + mv.visitVarInsn(ALOAD, existingLocal); + mv.visitInsn(ARETURN); + + mv.visitLabel(fallback); + emitFallback(mv, method, descriptor, hasValueArg, ARETURN); + + mv.visitMaxs(-1, -1); + mv.visitEnd(); + } + + /** + * Emits {@code return GlobalObjectStore.method(key, storeId[, value]);} reusing the key reference + * already left on the stack by {@link #emitInstanceOfCheck}. + */ + private static void emitFallback( + MethodVisitor mv, String method, String descriptor, boolean hasValueArg, int returnOpcode) { + mv.visitVarInsn(ILOAD, 1); + if (hasValueArg) { + mv.visitVarInsn(ALOAD, 2); + } + mv.visitMethodInsn(INVOKESTATIC, GLOBALOBJECTSTORE_CLASS, method, descriptor, false); + mv.visitInsn(returnOpcode); + } + + /** + * Emits {@code if (!(key instanceof KeyWithValue)) goto notInjected;} leaving the key reference + * on the stack, so callers don't need to reload it with {@code ALOAD 0}. + */ + private static void emitInstanceOfCheck(MethodVisitor mv, Label notInjected) { + mv.visitVarInsn(ALOAD, 0); + mv.visitInsn(DUP); + mv.visitTypeInsn(INSTANCEOF, KEYWITHVALUE_CLASS); + mv.visitJumpInsn(IFEQ, notInjected); + } + + /** Emits {@code existing = accessor.$dd_instrument_$get$(storeId);} */ + private static void emitAccessorGet(MethodVisitor mv, int accessorLocal, int existingLocal) { + mv.visitVarInsn(ALOAD, accessorLocal); + mv.visitVarInsn(ILOAD, 1); + mv.visitMethodInsn( + INVOKEINTERFACE, KEYWITHVALUE_CLASS, GET_ACCESSOR, GET_ACCESSOR_DESCRIPTOR, true); + mv.visitVarInsn(ASTORE, existingLocal); + } + + /** Packs a relocated class as a string constant into the generated Java source lines. */ + private static void packRelocatedClass( + List lines, Remapper remapper, RelocatedClass relocated) { + String constantName = relocated.constantName; + lines.add(" /** Packed bytecode for relocated " + constantName + " */"); + lines.add(" String " + constantName + " ="); + packBytecode(lines, relocate(remapper, relocated.originalBytecode)); + } + + /** Relocates the given bytecode with the given remapper, dropping non-essential metadata. */ + private static byte[] relocate(Remapper remapper, byte[] originalBytecode) { + ClassReader cr = new ClassReader(originalBytecode); + ClassWriter cw = new ClassWriter(0); + cr.accept(new MetadataPruner(new ClassRemapper(cw, remapper)), ClassReader.SKIP_DEBUG); + return cw.toByteArray(); + } + + /** Drops annotations and generic signatures, which are unused by the relocated glue. */ + private static final class MetadataPruner extends ClassVisitor { + MetadataPruner(ClassVisitor cv) { + super(ASM9, cv); + } + + @Override + public void visit( + int version, + int access, + String name, + String signature, + String superName, + String[] interfaces) { + super.visit(version, access, name, null, superName, interfaces); + } + + @Override + public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) { + return null; + } + + @Override + public AnnotationVisitor visitTypeAnnotation( + int typeRef, TypePath typePath, String descriptor, boolean visible) { + return null; + } + + @Override + public FieldVisitor visitField( + int access, String name, String descriptor, String signature, Object value) { + return new FieldVisitor(ASM9, super.visitField(access, name, descriptor, null, value)) { + @Override + public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) { + return null; + } + + @Override + public AnnotationVisitor visitTypeAnnotation( + int typeRef, TypePath typePath, String descriptor, boolean visible) { + return null; + } + }; + } + + @Override + public MethodVisitor visitMethod( + int access, String name, String descriptor, String signature, String[] exceptions) { + return new MethodVisitor( + ASM9, super.visitMethod(access, name, descriptor, null, exceptions)) { + @Override + public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) { + return null; + } + + @Override + public AnnotationVisitor visitTypeAnnotation( + int typeRef, TypePath typePath, String descriptor, boolean visible) { + return null; + } + + @Override + public void visitAnnotableParameterCount(int parameterCount, boolean visible) {} + + @Override + public AnnotationVisitor visitParameterAnnotation( + int parameter, String descriptor, boolean visible) { + return null; + } + }; + } + } + + /** Reads the bytecode of an already-compiled class off the classpath. */ + private static byte[] readClassResource(String internalName) throws IOException { + try (InputStream in = + ObjectStoreGlueGenerator.class + .getClassLoader() + .getResourceAsStream(internalName + ".class")) { + if (in == null) { + throw new IOException("Could not find class on classpath: " + internalName); + } + return JVM.readAllBytes(in); + } + } + + /** Writes the given bytecode as a resource class file under the given package path. */ + private static void writeClassResource(Path packagePath, String internalName, byte[] bytecode) + throws IOException { + String simpleName = internalName.substring(internalName.lastIndexOf('/') + 1); + Files.write(packagePath.resolve(simpleName + ".class"), bytecode); + } +} diff --git a/field-inject/src/main/java/datadog/instrument/fieldinject/GlobalObjectStore.java b/field-inject/src/main/java/datadog/instrument/fieldinject/GlobalObjectStore.java index 87f4400..657b3b1 100644 --- a/field-inject/src/main/java/datadog/instrument/fieldinject/GlobalObjectStore.java +++ b/field-inject/src/main/java/datadog/instrument/fieldinject/GlobalObjectStore.java @@ -27,7 +27,7 @@ *

Each shard is split into two maps with separate reference queues: young and old. Ageing a * shard by one generation creates a new young map; the previous young map becomes the old map. */ -public final class GlobalObjectStore { +final class GlobalObjectStore { /** Target ceiling for the total number of objects in a shard, young and old. */ static final int SHARD_HARD_LIMIT = 32_000; diff --git a/field-inject/src/main/java/datadog/instrument/fieldinject/ObjectStoreIds.java b/field-inject/src/main/java/datadog/instrument/fieldinject/ObjectStoreIds.java index 9b78578..c35773d 100644 --- a/field-inject/src/main/java/datadog/instrument/fieldinject/ObjectStoreIds.java +++ b/field-inject/src/main/java/datadog/instrument/fieldinject/ObjectStoreIds.java @@ -12,7 +12,7 @@ import java.util.function.Function; /** Manages unique {@link ObjectStore} ids for each key-value type combination. */ -public final class ObjectStoreIds { +final class ObjectStoreIds { private static final TypeIdGenerator typeIdGenerator = new TypeIdGenerator(); private static final Map typeIds = new ConcurrentHashMap<>(); diff --git a/field-inject/src/main/java/datadog/instrument/fieldinject/package-info.java b/field-inject/src/main/java/datadog/instrument/fieldinject/package-info.java index 7b78b2f..46be493 100644 --- a/field-inject/src/main/java/datadog/instrument/fieldinject/package-info.java +++ b/field-inject/src/main/java/datadog/instrument/fieldinject/package-info.java @@ -1,4 +1,4 @@ -/** Supports Key-value object stores where keys and values have specific types. */ +/** Supports key-value object stores where keys and values have specific types. */ @ParametersAreNonnullByDefault package datadog.instrument.fieldinject; diff --git a/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreDispatchFunctionalTest.java b/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreDispatchFunctionalTest.java new file mode 100644 index 0000000..979e32b --- /dev/null +++ b/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreDispatchFunctionalTest.java @@ -0,0 +1,324 @@ +package datadog.instrument.fieldinject; + +import static datadog.instrument.fieldinject.ObjectStoreIds.objectStoreId; +import static org.assertj.core.api.Assertions.assertThat; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +/** + * Loads and actually invokes the generated {@code ObjectStoreDispatch} bytecode, to check its + * fast-path/fallback dispatch and double-checked locking behave correctly - not just that the + * bytecode passes structural verification (see {@code ObjectStoreGlueGeneratorTest}). + * + *

The generated classes are loaded from the build-time glue on the test classpath, so they share + * the same class loader (and runtime package) as the package-private {@link GlobalObjectStore}. + */ +class ObjectStoreDispatchFunctionalTest { + + private static Class keyWithValueClass; + private static Method getMethod; + private static Method putMethod; + private static Method getOrPutMethod; + private static Method getOrComputeMethod; + private static Method removeMethod; + private static Method weakGetMethod; + private static Method weakPutMethod; + + @BeforeAll + static void loadGeneratedClasses() throws Exception { + keyWithValueClass = Class.forName("datadog.instrument.fieldinject.KeyWithValue"); + Class dispatchClass = Class.forName("datadog.instrument.fieldinject.ObjectStoreDispatch"); + + getMethod = dispatchClass.getMethod("get", Object.class, int.class); + putMethod = dispatchClass.getMethod("put", Object.class, int.class, Object.class); + getOrPutMethod = dispatchClass.getMethod("getOrPut", Object.class, int.class, Object.class); + getOrComputeMethod = + dispatchClass.getMethod("getOrCompute", Object.class, int.class, Function.class); + removeMethod = dispatchClass.getMethod("remove", Object.class, int.class); + weakGetMethod = dispatchClass.getMethod("weakGet", Object.class, int.class); + weakPutMethod = dispatchClass.getMethod("weakPut", Object.class, int.class, Object.class); + } + + @Test + void fastPathRoundTripsThroughAccessorAndBypassesGlobalStore() throws Exception { + int storeId = objectStoreId("test.Dispatch.FastPath.Key", "test.Dispatch.FastPath.Value"); + Map backing = Collections.synchronizedMap(new HashMap<>()); + Object key = fastPathKey(backing); + + // seed the global store, so we can tell if the fast path ever reads from or writes to it + GlobalObjectStore.put(key, storeId, "global"); + + assertThat(getMethod.invoke(null, key, storeId)) + .as("fast path get must not fall back to the global store") + .isNull(); + + putMethod.invoke(null, key, storeId, "value"); + + assertThat(backing.get(storeId)).isEqualTo("value"); + assertThat(getMethod.invoke(null, key, storeId)).isEqualTo("value"); + assertThat(GlobalObjectStore.get(key, storeId)) + .as("fast path put must not write to the global store") + .isEqualTo("global"); + } + + @Test + void fallbackPathDelegatesAllOperationsToGlobalObjectStoreForNonInjectedKey() throws Exception { + int storeId = objectStoreId("test.Dispatch.Fallback.Key", "test.Dispatch.Fallback.Value"); + Object key = new Object(); + + putMethod.invoke(null, key, storeId, "value"); + + assertThat(GlobalObjectStore.get(key, storeId)).isEqualTo("value"); + assertThat(getMethod.invoke(null, key, storeId)).isEqualTo("value"); + assertThat(getOrPutMethod.invoke(null, key, storeId, "ignored")).isEqualTo("value"); + assertThat(removeMethod.invoke(null, key, storeId)).isEqualTo("value"); + assertThat(GlobalObjectStore.get(key, storeId)).isNull(); + + assertThat(getOrPutMethod.invoke(null, key, storeId, "put-if-absent")) + .isEqualTo("put-if-absent"); + assertThat(GlobalObjectStore.get(key, storeId)).isEqualTo("put-if-absent"); + assertThat(removeMethod.invoke(null, key, storeId)).isEqualTo("put-if-absent"); + + AtomicInteger computeCalls = new AtomicInteger(); + Function compute = + k -> { + computeCalls.incrementAndGet(); + return "computed"; + }; + assertThat(getOrComputeMethod.invoke(null, key, storeId, compute)).isEqualTo("computed"); + assertThat(getOrComputeMethod.invoke(null, key, storeId, compute)).isEqualTo("computed"); + assertThat(computeCalls).hasValue(1); + assertThat(GlobalObjectStore.get(key, storeId)).isEqualTo("computed"); + } + + @Test + void getOrPutOnlySetsAccessorWhenAbsent() throws Exception { + int storeId = objectStoreId("test.Dispatch.GetOrPut.Key", "test.Dispatch.GetOrPut.Value"); + Map backing = Collections.synchronizedMap(new HashMap<>()); + Object key = fastPathKey(backing); + + assertThat(getOrPutMethod.invoke(null, key, storeId, "first")).isEqualTo("first"); + assertThat(backing.get(storeId)).isEqualTo("first"); + + assertThat(getOrPutMethod.invoke(null, key, storeId, "second")).isEqualTo("first"); + assertThat(backing.get(storeId)) + .as("existing accessor value must not be overwritten") + .isEqualTo("first"); + } + + @Test + void getOrComputeInvokesFunctionWithOriginalKeyOnlyWhenAbsent() throws Exception { + int storeId = + objectStoreId("test.Dispatch.GetOrCompute.Key", "test.Dispatch.GetOrCompute.Value"); + Map backing = Collections.synchronizedMap(new HashMap<>()); + Object key = fastPathKey(backing); + AtomicReference receivedKey = new AtomicReference<>(); + + Function computeOnce = + k -> { + receivedKey.set(k); + return "computed"; + }; + assertThat(getOrComputeMethod.invoke(null, key, storeId, computeOnce)).isEqualTo("computed"); + assertThat(receivedKey.get()) + .as("function should receive the original key, not the accessor") + .isSameAs(key); + assertThat(backing.get(storeId)).isEqualTo("computed"); + + AtomicBoolean called = new AtomicBoolean(false); + Function shouldNotRun = + k -> { + called.set(true); + return "should-not-be-used"; + }; + assertThat(getOrComputeMethod.invoke(null, key, storeId, shouldNotRun)).isEqualTo("computed"); + assertThat(called).as("compute function must not run when a value already exists").isFalse(); + } + + @Test + void removeReturnsPreviousValueAndClearsAccessor() throws Exception { + int storeId = objectStoreId("test.Dispatch.Remove.Key", "test.Dispatch.Remove.Value"); + Map backing = Collections.synchronizedMap(new HashMap<>()); + Object key = fastPathKey(backing); + backing.put(storeId, "existing"); + + assertThat(removeMethod.invoke(null, key, storeId)).isEqualTo("existing"); + assertThat(backing.get(storeId)).isNull(); + } + + @Test + void removeDoesNotWriteToAccessorWhenValueIsAbsent() throws Exception { + int storeId = + objectStoreId("test.Dispatch.RemoveAbsent.Key", "test.Dispatch.RemoveAbsent.Value"); + AtomicInteger setterCalls = new AtomicInteger(); + InvocationHandler handler = + (proxy, method, args) -> { + if (args.length == 1) { + return null; + } + setterCalls.incrementAndGet(); + return null; + }; + Object key = + Proxy.newProxyInstance( + keyWithValueClass.getClassLoader(), new Class[] {keyWithValueClass}, handler); + + assertThat(removeMethod.invoke(null, key, storeId)).isNull(); + assertThat(setterCalls).hasValue(0); + } + + @Test + void injectedAccessorFallsBackViaWeakGetPutForStoresItDoesNotHave() throws Exception { + int injectedId = objectStoreId("test.Dispatch.Weak.Key", "test.Dispatch.Weak.Injected"); + int fallbackId = objectStoreId("test.Dispatch.Weak.Key", "test.Dispatch.Weak.Fallback"); + Map backing = Collections.synchronizedMap(new HashMap<>()); + Object key = injectedKey(backing, Collections.singleton(injectedId)); + + putMethod.invoke(null, key, fallbackId, "fallback"); + assertThat(getMethod.invoke(null, key, fallbackId)).isEqualTo("fallback"); + assertThat(GlobalObjectStore.get(key, fallbackId)).isEqualTo("fallback"); + + assertThat(getOrPutMethod.invoke(null, key, fallbackId, "ignored")).isEqualTo("fallback"); + assertThat(removeMethod.invoke(null, key, fallbackId)).isEqualTo("fallback"); + assertThat(GlobalObjectStore.get(key, fallbackId)) + .as("remove must clear the global store via weakPut") + .isNull(); + assertThat( + getOrComputeMethod.invoke( + null, key, fallbackId, (Function) k -> "computed")) + .isEqualTo("computed"); + assertThat(GlobalObjectStore.get(key, fallbackId)).isEqualTo("computed"); + + // stores the accessor does have should still use the injected field + putMethod.invoke(null, key, injectedId, "injected"); + assertThat(backing).containsOnlyKeys(injectedId); + assertThat(GlobalObjectStore.get(key, injectedId)).isNull(); + } + + @Test + void concurrentGetOrPutOnSameKeyConvergesToASingleWinner() throws Exception { + int storeId = objectStoreId("test.Dispatch.Concurrent.Key", "test.Dispatch.Concurrent.Value"); + Map backing = Collections.synchronizedMap(new HashMap<>()); + int threadCount = 8; + CountDownLatch initialReads = new CountDownLatch(threadCount); + AtomicInteger getterCalls = new AtomicInteger(); + AtomicInteger setterCalls = new AtomicInteger(); + InvocationHandler handler = + (proxy, method, args) -> { + if (args.length == 1) { + if (getterCalls.incrementAndGet() <= threadCount) { + initialReads.countDown(); + if (!initialReads.await(10, TimeUnit.SECONDS)) { + throw new AssertionError("not every contender reached the initial accessor read"); + } + } + return backing.get((Integer) args[0]); + } + setterCalls.incrementAndGet(); + backing.put((Integer) args[0], args[1]); + return null; + }; + Object key = + Proxy.newProxyInstance( + keyWithValueClass.getClassLoader(), new Class[] {keyWithValueClass}, handler); + + CountDownLatch start = new CountDownLatch(1); + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + List results = new ArrayList<>(); + try { + List> futures = new ArrayList<>(); + for (int t = 0; t < threadCount; t++) { + String candidate = "candidate-" + t; + futures.add( + executor.submit( + () -> { + start.await(); + return getOrPutMethod.invoke(null, key, storeId, candidate); + })); + } + start.countDown(); + for (java.util.concurrent.Future future : futures) { + results.add(future.get(10, TimeUnit.SECONDS)); + } + } finally { + executor.shutdownNow(); + } + + Object winner = backing.get(storeId); + assertThat(setterCalls) + .as("only the first contender inside the lock should set the value") + .hasValue(1); + assertThat(results) + .as("double-checked locking should make every caller see the same winning value") + .allMatch(winner::equals); + } + + /** + * Creates a {@code KeyWithValue} proxy backed by a plain map, simulating a field-injected key. + */ + private static Object fastPathKey(Map backing) { + InvocationHandler handler = + (proxy, method, args) -> { + if (args.length == 1) { + return backing.get((Integer) args[0]); + } + backing.put((Integer) args[0], args[1]); + return null; + }; + return Proxy.newProxyInstance( + keyWithValueClass.getClassLoader(), new Class[] {keyWithValueClass}, handler); + } + + /** + * Mimics a field-injected key: stores it has fields for use the backing map, anything else falls + * back via {@code ObjectStoreDispatch.weakGet/weakPut} (injected code can't assume it has access + * to {@code GlobalObjectStore}). Fails fast if the accessor is re-entered by a recursive + * fallback. + */ + private static Object injectedKey(Map backing, Set injectedIds) { + AtomicInteger depth = new AtomicInteger(); + InvocationHandler handler = + (proxy, method, args) -> { + if (depth.incrementAndGet() > 1) { + // fail fast, rather than recursing until the stack overflows + throw new AssertionError("weakGet/weakPut must not re-enter the accessor"); + } + try { + Integer storeId = (Integer) args[0]; + if (args.length == 1) { + return injectedIds.contains(storeId) + ? backing.get(storeId) + : weakGetMethod.invoke(null, proxy, storeId); + } + if (injectedIds.contains(storeId)) { + backing.put(storeId, args[1]); + } else { + weakPutMethod.invoke(null, proxy, storeId, args[1]); + } + return null; + } finally { + depth.decrementAndGet(); + } + }; + return Proxy.newProxyInstance( + keyWithValueClass.getClassLoader(), new Class[] {keyWithValueClass}, handler); + } +} diff --git a/field-inject/src/test/java/datadog/instrument/glue/ObjectStoreGlueGeneratorTest.java b/field-inject/src/test/java/datadog/instrument/glue/ObjectStoreGlueGeneratorTest.java new file mode 100644 index 0000000..aef388a --- /dev/null +++ b/field-inject/src/test/java/datadog/instrument/glue/ObjectStoreGlueGeneratorTest.java @@ -0,0 +1,150 @@ +package datadog.instrument.glue; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.objectweb.asm.Opcodes.ASM9; + +import java.io.PrintWriter; +import java.io.StringWriter; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.lang.reflect.Modifier; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.TreeSet; +import org.junit.jupiter.api.Test; +import org.objectweb.asm.ClassReader; +import org.objectweb.asm.ClassVisitor; +import org.objectweb.asm.ClassWriter; +import org.objectweb.asm.MethodVisitor; +import org.objectweb.asm.Opcodes; +import org.objectweb.asm.Type; +import org.objectweb.asm.commons.ClassRemapper; +import org.objectweb.asm.commons.Remapper; +import org.objectweb.asm.util.CheckClassAdapter; + +class ObjectStoreGlueGeneratorTest { + + @Test + void keyWithValueBytecodePassesVerification() { + verifyAll(ObjectStoreGlueGenerator.generateKeyWithValueBytecode()); + } + + @Test + void objectStoreDispatchBytecodePassesVerification() { + verifyAll(ObjectStoreGlueGenerator.generateObjectStoreDispatchBytecode()); + } + + @Test + void objectStoreDispatchFallbacksMatchGlobalObjectStoreSignatures() throws Exception { + // GlobalObjectStore is package-private, so look it up by name + Class globalObjectStore = Class.forName("datadog.instrument.fieldinject.GlobalObjectStore"); + + // collect every static call the generated dispatch makes to GlobalObjectStore + String globalObjectStoreClass = Type.getInternalName(globalObjectStore); + Set fallbackCalls = new TreeSet<>(); + new ClassReader(ObjectStoreGlueGenerator.generateObjectStoreDispatchBytecode()) + .accept( + new ClassVisitor(ASM9) { + @Override + public MethodVisitor visitMethod( + int access, + String name, + String descriptor, + String signature, + String[] exceptions) { + return new MethodVisitor(ASM9) { + @Override + public void visitMethodInsn( + int opcode, + String owner, + String name, + String descriptor, + boolean isInterface) { + if (opcode == Opcodes.INVOKESTATIC && owner.equals(globalObjectStoreClass)) { + fallbackCalls.add(name + descriptor); + } + } + }; + } + }, + 0); + + // the dispatch lives in the same package, so it can call any non-private static method + Set accessibleStaticMethods = new TreeSet<>(); + for (Method method : globalObjectStore.getDeclaredMethods()) { + int modifiers = method.getModifiers(); + if (Modifier.isStatic(modifiers) && !Modifier.isPrivate(modifiers)) { + accessibleStaticMethods.add(method.getName() + Type.getMethodDescriptor(method)); + } + } + + assertThat(fallbackCalls).isNotEmpty(); + assertThat(accessibleStaticMethods) + .as("generated fallbacks must match non-private static methods on GlobalObjectStore") + .containsAll(fallbackCalls); + } + + /** Perform strict class verification, like -Xverify:all */ + private static void verifyAll(byte[] bytecode) { + StringWriter errors = new StringWriter(); + CheckClassAdapter.verify(new ClassReader(bytecode), false, new PrintWriter(errors)); + assertThat(errors.toString()).isEmpty(); + } + + @Test + void relocatedGlueClassesPassVerification() throws Exception { + List relocatedBytecode = relocatedBytecode(); + + // KeyWithValue, ObjectStoreDispatch, GlobalObjectStore + $StoreKey + $LookupKey + assertThat(relocatedBytecode).hasSize(7); + + // the relocated classes reference each other by their new java.lang names, which can't be + // loaded via a normal class loader (the JVM forbids defining classes into java.* packages + // outside the bootstrap loader), so we can't run full type-resolving verification here; + // fall back to structural verification, which still catches corrupted/inconsistent bytecode + for (byte[] bytecode : relocatedBytecode) { + verifyStructure(bytecode); + // relocated classes must live under the shared java.lang bootstrap namespace + assertThat(new ClassReader(bytecode).getClassName()).startsWith("java/lang/$Datadog$"); + } + } + + @Test + void relocatedGlueClassesOnlyReferenceRelocatedTypes() throws Exception { + // catches nested/helper classes missing from the relocation list, which would otherwise + // only fail at runtime with NoClassDefFoundError when loaded from the bootstrap classpath + Set unrelocatedTypes = new TreeSet<>(); + Remapper collector = + new Remapper(ASM9) { + @Override + public String map(String internalName) { + if (internalName.startsWith("datadog/")) { + unrelocatedTypes.add(internalName); + } + return internalName; + } + }; + for (byte[] bytecode : relocatedBytecode()) { + new ClassReader(bytecode).accept(new ClassRemapper(new ClassWriter(0), collector), 0); + } + assertThat(unrelocatedTypes).isEmpty(); + } + + /** Unpacks the relocated bytecode from the build-time generated {@code ObjectStoreGlue}. */ + private static List relocatedBytecode() throws IllegalAccessException { + // ObjectStoreGlue is generated at build-time and already compiled onto the test classpath + List relocatedBytecode = new ArrayList<>(); + for (Field field : ObjectStoreGlue.class.getFields()) { + if (field.getType() == String.class && !field.getName().equals("PREFIX")) { + relocatedBytecode.add(Glue.unpackBytecode((String) field.get(null))); + } + } + return relocatedBytecode; + } + + /** Performs structural checks without resolving referenced types */ + private static void verifyStructure(byte[] bytecode) { + new ClassReader(bytecode).accept(new CheckClassAdapter(new ClassWriter(0), true), 0); + } +}