chore: add otel integration to retry context

BenWhitehead · BenWhitehead · commit 2e8a0eef40df · 2025-03-14T17:38:20.000-04:00
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java
@@ -119,7 +119,6 @@
 import java.util.Set;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Function;
-import java.util.function.UnaryOperator;
 import java.util.logging.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
 
@@ -838,7 +837,9 @@ public Storage create(StorageOptions options) {
                       grpcStorageOptions.getRetryAlgorithmManager().idempotent()));
 
           OpenTelemetry otel = options.getOpenTelemetry();
-          DefaultRetrier retrier = new DefaultRetrier(UnaryOperator.identity(), grpcStorageOptions);
+          DefaultRetrier retrier =
+              new DefaultRetrier(
+                  OtelStorageDecorator.retryContextDecorator(otel), grpcStorageOptions);
           if (ZeroCopyReadinessChecker.isReady()) {
             LOGGER.config("zero-copy protobuf deserialization available, using it");
             StorageStubSettings baseSettings =
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/HttpStorageOptions.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/HttpStorageOptions.java
@@ -49,7 +49,6 @@
 import java.time.Clock;
 import java.util.Objects;
 import java.util.Set;
-import java.util.function.UnaryOperator;
 import org.checkerframework.checker.nullness.qual.NonNull;
 
 /** @since 2.14.0 */
@@ -405,7 +404,7 @@ public Storage create(StorageOptions options) {
                   factory,
                   new HttpRetrier(
                       new DefaultRetrier(
-                          UnaryOperator.identity(),
+                          OtelStorageDecorator.retryContextDecorator(otel),
                           RetryingDependencies.simple(
                               options.getClock(), options.getRetrySettings()))));
           return OtelStorageDecorator.decorate(storage, otel, Transport.HTTP);
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/OtelStorageDecorator.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/OtelStorageDecorator.java
@@ -59,6 +59,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.concurrent.TimeUnit;
+import java.util.function.UnaryOperator;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -1520,6 +1521,14 @@ static Storage decorate(Storage delegate, OpenTelemetry otel, Transport transpor
     return new OtelStorageDecorator(delegate, otel, baseAttributes);
   }
 
+  static UnaryOperator<RetryContext> retryContextDecorator(OpenTelemetry otel) {
+    requireNonNull(otel, "otel must be non null");
+    if (otel == OpenTelemetry.noop()) {
+      return UnaryOperator.identity();
+    }
+    return ctx -> new OtelRetryContextDecorator(ctx, Span.current());
+  }
+
   private static @NonNull String fmtBucket(String bucket) {
     return String.format(Locale.US, "gs://%s/", bucket);
   }
@@ -1826,8 +1835,10 @@ private class OtelBaseConfigDecorator
       @Override
       ObjectReadSessionStreamRead<Projection> newRead(
           long readId, RangeSpec range, RetryContext retryContext) {
+        OtelRetryContextDecorator otelRetryContext =
+            new OtelRetryContextDecorator(retryContext, parentSpan);
         ObjectReadSessionStreamRead<Projection> read =
-            delegate.newRead(readId, range, retryContext);
+            delegate.newRead(readId, range, otelRetryContext);
         read.setOnCloseCallback(parentSpan::end);
         return new OtelDecoratingObjectReadSessionStreamRead<>(read, parentSpan);
       }
@@ -1839,6 +1850,42 @@ ObjectReadSessionStreamRead<Projection> newRead(
     }
   }
 
+  private static final class OtelRetryContextDecorator implements RetryContext {
+    private final RetryContext delegate;
+    private final Span span;
+
+    private OtelRetryContextDecorator(RetryContext delegate, Span span) {
+      this.delegate = delegate;
+      this.span = span;
+    }
+
+    @Override
+    public boolean inBackoff() {
+      return delegate.inBackoff();
+    }
+
+    @Override
+    public void reset() {
+      delegate.reset();
+    }
+
+    @Override
+    public <T extends Throwable> void recordError(
+        T t, OnSuccess onSuccess, OnFailure<T> onFailure) {
+      span.recordException(t);
+      delegate.recordError(
+          t,
+          () -> {
+            span.addEvent("retrying");
+            onSuccess.onSuccess();
+          },
+          (tt) -> {
+            span.addEvent("terminal_failure");
+            onFailure.onFailure(tt);
+          });
+    }
+  }
+
   @VisibleForTesting
   class OtelDecoratingBlobReadSession implements BlobReadSession {
 
