impl(o11y): introduce http.response.status_code

gax-java/gax/src/main/java/com/google/api/gax/tracing/ObservabilityAttributes.java

@@ -84,4 +84,10 @@ public class ObservabilityAttributes {
 
   /** The url template of the request (e.g. /v1/{name}:access). */
   public static final String URL_TEMPLATE_ATTRIBUTE = "url.template";
+
+  /** The HTTP status code of the request (e.g., 200, 404). */
+  public static final String HTTP_RESPONSE_STATUS_ATTRIBUTE = "http.response.status_code";
+
+  /** The gRPC status code of the request (e.g., 0, 14). */
+  public static final String RPC_GRPC_STATUS_ATTRIBUTE = "rpc.grpc.status_code";
 }

gax-java/gax/src/main/java/com/google/api/gax/tracing/ObservabilityUtils.java

@@ -56,6 +56,66 @@ static String extractStatus(@Nullable Throwable error) {
     return statusString;
   }
 
+  static void populateStatusAttributes(
+      Map<String, Object> attributes,
+      @Nullable Throwable error,
+      ApiTracerContext.Transport transport) {
+    if (transport == ApiTracerContext.Transport.GRPC) {
+      populateGrpcStatusAttributes(attributes, error);
+    } else if (transport == ApiTracerContext.Transport.HTTP) {
+      populateHttpStatusAttributes(attributes, error);
+    }
+  }
+
+  private static void populateGrpcStatusAttributes(
+      Map<String, Object> attributes, @Nullable Throwable error) {
+    if (error == null) {
+      attributes.put(
+          ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE, StatusCode.Code.OK.toString());
+      attributes.put(
+          ObservabilityAttributes.RPC_GRPC_STATUS_ATTRIBUTE, (long) StatusCode.Code.OK.ordinal());
+    } else if (error instanceof ApiException) {
+      attributes.put(
+          ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE,
+          ((ApiException) error).getStatusCode().getCode().toString());
+      attributes.put(
+          ObservabilityAttributes.RPC_GRPC_STATUS_ATTRIBUTE,
+          (long) ((ApiException) error).getStatusCode().getCode().ordinal());
+    } else {
+      attributes.put(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE, extractStatus(error));
+      StatusCode.Code code = StatusCode.Code.UNKNOWN;
+      if (error instanceof CancellationException) {
+        code = StatusCode.Code.CANCELLED;
+      }
+      attributes.put(ObservabilityAttributes.RPC_GRPC_STATUS_ATTRIBUTE, (long) code.ordinal());
+    }
+  }
+
+  private static void populateHttpStatusAttributes(
+      Map<String, Object> attributes, @Nullable Throwable error) {
+    if (error == null) {
+      attributes.put(ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE, 200L);
+    } else if (error instanceof ApiException) {
+      Object transportCode = ((ApiException) error).getStatusCode().getTransportCode();
+      if (transportCode instanceof Integer) {
+        attributes.put(
+            ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE,
+            ((Integer) transportCode).longValue());
+      } else {
+        attributes.put(
+            ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE,
+            (long) ((ApiException) error).getStatusCode().getCode().getHttpStatusCode());
+      }
+    } else {
+      StatusCode.Code code = StatusCode.Code.UNKNOWN;
+      if (error instanceof CancellationException) {
+        code = StatusCode.Code.CANCELLED;
+      }
+      attributes.put(
+          ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE, (long) code.getHttpStatusCode());
+    }
+  }
+
   static Attributes toOtelAttributes(Map<String, Object> attributes) {
     AttributesBuilder attributesBuilder = Attributes.builder();
     if (attributes == null) {
@@ -67,6 +127,8 @@ static Attributes toOtelAttributes(Map<String, Object> attributes) {
             attributesBuilder.put(k, (String) v);
           } else if (v instanceof Integer) {
             attributesBuilder.put(k, (long) (Integer) v);
+          } else if (v instanceof Long) {
+            attributesBuilder.put(k, (Long) v);
           }
         });
     return attributesBuilder.build();

gax-java/gax/src/main/java/com/google/api/gax/tracing/OpenTelemetryTraceManager.java

@@ -71,6 +71,13 @@ private OtelSpan(io.opentelemetry.api.trace.Span span) {
       this.span = span;
     }
 
+    @Override
+    public void addAttributes(Map<String, Object> attributes) {
+      if (attributes != null) {
+        span.setAllAttributes(ObservabilityUtils.toOtelAttributes(attributes));
+      }
+    }
+
     @Override
     public void end() {
       span.end();

gax-java/gax/src/main/java/com/google/api/gax/tracing/SpanTracer.java

@@ -34,6 +34,7 @@
 import com.google.api.core.InternalApi;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.CancellationException;
 
 /**
  * An implementation of {@link ApiTracer} that uses a {@link TraceManager} to record traces. This
@@ -82,11 +83,44 @@ public void attemptStarted(Object request, int attemptNumber) {
 
   @Override
   public void attemptSucceeded() {
-    endAttempt();
+    endAttempt(null);
   }
 
-  private void endAttempt() {
+  @Override
+  public void attemptCancelled() {
+    endAttempt(new CancellationException());
+  }
+
+  @Override
+  public void attemptFailedRetriesExhausted(Throwable error) {
+    endAttempt(error);
+  }
+
+  @Override
+  public void attemptPermanentFailure(Throwable error) {
+    endAttempt(error);
+  }
+
+  @Override
+  public void attemptFailedDuration(Throwable error, java.time.Duration delay) {
+    endAttempt(error);
+  }
+
+  @Override
+  public void attemptFailed(Throwable error, org.threeten.bp.Duration delay) {
+    endAttempt(error);
+  }
+
+  private void endAttempt(Throwable error) {
     if (attemptHandle != null) {
+      Map<String, Object> endAttributes = new HashMap<>();
+      ObservabilityUtils.populateStatusAttributes(
+          endAttributes, error, this.apiTracerContext.transport());
+
+      if (!endAttributes.isEmpty()) {
+        attemptHandle.addAttributes(endAttributes);
+      }
+
       attemptHandle.end();
       attemptHandle = null;
     }

gax-java/gax/src/main/java/com/google/api/gax/tracing/TraceManager.java

@@ -45,6 +45,12 @@ public interface TraceManager {
   Span createSpan(String name, Map<String, Object> attributes);
 
   interface Span {
+    /**
+     * Adds an attribute to the span. This is useful for adding attributes that are only available
+     * after the span has started, e.g. status codes.
+     */
+    void addAttributes(Map<String, Object> attributes);
+
     void end();
   }
 }

gax-java/gax/src/test/java/com/google/api/gax/tracing/ObservabilityUtilsTest.java

@@ -114,6 +114,102 @@ void testToOtelAttributes_shouldMapIntAttributes() {
         .isEqualTo((long) attribute2Value);
   }
 
+  @Test
+  void testPopulateStatusAttributes_grpc_success() {
+    Map<String, Object> attributes = new java.util.HashMap<>();
+    ObservabilityUtils.populateStatusAttributes(attributes, null, ApiTracerContext.Transport.GRPC);
+    assertThat(attributes)
+        .containsEntry(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE, "OK");
+  }
+
+  @Test
+  void testPopulateStatusAttributes_grpc_apiException() {
+    Map<String, Object> attributes = new java.util.HashMap<>();
+    ApiException error =
+        new ApiException("fake_error", null, new FakeStatusCode(StatusCode.Code.NOT_FOUND), false);
+    ObservabilityUtils.populateStatusAttributes(attributes, error, ApiTracerContext.Transport.GRPC);
+    assertThat(attributes)
+        .containsEntry(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE, "NOT_FOUND");
+  }
+
+  @Test
+  void testPopulateStatusAttributes_grpc_cancellationException() {
+    Map<String, Object> attributes = new java.util.HashMap<>();
+    Throwable error = new java.util.concurrent.CancellationException();
+    ObservabilityUtils.populateStatusAttributes(attributes, error, ApiTracerContext.Transport.GRPC);
+    assertThat(attributes)
+        .containsEntry(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE, "CANCELLED");
+  }
+
+  @Test
+  void testPopulateStatusAttributes_http_success() {
+    Map<String, Object> attributes = new java.util.HashMap<>();
+    ObservabilityUtils.populateStatusAttributes(attributes, null, ApiTracerContext.Transport.HTTP);
+    assertThat(attributes)
+        .containsEntry(ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE, 200L);
+  }
+
+  @Test
+  void testPopulateStatusAttributes_http_apiExceptionWithIntegerTransportCode() {
+    Map<String, Object> attributes = new java.util.HashMap<>();
+    ApiException error =
+        new ApiException(
+            "fake_error",
+            null,
+            new com.google.api.gax.rpc.StatusCode() {
+              @Override
+              public Code getCode() {
+                return Code.NOT_FOUND;
+              }
+
+              @Override
+              public Object getTransportCode() {
+                return 404;
+              }
+            },
+            false);
+    ObservabilityUtils.populateStatusAttributes(attributes, error, ApiTracerContext.Transport.HTTP);
+    assertThat(attributes)
+        .containsEntry(ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE, 404L);
+  }
+
+  @Test
+  void testPopulateStatusAttributes_http_apiExceptionWithNonIntegerTransportCode() {
+    Map<String, Object> attributes = new java.util.HashMap<>();
+    ApiException error =
+        new ApiException(
+            "fake_error",
+            null,
+            new com.google.api.gax.rpc.StatusCode() {
+              @Override
+              public Code getCode() {
+                return Code.NOT_FOUND;
+              }
+
+              @Override
+              public Object getTransportCode() {
+                return "Not Found";
+              }
+            },
+            false);
+    ObservabilityUtils.populateStatusAttributes(attributes, error, ApiTracerContext.Transport.HTTP);
+    assertThat(attributes)
+        .containsEntry(
+            ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE,
+            (long) StatusCode.Code.NOT_FOUND.getHttpStatusCode());
+  }
+
+  @Test
+  void testPopulateStatusAttributes_http_cancellationException() {
+    Map<String, Object> attributes = new java.util.HashMap<>();
+    Throwable error = new java.util.concurrent.CancellationException();
+    ObservabilityUtils.populateStatusAttributes(attributes, error, ApiTracerContext.Transport.HTTP);
+    assertThat(attributes)
+        .containsEntry(
+            ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE,
+            (long) StatusCode.Code.CANCELLED.getHttpStatusCode());
+  }
+
   @Test
   void testToOtelAttributes_shouldReturnEmptyAttributes_nullInput() {
     assertThat(ObservabilityUtils.toOtelAttributes(null)).isEqualTo(Attributes.empty());

gax-java/gax/src/test/java/com/google/api/gax/tracing/SpanTracerTest.java

@@ -76,4 +76,122 @@ void testAttemptStarted_includesLanguageAttribute() {
     assertThat(attributesCaptor.getValue())
         .containsEntry(SpanTracer.LANGUAGE_ATTRIBUTE, SpanTracer.DEFAULT_LANGUAGE);
   }
+
+  @Test
+  void testAttemptSucceeded_grpc() {
+    ApiTracerContext context =
+        ApiTracerContext.newBuilder()
+            .setLibraryMetadata(com.google.api.gax.rpc.LibraryMetadata.empty())
+            .setTransport(ApiTracerContext.Transport.GRPC)
+            .build();
+    tracer = new SpanTracer(recorder, context, ATTEMPT_SPAN_NAME);
+    when(recorder.createSpan(eq(ATTEMPT_SPAN_NAME), anyMap())).thenReturn(attemptHandle);
+
+    tracer.attemptStarted(new Object(), 1);
+    tracer.attemptSucceeded();
+
+    ArgumentCaptor<Map<String, Object>> attrsCaptor = ArgumentCaptor.forClass(Map.class);
+    verify(attemptHandle).addAttributes(attrsCaptor.capture());
+    verify(attemptHandle).end();
+
+    assertThat(attrsCaptor.getValue())
+        .containsEntry(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE, "OK");
+  }
+
+  @Test
+  void testAttemptSucceeded_http() {
+    ApiTracerContext context =
+        ApiTracerContext.newBuilder()
+            .setLibraryMetadata(com.google.api.gax.rpc.LibraryMetadata.empty())
+            .setTransport(ApiTracerContext.Transport.HTTP)
+            .build();
+    tracer = new SpanTracer(recorder, context, ATTEMPT_SPAN_NAME);
+    when(recorder.createSpan(eq(ATTEMPT_SPAN_NAME), anyMap())).thenReturn(attemptHandle);
+
+    tracer.attemptStarted(new Object(), 1);
+    tracer.attemptSucceeded();
+
+    ArgumentCaptor<Map<String, Object>> attrsCaptor = ArgumentCaptor.forClass(Map.class);
+    verify(attemptHandle).addAttributes(attrsCaptor.capture());
+    verify(attemptHandle).end();
+
+    assertThat(attrsCaptor.getValue())
+        .containsEntry(ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE, 200L);
+  }
+
+  @Test
+  void testAttemptFailed_grpc() {
+    ApiTracerContext context =
+        ApiTracerContext.newBuilder()
+            .setLibraryMetadata(com.google.api.gax.rpc.LibraryMetadata.empty())
+            .setTransport(ApiTracerContext.Transport.GRPC)
+            .build();
+    tracer = new SpanTracer(recorder, context, ATTEMPT_SPAN_NAME);
+    when(recorder.createSpan(eq(ATTEMPT_SPAN_NAME), anyMap())).thenReturn(attemptHandle);
+
+    com.google.api.gax.rpc.ApiException exception =
+        new com.google.api.gax.rpc.ApiException(
+            "error",
+            null,
+            new com.google.api.gax.rpc.StatusCode() {
+              @Override
+              public Code getCode() {
+                return Code.NOT_FOUND;
+              }
+
+              @Override
+              public Object getTransportCode() {
+                return null;
+              }
+            },
+            false);
+
+    tracer.attemptStarted(new Object(), 1);
+    tracer.attemptFailedRetriesExhausted(exception);
+
+    ArgumentCaptor<Map<String, Object>> attrsCaptor = ArgumentCaptor.forClass(Map.class);
+    verify(attemptHandle).addAttributes(attrsCaptor.capture());
+    verify(attemptHandle).end();
+
+    assertThat(attrsCaptor.getValue())
+        .containsEntry(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE, "NOT_FOUND");
+  }
+
+  @Test
+  void testAttemptFailed_http() {
+    ApiTracerContext context =
+        ApiTracerContext.newBuilder()
+            .setLibraryMetadata(com.google.api.gax.rpc.LibraryMetadata.empty())
+            .setTransport(ApiTracerContext.Transport.HTTP)
+            .build();
+    tracer = new SpanTracer(recorder, context, ATTEMPT_SPAN_NAME);
+    when(recorder.createSpan(eq(ATTEMPT_SPAN_NAME), anyMap())).thenReturn(attemptHandle);
+
+    com.google.api.gax.rpc.ApiException exception =
+        new com.google.api.gax.rpc.ApiException(
+            "error",
+            null,
+            new com.google.api.gax.rpc.StatusCode() {
+              @Override
+              public Code getCode() {
+                return Code.NOT_FOUND;
+              }
+
+              @Override
+              public Object getTransportCode() {
+                return 404;
+              }
+            },
+            false);
+
+    tracer.attemptStarted(new Object(), 1);
+    tracer.attemptFailedRetriesExhausted(exception);
+
+    ArgumentCaptor<Map<String, Object>> attrsCaptor = ArgumentCaptor.forClass(Map.class);
+    verify(attemptHandle).addAttributes(attrsCaptor.capture());
+    verify(attemptHandle).end();
+
+    assertThat(attrsCaptor.getValue())
+        .containsEntry(ObservabilityAttributes.HTTP_RESPONSE_STATUS_ATTRIBUTE, 404L);
+  }
 }