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 bd3e72a15c3..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,6 +3,7 @@ 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; @@ -93,6 +94,9 @@ static final class Renamer extends Remapper { 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/OtelMethodCallMapper.java b/dd-java-agent/agent-otel/otel-tooling/src/main/java/datadog/opentelemetry/tooling/OtelMethodCallMapper.java index 7e993d90305..3d8021acf84 100644 --- 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 @@ -1,12 +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); + } + } }