Thrift instrumentation simplifications

instrumentation/thrift-0.13/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/thrift/v0_13/ThriftAsyncClientInstrumentation.java

@@ -70,14 +70,9 @@ public static TProtocolFactory onEnter(
         @Advice.Origin("#t") Class<?> declaringClass,
         @Advice.Argument(0) TProtocolFactory protocolFactory,
         @Advice.Argument(2) TNonblockingTransport transport) {
-      Class<?> serviceClass = declaringClass;
-      if (serviceClass.getDeclaringClass() != null) {
-        serviceClass = serviceClass.getDeclaringClass();
-      }
-
       return new ClientProtocolDecorator.Factory(
           protocolFactory,
-          serviceClass.getName(),
+          ThriftSingletons.thriftServiceName(declaringClass),
           ThriftSingletons.clientInstrumenter(),
           getPropagators(),
           transport);
@@ -96,7 +91,6 @@ public static Object[] onEnter(@Advice.AllArguments Object[] arguments) {
       if (args[args.length - 1] != null) {
         AsyncMethodCallback<?> callback = (AsyncMethodCallback<?>) args[args.length - 1];
         args[args.length - 1] = AsyncMethodCallbackUtil.wrap(callback, clientContext);
-        clientContext.setHasAsyncCallback();
       }
       return new Object[] {args, clientContext};
     }

instrumentation/thrift-0.13/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/thrift/v0_13/ThriftServiceClientInstrumentation.java

@@ -48,13 +48,11 @@ public static Object[] onExit(
         @Advice.This TServiceClient serviceClient,
         @Advice.FieldValue("iprot_") TProtocol inProtocol,
         @Advice.FieldValue("oprot_") TProtocol outProtocol) {
-      Class<?> serviceClass = serviceClient.getClass();
-      if (serviceClass.getDeclaringClass() != null) {
-        serviceClass = serviceClass.getDeclaringClass();
-      }
       ClientProtocolDecorator.AgentDecorator agentDecorator =
           new ClientProtocolDecorator.AgentDecorator(
-              serviceClass.getName(), ThriftSingletons.clientInstrumenter(), getPropagators());
+              ThriftSingletons.thriftServiceName(serviceClient.getClass()),
+              ThriftSingletons.clientInstrumenter(),
+              getPropagators());
       return new Object[] {
         agentDecorator.decorate(inProtocol), agentDecorator.decorate(outProtocol)
       };

instrumentation/thrift-0.13/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/thrift/v0_13/ThriftSingletons.java

@@ -36,5 +36,15 @@ public static ContextPropagators getPropagators() {
     return GlobalOpenTelemetry.get().getPropagators();
   }
 
+  /**
+   * Returns the enclosing thrift service class name (e.g. {@code com.example.MyService}) for an
+   * inner client class such as {@code com.example.MyService$Client}; falls back to the class itself
+   * when there is no enclosing class.
+   */
+  public static String thriftServiceName(Class<?> clientClass) {
+    Class<?> declaring = clientClass.getDeclaringClass();
+    return declaring != null ? declaring.getName() : clientClass.getName();
+  }
+
   private ThriftSingletons() {}
 }

instrumentation/thrift-0.13/library/src/main/java/io/opentelemetry/instrumentation/thrift/v0_13/ThriftTelemetry.java

@@ -89,12 +89,11 @@ public <T> T wrapAsyncClient(TAsyncClient client, Class<T> asyncInterface) {
                 if (args.length > 0 && args[args.length - 1] instanceof AsyncMethodCallback) {
                   AsyncMethodCallback<?> callback = (AsyncMethodCallback<?>) args[args.length - 1];
                   args[args.length - 1] = AsyncMethodCallbackUtil.wrap(callback, clientContext);
-                  clientContext.setHasAsyncCallback();
                 }
                 return method.invoke(client, args);
               } catch (InvocationTargetException e) {
                 error = e.getCause();
-                throw e.getCause();
+                throw error;
               } catch (Throwable t) {
                 error = t;
                 throw t;

instrumentation/thrift-0.13/library/src/main/java/io/opentelemetry/instrumentation/thrift/v0_13/internal/AsyncMethodCallbackUtil.java

@@ -7,7 +7,6 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
-import javax.annotation.Nullable;
 import org.apache.thrift.async.AsyncMethodCallback;
 
 /**
@@ -16,13 +15,8 @@
  */
 public final class AsyncMethodCallbackUtil {
 
-  @Nullable
   public static <T> AsyncMethodCallback<T> wrap(
-      @Nullable AsyncMethodCallback<T> callback, ClientCallContext clientCallContext) {
-    if (callback == null) {
-      return null;
-    }
-
+      AsyncMethodCallback<T> callback, ClientCallContext clientCallContext) {
     Context context = Context.current();
     return new AsyncMethodCallback<T>() {
 

instrumentation/thrift-0.13/library/src/main/java/io/opentelemetry/instrumentation/thrift/v0_13/internal/ClientCallContext.java

@@ -15,7 +15,6 @@ public final class ClientCallContext {
   private static final ThreadLocal<ClientCallContext> current = new ThreadLocal<>();
 
   @Nullable private ClientProtocolDecorator protocolDecorator;
-  private boolean hasAsyncCallback;
 
   public static ClientCallContext start() {
     ClientCallContext context = new ClientCallContext();
@@ -30,17 +29,6 @@ static void setClientProtocolDecorator(ClientProtocolDecorator protocolDecorator
     }
   }
 
-  /** Returns {@code true} when a callback has been installed for ending the span. */
-  static boolean hasAsyncCallback() {
-    ClientCallContext context = current.get();
-    return context != null && context.hasAsyncCallback;
-  }
-
-  /** Notify that a callback has been installed for ending the span. */
-  public void setHasAsyncCallback() {
-    this.hasAsyncCallback = true;
-  }
-
   public void endSpan(@Nullable Throwable throwable) {
     if (protocolDecorator != null) {
       protocolDecorator.endSpan(throwable);

instrumentation/thrift-0.13/library/src/main/java/io/opentelemetry/instrumentation/thrift/v0_13/internal/ClientProtocolDecorator.java

@@ -42,7 +42,7 @@ public ClientProtocolDecorator(
       String serviceName,
       Instrumenter<ThriftRequest, ThriftResponse> instrumenter,
       ContextPropagators propagators) {
-    this(protocol, serviceName, instrumenter, propagators, null, new State());
+    this(protocol, serviceName, instrumenter, propagators, null, new State(), true);
   }
 
   private ClientProtocolDecorator(
@@ -51,17 +51,16 @@ private ClientProtocolDecorator(
       Instrumenter<ThriftRequest, ThriftResponse> instrumenter,
       ContextPropagators propagators,
       @Nullable TTransport transport,
-      State state) {
+      State state,
+      boolean endSpanInline) {
     super(protocol);
     this.protocol = protocol;
     this.serviceName = serviceName;
     this.instrumenter = instrumenter;
     this.propagators = propagators;
     this.transport = transport;
     this.state = state;
-    // if there's an async callback, the span will be ended in the callback
-    this.endSpan = !ClientCallContext.hasAsyncCallback();
-    ClientCallContext.setClientProtocolDecorator(this);
+    this.endSpan = endSpanInline;
   }
 
   @Override
@@ -226,13 +225,18 @@ public Factory(
 
     @Override
     public TProtocol getProtocol(TTransport transport) {
-      return new ClientProtocolDecorator(
-          protocolFactory.getProtocol(transport),
-          serviceName,
-          instrumenter,
-          propagators,
-          configuredTransport,
-          state);
+      ClientProtocolDecorator decorator =
+          new ClientProtocolDecorator(
+              protocolFactory.getProtocol(transport),
+              serviceName,
+              instrumenter,
+              propagators,
+              configuredTransport,
+              state,
+              false);
+      // span will be ended by the wrapped async callback, which calls back into this decorator
+      ClientCallContext.setClientProtocolDecorator(decorator);
+      return decorator;
     }
   }
 
@@ -260,7 +264,7 @@ public TProtocol decorate(TProtocol protocol) {
         return protocol;
       }
       return new ClientProtocolDecorator(
-          protocol, serviceName, instrumenter, propagators, null, state);
+          protocol, serviceName, instrumenter, propagators, null, state, true);
     }
   }
 }

instrumentation/thrift-0.13/library/src/main/java/io/opentelemetry/instrumentation/thrift/v0_13/internal/ThriftRpcAttributesGetter.java

@@ -48,7 +48,7 @@ public String getErrorType(
       ThriftRequest request, @Nullable ThriftResponse response, @Nullable Throwable error) {
     if (error instanceof TTransportException) {
       int errorCode = ((TTransportException) error).getType();
-      switch (((TTransportException) error).getType()) {
+      switch (errorCode) {
         case 0:
           return "UNKNOWN";
         case 1: