diff --git a/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelInstrumentationMapper.java b/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelInstrumentationMapper.java index eb48454ba6d..80b84df839f 100644 --- a/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelInstrumentationMapper.java +++ b/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelInstrumentationMapper.java @@ -3,15 +3,17 @@ import static datadog.trace.agent.tooling.ExtensionHandler.MAP_LOGGING; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.bootstrap.ContextStore; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; import net.bytebuddy.jar.asm.ClassVisitor; import net.bytebuddy.jar.asm.MethodVisitor; +import net.bytebuddy.jar.asm.Type; import net.bytebuddy.jar.asm.commons.ClassRemapper; import net.bytebuddy.jar.asm.commons.Remapper; @@ -19,16 +21,21 @@ public final class OtelInstrumentationMapper extends ClassRemapper { private static final Set UNSUPPORTED_TYPES = - new HashSet<>( - Arrays.asList("io/opentelemetry/javaagent/tooling/muzzle/InstrumentationModuleMuzzle")); + Collections.singleton( + "io/opentelemetry/javaagent/tooling/muzzle/InstrumentationModuleMuzzle"); private static final Set UNSUPPORTED_METHODS = - new HashSet<>(Arrays.asList("getMuzzleReferences", "registerMuzzleVirtualFields")); + Collections.singleton("getMuzzleReferences"); public OtelInstrumentationMapper(ClassVisitor classVisitor) { super(classVisitor, Renamer.INSTANCE); } + @Override + protected MethodVisitor createMethodRemapper(MethodVisitor methodVisitor) { + return new OtelMethodCallMapper(methodVisitor, remapper); + } + @Override public void visit( int version, @@ -74,16 +81,22 @@ static final class Renamer extends Remapper { static { RENAMED_TYPES.put( "io/opentelemetry/javaagent/extension/instrumentation/InstrumentationModule", - "datadog/opentelemetry/tooling/OtelInstrumenterModule"); + Type.getInternalName(OtelInstrumenterModule.class)); RENAMED_TYPES.put( "io/opentelemetry/javaagent/extension/instrumentation/TypeInstrumentation", - "datadog/opentelemetry/tooling/OtelInstrumenter"); + Type.getInternalName(OtelInstrumenter.class)); RENAMED_TYPES.put( "io/opentelemetry/javaagent/extension/instrumentation/TypeTransformer", - "datadog/opentelemetry/tooling/OtelTransformer"); + Type.getInternalName(OtelTransformer.class)); RENAMED_TYPES.put( "io/opentelemetry/javaagent/extension/matcher/AgentElementMatchers", - "datadog/opentelemetry/tooling/OtelElementMatchers"); + Type.getInternalName(OtelElementMatchers.class)); + RENAMED_TYPES.put( + "io/opentelemetry/javaagent/tooling/muzzle/VirtualFieldMappingsBuilder", + Type.getInternalName(OtelInstrumenterModule.VirtualFieldBuilder.class)); + RENAMED_TYPES.put( + "io/opentelemetry/javaagent/shaded/instrumentation/api/util/VirtualField", + Type.getInternalName(ContextStore.class)); RENAMED_TYPES.put( "io/opentelemetry/javaagent/bootstrap/Java8BytecodeBridge", "datadog/trace/bootstrap/otel/Java8BytecodeBridge"); diff --git a/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelInstrumenterModule.java b/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelInstrumenterModule.java index 7011106f192..2477518cc22 100644 --- a/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelInstrumenterModule.java +++ b/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelInstrumenterModule.java @@ -4,7 +4,9 @@ import datadog.trace.api.InstrumenterConfig; import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; /** * Replaces OpenTelemetry's {@code InstrumentationModule} when mapping extensions. @@ -62,4 +64,25 @@ public List getMuzzleHelperClassNames() { public List getAdditionalHelperClassNames() { return Collections.emptyList(); } + + @Override + public Map contextStore() { + Map virtualFields = new HashMap<>(); + registerMuzzleVirtualFields( + new VirtualFieldBuilder() { + @Override + public VirtualFieldBuilder register(String typeName, String fieldTypeName) { + virtualFields.put(typeName, fieldTypeName); + return this; + } + }); + return virtualFields; + } + + /** This method is generated by OpenTelemetry's muzzle plugin at build time. */ + public void registerMuzzleVirtualFields(VirtualFieldBuilder builder) {} + + public interface VirtualFieldBuilder { + VirtualFieldBuilder register(String typeName, String fieldTypeName); + } } diff --git a/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelMethodCallMapper.java b/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelMethodCallMapper.java new file mode 100644 index 00000000000..3d8021acf84 --- /dev/null +++ b/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelMethodCallMapper.java @@ -0,0 +1,114 @@ +package datadog.opentelemetry.tooling; + +import static datadog.trace.bootstrap.FieldBackedContextStores.FAST_STORE_ID_LIMIT; +import static datadog.trace.bootstrap.FieldBackedContextStores.getContextStoreId; + +import datadog.trace.bootstrap.FieldBackedContextStore; +import datadog.trace.bootstrap.FieldBackedContextStores; +import net.bytebuddy.jar.asm.MethodVisitor; +import net.bytebuddy.jar.asm.Opcodes; +import net.bytebuddy.jar.asm.Type; +import net.bytebuddy.jar.asm.commons.MethodRemapper; +import net.bytebuddy.jar.asm.commons.Remapper; + +/** Maps OpenTelemetry method calls to use the Datadog equivalent API. */ +public final class OtelMethodCallMapper extends MethodRemapper { + + private static final String VIRTUAL_FIELD_CLASS = + "io/opentelemetry/javaagent/shaded/instrumentation/api/util/VirtualField"; + + private static final String FIELD_BACKED_CONTEXT_STORES_CLASS = + Type.getInternalName(FieldBackedContextStores.class); + + private static final String GET_CONTENT_STORE_METHOD = "getContextStore"; + private static final String GET_CONTENT_STORE_METHOD_DESCRIPTOR = + Type.getMethodDescriptor(Type.getType(FieldBackedContextStore.class), Type.INT_TYPE); + + private static final String FIELD_BACKED_CONTENT_STORE_DESCRIPTOR = + Type.getDescriptor(FieldBackedContextStore.class); + + private static final String FAST_CONTENT_STORE_PREFIX = "contextStore"; + + /** The last two constants pushed onto the stack. */ + private Object constant1, constant2; + + public OtelMethodCallMapper(MethodVisitor methodVisitor, Remapper remapper) { + super(methodVisitor, remapper); + } + + @Override + public void visitLdcInsn(final Object value) { + // 'first' constant gets moved over once 'second' one comes in + constant1 = constant2; + constant2 = value; + super.visitLdcInsn(value); + } + + @Override + public void visitMethodInsn( + final int opcode, + final String owner, + final String name, + final String descriptor, + final boolean isInterface) { + + if (VIRTUAL_FIELD_CLASS.equals(owner)) { + // replace VirtualField calls with their equivalent in the ContextStore API + if ("find".equals(name)) { + redirectVirtualFieldLookup(); + return; + } else if ("set".equals(name)) { + super.visitMethodInsn(Opcodes.INVOKEINTERFACE, owner, "put", descriptor, true); + return; + } else if ("get".equals(name)) { + super.visitMethodInsn(Opcodes.INVOKEINTERFACE, owner, "get", descriptor, true); + return; + } + } + + super.visitMethodInsn(opcode, owner, name, descriptor, isInterface); + + // reset constants for next method check + constant1 = null; + constant2 = null; + } + + private void redirectVirtualFieldLookup() { + // We track the last two constants pushed onto the stack to make sure they match + // the expected key and context types. Matching calls are rewritten to call the + // dynamically injected context store implementation instead. + String keyClassName = null; + String contextClassName = null; + if (constant1 instanceof Type && constant2 instanceof Type) { + keyClassName = ((Type) constant1).getClassName(); + contextClassName = ((Type) constant2).getClassName(); + } + + if (null == keyClassName || null == contextClassName) { + throw new IllegalStateException( + "Incorrect VirtualField usage detected. Type and fieldType must be class-literals. " + + "Example of correct usage: VirtualField.find(Runnable.class, RunnableContext.class)"); + } + + // discard original parameters so we can use numeric id instead + mv.visitInsn(Opcodes.POP2); + + int storeId = getContextStoreId(keyClassName, contextClassName); + // use fast direct field access for a small number of stores + if (storeId < FAST_STORE_ID_LIMIT) { + mv.visitFieldInsn( + Opcodes.GETSTATIC, + FIELD_BACKED_CONTEXT_STORES_CLASS, + FAST_CONTENT_STORE_PREFIX + storeId, + FIELD_BACKED_CONTENT_STORE_DESCRIPTOR); + } else { + mv.visitLdcInsn(storeId); + mv.visitMethodInsn( + Opcodes.INVOKESTATIC, + FIELD_BACKED_CONTEXT_STORES_CLASS, + GET_CONTENT_STORE_METHOD, + GET_CONTENT_STORE_METHOD_DESCRIPTOR, + false); + } + } +}