From 3e63ae097ccdd3ea44ccccc5183e187295ddb726 Mon Sep 17 00:00:00 2001
From: Andrea Marziali <andrea.marziali@datadoghq.com>
Date: Mon, 26 Aug 2024 11:28:43 +0200
Subject: [PATCH] Fix grpc server error mark

---
 .../grpc/client/GrpcClientDecorator.java      |  5 +---
 .../grpc/server/GrpcServerDecorator.java      | 20 ++++++++++++---
 .../src/test/groovy/ArmeriaGrpcTest.groovy    | 25 ++++++++++++-------
 .../grpc/client/GrpcClientDecorator.java      |  5 +---
 .../grpc/server/GrpcServerDecorator.java      | 20 ++++++++++++---
 .../grpc-1.5/src/test/groovy/GrpcTest.groovy  | 25 ++++++++++++-------
 6 files changed, 68 insertions(+), 32 deletions(-)

diff --git a/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java b/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java
index 7a6f76d200a..060f88eaa43 100644
--- a/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java
+++ b/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java
@@ -115,10 +115,7 @@ public AgentSpan onClose(final AgentSpan span, final Status status) {
 
     // TODO why is there a mismatch between client / server for calling the onError method?
     onError(span, status.getCause());
-    if (CLIENT_ERROR_STATUSES.get(status.getCode().value())) {
-      span.setError(true);
-    }
-
+    span.setError(CLIENT_ERROR_STATUSES.get(status.getCode().value()));
     return span;
   }
 }
diff --git a/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/GrpcServerDecorator.java b/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/GrpcServerDecorator.java
index de5058948e2..7d51b672439 100644
--- a/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/GrpcServerDecorator.java
+++ b/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/GrpcServerDecorator.java
@@ -14,6 +14,8 @@
 import datadog.trace.bootstrap.instrumentation.decorator.ServerDecorator;
 import io.grpc.ServerCall;
 import io.grpc.Status;
+import io.grpc.StatusException;
+import io.grpc.StatusRuntimeException;
 import java.util.BitSet;
 import java.util.LinkedHashMap;
 import java.util.function.Function;
@@ -97,15 +99,27 @@ public <RespT, ReqT> AgentSpan onCall(final AgentSpan span, ServerCall<ReqT, Res
     return span;
   }
 
-  public AgentSpan onClose(final AgentSpan span, final Status status) {
+  public AgentSpan onStatus(final AgentSpan span, final Status status) {
     span.setTag("status.code", status.getCode().name());
     span.setTag("status.description", status.getDescription());
+    return span.setError(SERVER_ERROR_STATUSES.get(status.getCode().value()));
+  }
 
-    if (SERVER_ERROR_STATUSES.get(status.getCode().value())) {
+  public AgentSpan onClose(final AgentSpan span, final Status status) {
+    if (status.getCause() != null) {
       onError(span, status.getCause());
-      span.setError(true);
     }
+    return onStatus(span, status);
+  }
 
+  @Override
+  public AgentSpan onError(AgentSpan span, Throwable throwable) {
+    super.onError(span, throwable);
+    if (throwable instanceof StatusRuntimeException) {
+      onStatus(span, ((StatusRuntimeException) throwable).getStatus());
+    } else if (throwable instanceof StatusException) {
+      onStatus(span, ((StatusException) throwable).getStatus());
+    }
     return span;
   }
 }
diff --git a/dd-java-agent/instrumentation/armeria-grpc/src/test/groovy/ArmeriaGrpcTest.groovy b/dd-java-agent/instrumentation/armeria-grpc/src/test/groovy/ArmeriaGrpcTest.groovy
index 9aa49c386d8..5cb57f6f331 100644
--- a/dd-java-agent/instrumentation/armeria-grpc/src/test/groovy/ArmeriaGrpcTest.groovy
+++ b/dd-java-agent/instrumentation/armeria-grpc/src/test/groovy/ArmeriaGrpcTest.groovy
@@ -1,3 +1,5 @@
+import static datadog.trace.api.config.TraceInstrumentationConfig.GRPC_SERVER_ERROR_STATUSES
+
 import com.google.common.util.concurrent.ListenableFuture
 import com.linecorp.armeria.client.Clients
 import com.linecorp.armeria.common.SessionProtocol
@@ -74,6 +76,7 @@ abstract class ArmeriaGrpcTest extends VersionedNamingTestBase {
     // here to trigger wrapping to record scheduling time - the logic is trivial so it's enough to verify
     // that ClassCastExceptions do not arise from the wrapping
     injectSysConfig("dd.profiling.enabled", "true")
+    injectSysConfig(GRPC_SERVER_ERROR_STATUSES, "2-14", true)
   }
 
   @Override
@@ -436,12 +439,14 @@ abstract class ArmeriaGrpcTest extends VersionedNamingTestBase {
           resourceName "example.Greeter/SayHello"
           spanType DDSpanTypes.RPC
           childOf trace(0).get(0)
-          errored true
+          errored errorFlag
           measured true
           tags {
             "$Tags.COMPONENT" "armeria-grpc-server"
             "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
             errorTags error.class, error.message
+            "status.code" "${status.code.name()}"
+            "status.description"  { it == null || String}
             "canceled" { true } // 1.0.0 handles cancellation incorrectly so accesting any value
             if ({ isDataStreamsEnabled() }) {
               "$DDTags.PATHWAY_HASH" { String }
@@ -470,13 +475,15 @@ abstract class ArmeriaGrpcTest extends VersionedNamingTestBase {
     serverRule.stop().get()
 
     where:
-    name                          | status
-    "Runtime - cause"             | Status.UNKNOWN.withCause(new RuntimeException("some error"))
-    "Status - cause"              | Status.PERMISSION_DENIED.withCause(new RuntimeException("some error"))
-    "StatusRuntime - cause"       | Status.UNIMPLEMENTED.withCause(new RuntimeException("some error"))
-    "Runtime - description"       | Status.UNKNOWN.withDescription("some description")
-    "Status - description"        | Status.PERMISSION_DENIED.withDescription("some description")
-    "StatusRuntime - description" | Status.UNIMPLEMENTED.withDescription("some description")
+    name                                    | status                                                                  | errorFlag
+    "Runtime - cause"                       | Status.UNKNOWN.withCause(new RuntimeException("some error"))            | true
+    "Status - cause"                        | Status.PERMISSION_DENIED.withCause(new RuntimeException("some error"))  | true
+    "StatusRuntime - cause"                 | Status.UNIMPLEMENTED.withCause(new RuntimeException("some error"))      | true
+    "Runtime - description"                 | Status.UNKNOWN.withDescription("some description")                      | true
+    "Status - description"                  | Status.PERMISSION_DENIED.withDescription("some description")            | true
+    "StatusRuntime - description"           | Status.UNIMPLEMENTED.withDescription("some description")                | true
+    "StatusRuntime - Not errored no cause"   | Status.fromCodeValue(15).withDescription("some description")           | false
+    "StatusRuntime - Not errored with cause" | Status.fromCodeValue(15).withCause(new RuntimeException("some error")) | false
   }
 
   def "skip binary headers"() {
@@ -672,7 +679,7 @@ abstract class ArmeriaGrpcDataStreamsEnabledForkedTest extends ArmeriaGrpcTest {
   }
 }
 
-class ArmeriaGrpcDataStreamsEnabledV0ForkedTest extends ArmeriaGrpcDataStreamsEnabledForkedTest {
+class ArmeriaGrpcDataStreamsEnabledV0Test extends ArmeriaGrpcDataStreamsEnabledForkedTest {
 
   @Override
   int version() {
diff --git a/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/GrpcClientDecorator.java b/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/GrpcClientDecorator.java
index ac617282553..cded71d6c8f 100644
--- a/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/GrpcClientDecorator.java
+++ b/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/GrpcClientDecorator.java
@@ -115,10 +115,7 @@ public AgentSpan onClose(final AgentSpan span, final Status status) {
 
     // TODO why is there a mismatch between client / server for calling the onError method?
     onError(span, status.getCause());
-    if (CLIENT_ERROR_STATUSES.get(status.getCode().value())) {
-      span.setError(true);
-    }
-
+    span.setError(CLIENT_ERROR_STATUSES.get(status.getCode().value()));
     return span;
   }
 }
diff --git a/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/GrpcServerDecorator.java b/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/GrpcServerDecorator.java
index 0c2fbd2b519..7d53598b20e 100644
--- a/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/GrpcServerDecorator.java
+++ b/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/GrpcServerDecorator.java
@@ -14,6 +14,8 @@
 import datadog.trace.bootstrap.instrumentation.decorator.ServerDecorator;
 import io.grpc.ServerCall;
 import io.grpc.Status;
+import io.grpc.StatusException;
+import io.grpc.StatusRuntimeException;
 import java.util.BitSet;
 import java.util.LinkedHashMap;
 import java.util.function.Function;
@@ -97,15 +99,27 @@ public <RespT, ReqT> AgentSpan onCall(final AgentSpan span, ServerCall<ReqT, Res
     return span;
   }
 
-  public AgentSpan onClose(final AgentSpan span, final Status status) {
+  public AgentSpan onStatus(final AgentSpan span, final Status status) {
     span.setTag("status.code", status.getCode().name());
     span.setTag("status.description", status.getDescription());
+    return span.setError(SERVER_ERROR_STATUSES.get(status.getCode().value()));
+  }
 
-    if (SERVER_ERROR_STATUSES.get(status.getCode().value())) {
+  public AgentSpan onClose(final AgentSpan span, final Status status) {
+    if (status.getCause() != null) {
       onError(span, status.getCause());
-      span.setError(true);
     }
+    return onStatus(span, status);
+  }
 
+  @Override
+  public AgentSpan onError(AgentSpan span, Throwable throwable) {
+    super.onError(span, throwable);
+    if (throwable instanceof StatusRuntimeException) {
+      onStatus(span, ((StatusRuntimeException) throwable).getStatus());
+    } else if (throwable instanceof StatusException) {
+      onStatus(span, ((StatusException) throwable).getStatus());
+    }
     return span;
   }
 }
diff --git a/dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/GrpcTest.groovy b/dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/GrpcTest.groovy
index 49f76f430a4..03749eae614 100644
--- a/dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/GrpcTest.groovy
+++ b/dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/GrpcTest.groovy
@@ -1,3 +1,5 @@
+import static datadog.trace.api.config.TraceInstrumentationConfig.GRPC_SERVER_ERROR_STATUSES
+
 import com.google.common.util.concurrent.MoreExecutors
 import datadog.trace.agent.test.naming.VersionedNamingTestBase
 import datadog.trace.api.DDSpanId
@@ -71,6 +73,7 @@ abstract class GrpcTest extends VersionedNamingTestBase {
     // here to trigger wrapping to record scheduling time - the logic is trivial so it's enough to verify
     // that ClassCastExceptions do not arise from the wrapping
     injectSysConfig("dd.profiling.enabled", "true")
+    injectSysConfig(GRPC_SERVER_ERROR_STATUSES, "2-14", true)
   }
 
   def setupSpec() {
@@ -422,11 +425,13 @@ abstract class GrpcTest extends VersionedNamingTestBase {
           resourceName "example.Greeter/SayHello"
           spanType DDSpanTypes.RPC
           childOf trace(0).get(0)
-          errored true
+          errored errorFlag
           measured true
           tags {
             "$Tags.COMPONENT" "grpc-server"
             "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
+            "status.code" "${status.code.name()}"
+            "status.description"  { it == null || String}
             errorTags error.class, error.message
             if ({ isDataStreamsEnabled() }) {
               "$DDTags.PATHWAY_HASH" { String }
@@ -456,13 +461,15 @@ abstract class GrpcTest extends VersionedNamingTestBase {
     server?.shutdownNow()?.awaitTermination()
 
     where:
-    name                          | status
-    "Runtime - cause"             | Status.UNKNOWN.withCause(new RuntimeException("some error"))
-    "Status - cause"              | Status.PERMISSION_DENIED.withCause(new RuntimeException("some error"))
-    "StatusRuntime - cause"       | Status.UNIMPLEMENTED.withCause(new RuntimeException("some error"))
-    "Runtime - description"       | Status.UNKNOWN.withDescription("some description")
-    "Status - description"        | Status.PERMISSION_DENIED.withDescription("some description")
-    "StatusRuntime - description" | Status.UNIMPLEMENTED.withDescription("some description")
+    name                                    | status                                                                  | errorFlag
+    "Runtime - cause"                       | Status.UNKNOWN.withCause(new RuntimeException("some error"))            | true
+    "Status - cause"                        | Status.PERMISSION_DENIED.withCause(new RuntimeException("some error"))  | true
+    "StatusRuntime - cause"                 | Status.UNIMPLEMENTED.withCause(new RuntimeException("some error"))      | true
+    "Runtime - description"                 | Status.UNKNOWN.withDescription("some description")                      | true
+    "Status - description"                  | Status.PERMISSION_DENIED.withDescription("some description")            | true
+    "StatusRuntime - description"           | Status.UNIMPLEMENTED.withDescription("some description")                | true
+    "StatusRuntime - Not errored no cause"   | Status.fromCodeValue(15).withDescription("some description")           | false
+    "StatusRuntime - Not errored with cause" | Status.fromCodeValue(15).withCause(new RuntimeException("some error")) | false
   }
 
   def "skip binary headers"() {
@@ -642,7 +649,7 @@ abstract class GrpcDataStreamsEnabledForkedTest extends GrpcTest {
   }
 }
 
-class GrpcDataStreamsEnabledV0ForkedTest extends GrpcDataStreamsEnabledForkedTest {
+class GrpcDataStreamsEnabledV0Test extends GrpcDataStreamsEnabledForkedTest {
 
   @Override
   int version() {