Skip to content

Commit

Permalink
Map OpenTelemetry VirtualField to Datadog ContextStore (#7129)
Browse files Browse the repository at this point in the history
* Map virtual field usage into context store request
* Use Type.getInternalName to help future refactoring
* Allow advanced mapping of OpenTelemetry method calls
* Map OpenTelemetry VirtualField to Datadog ContextStore
  • Loading branch information
mcculls committed Jun 12, 2024
1 parent 97065ed commit 46e782c
Show file tree
Hide file tree
Showing 3 changed files with 158 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -3,32 +3,39 @@
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;

/** Maps OpenTelemetry instrumentations to use the Datadog {@link InstrumenterModule} API. */
public final class OtelInstrumentationMapper extends ClassRemapper {

private static final Set<String> UNSUPPORTED_TYPES =
new HashSet<>(
Arrays.asList("io/opentelemetry/javaagent/tooling/muzzle/InstrumentationModuleMuzzle"));
Collections.singleton(
"io/opentelemetry/javaagent/tooling/muzzle/InstrumentationModuleMuzzle");

private static final Set<String> 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,
Expand Down Expand Up @@ -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");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -62,4 +64,25 @@ public List<String> getMuzzleHelperClassNames() {
public List<String> getAdditionalHelperClassNames() {
return Collections.emptyList();
}

@Override
public Map<String, String> contextStore() {
Map<String, String> 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);
}
}
Original file line number Diff line number Diff line change
@@ -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);
}
}
}

0 comments on commit 46e782c

Please sign in to comment.