feat: Add support for Thrift 0.9.1 and later versions

instrumentation/thrift/thrift-0.9.1/javaagent/build.gradle.kts

@@ -0,0 +1,51 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("org.apache.thrift")
+    module.set("libthrift")
+    versions.set("[0.9.1,)")
+  }
+}
+val thriftExecutable = "./src/test/resources/thrift"
+val thriftInputFile = "$projectDir/src/test/resources/ThriftService.thrift"
+val thriftOutputDir = "$projectDir/src/test/java"
+
+var generateThrift = tasks.register<Exec>("generateThrift") {
+  group = "build"
+  description = "Generate Java code from Thrift IDL files"
+  commandLine(thriftExecutable, "--gen", "java", "-out", thriftOutputDir, thriftInputFile)
+}
+
+tasks.named<JavaCompile>("compileTestJava") {
+  dependsOn(generateThrift)
+
+  doFirst {
+    source.forEach { file ->
+      if (file.absolutePath.contains("$thriftOutputDir/io/opentelemetry/javaagent/instrumentation/thrift/v0_9_1/thrift/ThriftService.java")) {
+        options.compilerArgs.add("-nowarn")
+        options.compilerArgs.add("-Xlint:-unchecked")
+      }
+    }
+  }
+}
+
+tasks.named<Checkstyle>("checkstyleTest") {
+  exclude("**/thrift/ThriftService.java")
+}
+
+spotless {
+  java {
+    targetExclude("**/thrift/ThriftService.java")
+  }
+}
+
+dependencies {
+  compileOnly("org.apache.thrift:libthrift:0.9.1")
+  implementation(project(":instrumentation:thrift:thrift-common:library"))
+
+  testImplementation("org.apache.thrift:libthrift:0.9.1")
+  testImplementation("javax.annotation:javax.annotation-api:1.3.2")
+}

instrumentation/thrift/thrift-0.9.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/thrift/v0_9_1/AbstractProtocolWrapper.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.thrift.v0_9_1;
+
+import org.apache.thrift.protocol.TProtocol;
+import org.apache.thrift.protocol.TProtocolDecorator;
+
+/**
+ * Note that the 8888th field of record is reserved for transporting trace header. Because Thrift
+ * doesn't support to transport metadata.
+ */
+public abstract class AbstractProtocolWrapper extends TProtocolDecorator {
+  public static final String OT_MAGIC_FIELD = "OT_MAGIC_FIELD";
+  public static final short OT_MAGIC_FIELD_ID = 8888;
+
+  public AbstractProtocolWrapper(TProtocol protocol) {
+    super(protocol);
+  }
+}

instrumentation/thrift/thrift-0.9.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/thrift/v0_9_1/AsyncMethodCallbackWrapper.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.thrift.v0_9_1;
+
+import static io.opentelemetry.javaagent.instrumentation.thrift.v0_9_1.ThriftSingletons.clientInstrumenter;
+import static io.opentelemetry.javaagent.instrumentation.thrift.v0_9_1.ThriftSingletons.serverInstrumenter;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.thrift.common.RequestScopeContext;
+import io.opentelemetry.instrumentation.thrift.common.ThriftRequest;
+import org.apache.thrift.async.AsyncMethodCallback;
+
+public final class AsyncMethodCallbackWrapper<T> implements AsyncMethodCallback<T> {
+  private final AsyncMethodCallback<T> delegate;
+  private RequestScopeContext requestScopeContext;
+  private final boolean isServer;
+
+  public AsyncMethodCallbackWrapper(AsyncMethodCallback<T> methodCallback, boolean isServer) {
+    this.delegate = methodCallback;
+    this.isServer = isServer;
+  }
+
+  public void setRequestScopeContext(RequestScopeContext requestScopeContext) {
+    this.requestScopeContext = requestScopeContext;
+  }
+
+  @Override
+  public void onComplete(T t) {
+    try {
+      if (this.requestScopeContext == null) {
+        return;
+      }
+      this.requestScopeContext.close();
+      Context context = this.requestScopeContext.getContext();
+      ThriftRequest request = this.requestScopeContext.getRequest();
+      if (isServer) {
+        serverInstrumenter().end(context, request, 0, null);
+      } else {
+        clientInstrumenter().end(context, request, 0, null);
+      }
+    } finally {
+      this.delegate.onComplete(t);
+    }
+  }
+
+  @Override
+  public void onError(Exception e) {
+    try {
+      if (this.requestScopeContext == null) {
+        return;
+      }
+      this.requestScopeContext.close();
+      Context context = this.requestScopeContext.getContext();
+      ThriftRequest request = this.requestScopeContext.getRequest();
+      if (isServer) {
+        serverInstrumenter().end(context, request, 1, e);
+      } else {
+        clientInstrumenter().end(context, request, 1, e);
+      }
+    } finally {
+      this.delegate.onError(e);
+    }
+  }
+}

instrumentation/thrift/thrift-0.9.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/thrift/v0_9_1/ThriftSingletons.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.thrift.v0_9_1;
+
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.thrift.common.ThriftInstrumenterFactory;
+import io.opentelemetry.instrumentation.thrift.common.ThriftRequest;
+
+public final class ThriftSingletons {
+  private static final String INSTRUMENTATION_NAME = "io.opentelemetry.thrift-0.9.1";
+
+  private static final Instrumenter<ThriftRequest, Integer> CLIENT_INSTRUMENTER =
+      ThriftInstrumenterFactory.clientInstrumenter(INSTRUMENTATION_NAME);
+  private static final Instrumenter<ThriftRequest, Integer> SERVER_INSTRUMENTER =
+      ThriftInstrumenterFactory.serverInstrumenter(INSTRUMENTATION_NAME);
+
+  public static Instrumenter<ThriftRequest, Integer> clientInstrumenter() {
+    return CLIENT_INSTRUMENTER;
+  }
+
+  public static Instrumenter<ThriftRequest, Integer> serverInstrumenter() {
+    return SERVER_INSTRUMENTER;
+  }
+
+  private ThriftSingletons() {}
+}

instrumentation/thrift/thrift-0.9.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/thrift/v0_9_1/client/ClientOutProtocolWrapper.java

@@ -0,0 +1,143 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.thrift.v0_9_1.client;
+
+import static io.opentelemetry.javaagent.instrumentation.thrift.v0_9_1.ThriftSingletons.clientInstrumenter;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.thrift.common.RequestScopeContext;
+import io.opentelemetry.instrumentation.thrift.common.SocketAccessor;
+import io.opentelemetry.instrumentation.thrift.common.ThriftRequest;
+import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.instrumentation.thrift.v0_9_1.AbstractProtocolWrapper;
+import java.net.Socket;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import org.apache.thrift.TException;
+import org.apache.thrift.protocol.TField;
+import org.apache.thrift.protocol.TMap;
+import org.apache.thrift.protocol.TMessage;
+import org.apache.thrift.protocol.TMessageType;
+import org.apache.thrift.protocol.TProtocol;
+import org.apache.thrift.protocol.TType;
+import org.apache.thrift.transport.TTransport;
+
+@SuppressWarnings("all")
+public final class ClientOutProtocolWrapper extends AbstractProtocolWrapper {
+  public static final String ONE_WAY_METHOD_NAME_PREFIX = "recv_";
+  private volatile RequestScopeContext requestScopeContext;
+  public TTransport transport;
+  private boolean injected = true;
+  private String methodName;
+  private final Set<String> voidMethodNames;
+  private String serviceName;
+  private byte type = -1;
+  private byte originType;
+
+  public ClientOutProtocolWrapper(
+      TProtocol protocol, String serviceName, Set<String> voidMethodNames) {
+    super(protocol);
+    this.serviceName = serviceName;
+    this.voidMethodNames = voidMethodNames;
+  }
+
+  @Override
+  public void writeMessageBegin(TMessage message) throws TException {
+    this.injected = false;
+    this.methodName = message.name;
+    this.originType = message.type;
+    // Compatible with version 0.9.1 and 0.9.2 asynchronous logic
+    if (message.type == TMessageType.ONEWAY || this.type == -1) {
+      this.type = message.type;
+    }
+    if (!this.isOneway()) {
+      if (this.voidMethodNames != null
+          && this.voidMethodNames.contains(this.methodName)
+          && !this.voidMethodNames.contains(ONE_WAY_METHOD_NAME_PREFIX + this.methodName)) {
+        this.type = TMessageType.ONEWAY;
+      }
+    }
+    try {
+      if (this.requestScopeContext == null) {
+        Socket socket = SocketAccessor.getSocket(super.getTransport());
+        if (socket == null) {
+          socket = SocketAccessor.getSocket(this.transport);
+        }
+        ThriftRequest request =
+            ThriftRequest.create(this.serviceName, this.methodName, socket, new HashMap<>());
+        Context parentContext = Java8BytecodeBridge.currentContext();
+        if (!clientInstrumenter().shouldStart(parentContext, request)) {
+          return;
+        }
+        Context context = clientInstrumenter().start(parentContext, request);
+        this.requestScopeContext = RequestScopeContext.create(request, null, context);
+      }
+    } finally {
+      if (this.isOneway() && message.type != TMessageType.ONEWAY) {
+        TMessage onewayMessage = new TMessage(message.name, TMessageType.ONEWAY, message.seqid);
+        super.writeMessageBegin(onewayMessage);
+      } else {
+        super.writeMessageBegin(message);
+      }
+    }
+  }
+
+  @Override
+  public void writeFieldStop() throws TException {
+    try {
+      if (!this.injected && this.requestScopeContext != null) {
+        ThriftRequest request = this.requestScopeContext.getRequest();
+        this.writeHeader(request.getHeader());
+      }
+    } finally {
+      this.injected = true;
+      super.writeFieldStop();
+    }
+  }
+
+  public void writeHeader(Map<String, String> header) throws TException {
+    super.writeFieldBegin(new TField(OT_MAGIC_FIELD, TType.MAP, OT_MAGIC_FIELD_ID));
+    super.writeMapBegin(new TMap(TType.STRING, TType.STRING, header.size()));
+
+    Set<Map.Entry<String, String>> entries = header.entrySet();
+    for (Map.Entry<String, String> entry : entries) {
+      super.writeString(entry.getKey());
+      super.writeString(entry.getValue());
+    }
+
+    super.writeMapEnd();
+    super.writeFieldEnd();
+  }
+
+  public boolean isOneway() {
+    return this.type == TMessageType.ONEWAY;
+  }
+
+  public boolean isChangeToOneway() {
+    return this.type != this.originType;
+  }
+
+  public void updateTransport(TTransport transport) {
+    this.transport = transport;
+  }
+
+  public RequestScopeContext getRequestScopeContext() {
+    return requestScopeContext;
+  }
+
+  public void setRequestScopeContext(RequestScopeContext requestScopeContext) {
+    this.requestScopeContext = requestScopeContext;
+  }
+
+  public void setServiceName(String serviceName) {
+    this.serviceName = serviceName;
+  }
+
+  public void setType(byte type) {
+    this.type = type;
+  }
+}

instrumentation/thrift/thrift-0.9.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/thrift/v0_9_1/client/ClientProtocolFactoryWrapper.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.thrift.v0_9_1.client;
+
+import org.apache.thrift.protocol.TProtocol;
+import org.apache.thrift.protocol.TProtocolFactory;
+import org.apache.thrift.transport.TTransport;
+
+@SuppressWarnings({"serial"})
+public final class ClientProtocolFactoryWrapper implements TProtocolFactory {
+  public TProtocolFactory delegate;
+  public TTransport transport;
+  public String serviceName;
+
+  @Override
+  public TProtocol getProtocol(TTransport transport) {
+    TProtocol protocol = this.delegate.getProtocol(transport);
+    if (protocol instanceof ClientOutProtocolWrapper) {
+      if (transport != null) {
+        ((ClientOutProtocolWrapper) protocol).updateTransport(this.transport);
+      }
+      ((ClientOutProtocolWrapper) protocol).setServiceName(this.serviceName);
+      return protocol;
+    }
+    protocol = new ClientOutProtocolWrapper(protocol, this.serviceName, null);
+    if (transport != null) {
+      ((ClientOutProtocolWrapper) protocol).updateTransport(this.transport);
+    }
+    return protocol;
+  }
+
+  public ClientProtocolFactoryWrapper(
+      TProtocolFactory protocolFactory, TTransport transport, String serviceName) {
+    this.delegate = protocolFactory;
+    this.transport = transport;
+    this.serviceName = serviceName;
+  }
+}