feat: Instrument HTTP readAllBytes (#2812)

sydney-munro · web-flow · commit d7243f4f855c · 2024-11-11T12:07:27.000-08:00
* feat: instrument HTTP readAllBytes
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java
@@ -791,9 +791,11 @@ public StorageObject compose(
 
   @Override
   public byte[] load(StorageObject from, Map<Option, ?> options) {
+    OpenTelemetryTraceUtil.Span otelSpan =
+        openTelemetryTraceUtil.startSpan("load", this.getClass().getName());
     Span span = startSpan(HttpStorageRpcSpans.SPAN_NAME_LOAD);
     Scope scope = tracer.withSpan(span);
-    try {
+    try (OpenTelemetryTraceUtil.Scope unused = otelSpan.makeCurrent()) {
       Storage.Objects.Get getRequest =
           storage
               .objects()
@@ -812,9 +814,12 @@ public byte[] load(StorageObject from, Map<Option, ?> options) {
       getRequest.executeMedia().download(out);
       return out.toByteArray();
     } catch (IOException ex) {
+      otelSpan.recordException(ex);
+      otelSpan.setStatus(StatusCode.ERROR, ex.getClass().getSimpleName());
       span.setStatus(Status.UNKNOWN.withDescription(ex.getMessage()));
       throw translate(ex);
     } finally {
+      otelSpan.end();
       scope.close();
       span.end(HttpStorageRpcSpans.END_SPAN_OPTIONS);
     }
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/ITHttpOpenTelemetryTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/ITHttpOpenTelemetryTest.java
@@ -141,6 +141,17 @@ public void runCopy() {
     Assert.assertTrue(spanData.stream().anyMatch(x -> x.getName().contains("rewrite")));
   }
 
+  @Test
+  public void runReadAllBytes() {
+    BlobInfo blobInfo = BlobInfo.newBuilder(blobId).setContentType("text/plain").build();
+    storage.create(blobInfo, helloWorldTextBytes);
+    byte[] read = storage.readAllBytes(blobId);
+    TestExporter testExported = (TestExporter) exporter;
+    List<SpanData> spanData = testExported.getExportedSpans();
+    checkCommonAttributes(spanData);
+    Assert.assertTrue(spanData.stream().anyMatch(x -> x.getName().contains("load")));
+  }
+
   private void checkCommonAttributes(List<SpanData> spanData) {
     for (SpanData span : spanData) {
       Assert.assertEquals("Storage", getAttributeValue(span, "gcp.client.service"));
