Add Sofa RPC instrumentation (Bolt, H2C, REST, Triple)

dd-java-agent/build.gradle

@@ -467,7 +467,7 @@ tasks.register('checkAgentJarSize') {
   doLast {
     // Arbitrary limit to prevent unintentional increases to the agent jar size
     // Raise or lower as required
-    assert tasks.named("shadowJar", ShadowJar).get().archiveFile.get().getAsFile().length() <= 32 * 1024 * 1024
+    assert tasks.named("shadowJar", ShadowJar).get().archiveFile.get().getAsFile().length() <= 33 * 1024 * 1024
   }
 
   dependsOn "shadowJar"

dd-java-agent/instrumentation/sofarpc/build.gradle

@@ -0,0 +1 @@
+apply from: "$rootDir/gradle/java.gradle"

dd-java-agent/instrumentation/sofarpc/sofarpc-5.0/build.gradle

@@ -0,0 +1,35 @@
+muzzle {
+  pass {
+    group = "com.alipay.sofa"
+    module = "sofa-rpc-all"
+    versions = "[5.0.0,)"
+    assertInverse = true
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+addTestSuiteForDir('latestDepTest', 'test')
+
+configurations.testRuntimeClasspath {
+  resolutionStrategy.force "com.google.guava:guava:32.1.3-jre"
+}
+
+dependencies {
+  compileOnly group: "com.alipay.sofa", name: "sofa-rpc-all", version: "5.6.0"
+
+  testImplementation group: "com.alipay.sofa", name: "sofa-rpc-all", version: "5.14.2"
+  // JAX-RS annotations required for the REST protocol test interface (@Path, @GET, etc.)
+  testImplementation group: "javax.ws.rs", name: "javax.ws.rs-api", version: "2.1.1"
+  // Required so that GrpcServerModule / GrpcClientModule are discovered via ServiceLoader
+  // in SofaRpcTripleWithGrpcForkedTest.
+  testImplementation project(':dd-java-agent:instrumentation:grpc-1.5')
+  testImplementation group: "io.grpc", name: "grpc-netty", version: "1.53.0"
+  testImplementation group: "io.grpc", name: "grpc-core", version: "1.53.0"
+  testImplementation group: "io.grpc", name: "grpc-stub", version: "1.53.0"
+  testImplementation group: "com.google.protobuf", name: "protobuf-java", version: "3.25.3"
+  testImplementation group: "com.alibaba", name: "fastjson", version: "1.2.83"
+  testImplementation group: "com.google.guava", name: "guava", version: "32.1.3-jre"
+
+  latestDepTestImplementation group: "com.alipay.sofa", name: "sofa-rpc-all", version: "+"
+}

dd-java-agent/instrumentation/sofarpc/sofarpc-5.0/src/main/java/datadog/trace/instrumentation/sofarpc/AbstractClusterInstrumentation.java

@@ -0,0 +1,74 @@
+package datadog.trace.instrumentation.sofarpc;
+
+import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.sofarpc.SofaRpcClientDecorator.DECORATE;
+import static datadog.trace.instrumentation.sofarpc.SofaRpcClientDecorator.SOFA_RPC_CLIENT;
+import static datadog.trace.instrumentation.sofarpc.SofaRpcInjectAdapter.SETTER;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.alipay.sofa.rpc.client.AbstractCluster;
+import com.alipay.sofa.rpc.config.ConsumerConfig;
+import com.alipay.sofa.rpc.core.request.SofaRequest;
+import com.alipay.sofa.rpc.core.response.SofaResponse;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import net.bytebuddy.asm.Advice;
+
+public class AbstractClusterInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  @Override
+  public String instrumentedType() {
+    return "com.alipay.sofa.rpc.client.AbstractCluster";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("invoke"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("com.alipay.sofa.rpc.core.request.SofaRequest"))),
+        getClass().getName() + "$InvokeAdvice");
+  }
+
+  public static class InvokeAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope enter(
+        @Advice.This AbstractCluster self, @Advice.Argument(0) SofaRequest request) {
+      ConsumerConfig config = self.getConsumerConfig();
+      String protocol = config != null ? config.getProtocol() : null;
+      AgentSpan span = startSpan(SOFA_RPC_CLIENT);
+      DECORATE.afterStart(span);
+      DECORATE.onRequest(span, request);
+      if (protocol != null) {
+        span.setTag("sofarpc.protocol", protocol);
+      }
+      AgentScope scope = activateSpan(span);
+      defaultPropagator().inject(span, request, SETTER);
+      return scope;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(
+        @Advice.Enter AgentScope scope,
+        @Advice.Return SofaResponse response,
+        @Advice.Thrown Throwable throwable) {
+      if (scope == null) {
+        return;
+      }
+      AgentSpan span = scope.span();
+      DECORATE.onResponse(span, response);
+      DECORATE.onError(span, throwable);
+      DECORATE.beforeFinish(span);
+      scope.close();
+      span.finish();
+    }
+  }
+}

dd-java-agent/instrumentation/sofarpc/sofarpc-5.0/src/main/java/datadog/trace/instrumentation/sofarpc/BoltServerProcessorInstrumentation.java

@@ -0,0 +1,40 @@
+package datadog.trace.instrumentation.sofarpc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import datadog.trace.agent.tooling.Instrumenter;
+import net.bytebuddy.asm.Advice;
+
+public class BoltServerProcessorInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  @Override
+  public String instrumentedType() {
+    return "com.alipay.sofa.rpc.server.bolt.BoltServerProcessor";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("handleRequest"))
+            .and(takesArguments(3))
+            .and(takesArgument(2, named("com.alipay.sofa.rpc.core.request.SofaRequest"))),
+        getClass().getName() + "$HandleRequestAdvice");
+  }
+
+  public static class HandleRequestAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void enter() {
+      SofaRpcProtocolContext.set("bolt");
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit() {
+      SofaRpcProtocolContext.clear();
+    }
+  }
+}

dd-java-agent/instrumentation/sofarpc/sofarpc-5.0/src/main/java/datadog/trace/instrumentation/sofarpc/H2cServerTaskInstrumentation.java

@@ -0,0 +1,35 @@
+package datadog.trace.instrumentation.sofarpc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
+
+import datadog.trace.agent.tooling.Instrumenter;
+import net.bytebuddy.asm.Advice;
+
+public class H2cServerTaskInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  @Override
+  public String instrumentedType() {
+    return "com.alipay.sofa.rpc.server.http.AbstractHttpServerTask";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(named("run")).and(takesNoArguments()), getClass().getName() + "$RunAdvice");
+  }
+
+  public static class RunAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void enter() {
+      SofaRpcProtocolContext.set("h2c");
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit() {
+      SofaRpcProtocolContext.clear();
+    }
+  }
+}

dd-java-agent/instrumentation/sofarpc/sofarpc-5.0/src/main/java/datadog/trace/instrumentation/sofarpc/ProviderProxyInvokerInstrumentation.java

@@ -0,0 +1,80 @@
+package datadog.trace.instrumentation.sofarpc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.sofarpc.SofaRpcExtractAdapter.GETTER;
+import static datadog.trace.instrumentation.sofarpc.SofaRpcServerDecorator.DECORATE;
+import static datadog.trace.instrumentation.sofarpc.SofaRpcServerDecorator.SOFA_RPC_SERVER;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.alipay.sofa.rpc.core.request.SofaRequest;
+import com.alipay.sofa.rpc.core.response.SofaResponse;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import net.bytebuddy.asm.Advice;
+
+public class ProviderProxyInvokerInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  @Override
+  public String instrumentedType() {
+    return "com.alipay.sofa.rpc.server.ProviderProxyInvoker";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("invoke"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("com.alipay.sofa.rpc.core.request.SofaRequest"))),
+        getClass().getName() + "$InvokeAdvice");
+  }
+
+  public static class InvokeAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope enter(@Advice.Argument(0) SofaRequest request) {
+      // Protocol is set in thread-local by transport-specific instrumentation before this call.
+      // If null, the transport is not instrumented — skip.
+      String protocol = SofaRpcProtocolContext.get();
+      if (protocol == null) {
+        return null;
+      }
+      // For Bolt and H2C the client injects trace context into SofaRequest.requestProps;
+      // extract it here. For Triple, parentContext will be null and startSpan() without an
+      // explicit parent naturally attaches to the active grpc.server span. For REST,
+      // parentContext will also be null and the active netty.request span becomes the parent.
+      AgentSpanContext parentContext = extractContextAndGetSpanContext(request, GETTER);
+      AgentSpan span =
+          parentContext != null
+              ? startSpan(SOFA_RPC_SERVER, parentContext)
+              : startSpan(SOFA_RPC_SERVER);
+      DECORATE.afterStart(span);
+      DECORATE.onRequest(span, request);
+      span.setTag("sofarpc.protocol", protocol);
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(
+        @Advice.Enter AgentScope scope,
+        @Advice.Return SofaResponse response,
+        @Advice.Thrown Throwable throwable) {
+      if (scope == null) {
+        return;
+      }
+      AgentSpan span = scope.span();
+      DECORATE.onResponse(span, response);
+      DECORATE.onError(span, throwable);
+      DECORATE.beforeFinish(span);
+      scope.close();
+      span.finish();
+    }
+  }
+}

dd-java-agent/instrumentation/sofarpc/sofarpc-5.0/src/main/java/datadog/trace/instrumentation/sofarpc/RestServerHandlerInstrumentation.java

@@ -0,0 +1,36 @@
+package datadog.trace.instrumentation.sofarpc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import datadog.trace.agent.tooling.Instrumenter;
+import net.bytebuddy.asm.Advice;
+
+public class RestServerHandlerInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  @Override
+  public String instrumentedType() {
+    return "com.alipay.sofa.rpc.server.rest.SofaRestRequestHandler";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(named("channelRead0")).and(takesArguments(2)),
+        getClass().getName() + "$ChannelRead0Advice");
+  }
+
+  public static class ChannelRead0Advice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void enter() {
+      SofaRpcProtocolContext.set("rest");
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit() {
+      SofaRpcProtocolContext.clear();
+    }
+  }
+}

dd-java-agent/instrumentation/sofarpc/sofarpc-5.0/src/main/java/datadog/trace/instrumentation/sofarpc/SofaRpcClientDecorator.java

@@ -0,0 +1,67 @@
+package datadog.trace.instrumentation.sofarpc;
+
+import com.alipay.sofa.rpc.core.request.SofaRequest;
+import com.alipay.sofa.rpc.core.response.SofaResponse;
+import datadog.trace.api.naming.SpanNaming;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.bootstrap.instrumentation.decorator.ClientDecorator;
+
+public class SofaRpcClientDecorator extends ClientDecorator {
+
+  public static final CharSequence SOFA_RPC_CLIENT =
+      UTF8BytesString.create(
+          SpanNaming.instance().namingSchema().client().operationForProtocol("sofarpc"));
+
+  private static final CharSequence COMPONENT_NAME = UTF8BytesString.create("sofarpc-client");
+
+  public static final SofaRpcClientDecorator DECORATE = new SofaRpcClientDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"sofarpc"};
+  }
+
+  @Override
+  protected CharSequence component() {
+    return COMPONENT_NAME;
+  }
+
+  @Override
+  protected CharSequence spanType() {
+    return InternalSpanTypes.RPC;
+  }
+
+  @Override
+  protected String service() {
+    return null;
+  }
+
+  public AgentSpan onRequest(AgentSpan span, SofaRequest request) {
+    span.setTag("rpc.system", "sofarpc");
+    if (request == null) {
+      return span;
+    }
+    String serviceName = request.getTargetServiceUniqueName();
+    String methodName = request.getMethodName();
+    span.setTag(Tags.RPC_SERVICE, serviceName);
+    span.setTag("rpc.method", methodName);
+    // peer.service is derived automatically by PeerServiceCalculator from rpc.service.
+    if (serviceName != null && methodName != null) {
+      span.setResourceName(serviceName + "/" + methodName);
+    } else if (methodName != null) {
+      span.setResourceName(methodName);
+    }
+    return span;
+  }
+
+  public AgentSpan onResponse(AgentSpan span, SofaResponse response) {
+    if (response != null && response.isError()) {
+      span.setError(true);
+      span.setTag("error.message", response.getErrorMsg());
+    }
+    return span;
+  }
+}