feat: Add support for Thrift 0.9.1 and later versions

.fossa.yml

@@ -1024,6 +1024,12 @@ targets:
     - type: gradle
       path: ./
       target: ':instrumentation:struts:struts-7.0:javaagent'
+    - type: gradle
+      path: ./
+      target: ':instrumentation:thrift:thrift-0.9.1:javaagent'
+    - type: gradle
+      path: ./
+      target: ':instrumentation:thrift:thrift-common:javaagent'
     - type: gradle
       path: ./
       target: ':instrumentation:tomcat:tomcat-10.0:javaagent'

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

@@ -0,0 +1,60 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("org.apache.thrift")
+    module.set("libthrift")
+    versions.set("[0.9.1,)")
+    assertInverse.set(true)
+  }
+}
+
+dependencies {
+  library("org.apache.thrift:libthrift:0.9.1")
+  testLibrary("org.apache.thrift:libthrift:0.12.0")
+  implementation(project(":instrumentation:thrift:thrift-common:javaagent"))
+}
+
+val generatedThriftDir = layout.buildDirectory.dir("generated/thrift/test")
+
+val generateThriftSources by tasks.registering(Exec::class) {
+  val thriftFilePath =
+    layout.projectDirectory.file("src/test/resources/ThriftService.thrift").asFile.absolutePath
+  val outputDirPath = generatedThriftDir.get().asFile.also { it.mkdirs() }.absolutePath
+  inputs.file(thriftFilePath)
+  outputs.dir(outputDirPath)
+
+  standardOutput = System.out
+  executable = "docker"
+  args = listOf(
+    "run",
+    "--rm",
+    "--platform=linux/amd64",
+    "-v", "$thriftFilePath:/thrift/input/ThriftService.thrift:ro",
+    "-v", "$outputDirPath:/thrift/output",
+    "thrift:0.12.0",
+    "thrift",
+    "--gen",
+    "java",
+    "-out",
+    "/thrift/output",
+    "/thrift/input/ThriftService.thrift")
+}
+
+sourceSets {
+  test {
+    java {
+      srcDir(generatedThriftDir)
+    }
+  }
+}
+
+tasks.compileTestJava {
+  dependsOn(generateThriftSources)
+}
+
+tasks.named<Checkstyle>("checkstyleTest") {
+  exclude("**/thrift/**")
+}

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Includes work from:
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+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,148 @@
+/*
+ * 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.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;
+
+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 = Context.current();
+        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) {
+        // In Thrift 0.9.1 and 0.9.2 versions, the type of the TMessage for oneway requests is still
+        // TMessageType.CALL.
+        // This causes issues with the server-side instrumentation logic. Here, we are simply
+        // correcting the actual request type.
+        // Since it is a oneway request, the client does not need to handle the response,
+        // and the server does not use this type for any specific logic processing.
+        // Therefore, it has no impact on either the client or the server.
+        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;
+  }
+}