chore: refactor RangeProjectionConfig to support more than one type of ObjectReadSession integration

BenWhitehead · web-flow · commit 50e4589dc0e6 · 2025-03-17T15:22:59.000-04:00
This is pre-work to allow creating a SeekableByteChannel that will use the ObjectReadSession rather than reading from the stream directly.
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/ObjectReadSessionImpl.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/ObjectReadSessionImpl.java
@@ -21,6 +21,7 @@
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.cloud.storage.GrpcUtils.ZeroCopyBidiStreamingCallable;
+import com.google.cloud.storage.RangeProjectionConfig.ProjectionType;
 import com.google.cloud.storage.RetryContext.RetryContextProvider;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.storage.v2.BidiReadObjectRequest;
@@ -31,6 +32,7 @@
 import java.util.ArrayList;
 import java.util.IdentityHashMap;
 import java.util.Iterator;
+import java.util.Locale;
 import java.util.Map.Entry;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
@@ -83,11 +85,23 @@ public <Projection> Projection readRange(
     lock.lock();
     try {
       checkState(open, "stream already closed");
-      long readId = state.newReadId();
-      ObjectReadSessionStreamRead<Projection> read =
-          config.cast().newRead(readId, range, retryContextProvider.create());
-      registerReadInState(readId, read);
-      return read.project();
+      switch (config.getType()) {
+        case STREAM_READ:
+          long readId = state.newReadId();
+          ObjectReadSessionStreamRead<Projection> read =
+              config.cast().newRead(readId, range, retryContextProvider.create());
+          registerReadInState(readId, read);
+          return read.project();
+        case SESSION_USER:
+          return config.project(range, this, IOAutoCloseable.noOp());
+        default:
+          throw new IllegalStateException(
+              String.format(
+                  Locale.US,
+                  "Broken java enum %s value=%s",
+                  ProjectionType.class.getName(),
+                  config.getType().name()));
+      }
     } finally {
       lock.unlock();
     }
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
@@ -1816,21 +1816,43 @@ public void copyChunk() {
 
   private static final class OtelRangeProjectionConfig<Projection>
       extends RangeProjectionConfig<Projection> {
-    private final BaseConfig<Projection, ?> delegate;
+    private final RangeProjectionConfig<Projection> delegate;
     private final Span parentSpan;
 
     private OtelRangeProjectionConfig(RangeProjectionConfig<Projection> delegate, Span parentSpan) {
-      this.delegate = delegate.cast();
+      this.delegate = delegate;
       this.parentSpan = parentSpan;
     }
 
     @Override
     BaseConfig<Projection, ?> cast() {
-      return new OtelBaseConfigDecorator();
+      return new OtelBaseConfigDecorator(delegate.cast());
+    }
+
+    @Override
+    public ProjectionType getType() {
+      return delegate.getType();
+    }
+
+    @Override
+    Projection project(RangeSpec range, ObjectReadSession session, IOAutoCloseable closeAlongWith) {
+      try {
+        return delegate.project(range, session, closeAlongWith.andThen(parentSpan::end));
+      } catch (Throwable t) {
+        parentSpan.recordException(t);
+        parentSpan.setStatus(StatusCode.ERROR, t.getClass().getSimpleName());
+        parentSpan.end();
+        throw t;
+      }
     }
 
     private class OtelBaseConfigDecorator
         extends BaseConfig<Projection, ObjectReadSessionStreamRead<Projection>> {
+      private final BaseConfig<Projection, ?> delegate;
+
+      private OtelBaseConfigDecorator(BaseConfig<Projection, ?> delegate) {
+        this.delegate = delegate;
+      }
 
       @Override
       ObjectReadSessionStreamRead<Projection> newRead(
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/RangeProjectionConfig.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/RangeProjectionConfig.java
@@ -17,10 +17,30 @@
 package com.google.cloud.storage;
 
 import com.google.cloud.storage.RangeProjectionConfigs.BaseConfig;
+import java.util.Locale;
 
 public abstract class RangeProjectionConfig<Projection> {
 
   RangeProjectionConfig() {}
 
-  abstract BaseConfig<Projection, ?> cast();
+  BaseConfig<Projection, ?> cast() {
+    throw new UnsupportedOperationException(String.format("%s#cast()", this.getClass().getName()));
+  }
+
+  abstract ProjectionType getType();
+
+  Projection project(RangeSpec range, ObjectReadSession session, IOAutoCloseable closeAlongWith) {
+    throw new UnsupportedOperationException(
+        String.format(Locale.US, "%s#project()", this.getClass().getName()));
+  }
+
+  enum ProjectionType {
+    /** Those projections which translate to a direct read registered in the state of the stream */
+    STREAM_READ,
+    /**
+     * Those projections which use an ObjectReadSession rather than directly registering a read in
+     * the stream state.
+     */
+    SESSION_USER
+  }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/RangeProjectionConfigs.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/RangeProjectionConfigs.java
@@ -27,6 +27,11 @@ public final class RangeProjectionConfigs {
   abstract static class BaseConfig<Projection, Read extends ObjectReadSessionStreamRead<Projection>>
       extends RangeProjectionConfig<Projection> {
     abstract Read newRead(long readId, RangeSpec range, RetryContext retryContext);
+
+    @Override
+    ProjectionType getType() {
+      return ProjectionType.STREAM_READ;
+    }
   }
 
   /**
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageDataClient.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageDataClient.java
@@ -23,6 +23,7 @@
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.cloud.storage.GrpcUtils.ZeroCopyBidiStreamingCallable;
 import com.google.cloud.storage.ObjectReadSessionState.OpenArguments;
+import com.google.cloud.storage.RangeProjectionConfig.ProjectionType;
 import com.google.cloud.storage.RetryContext.RetryContextProvider;
 import com.google.storage.v2.BidiReadObjectRequest;
 import com.google.storage.v2.BidiReadObjectResponse;
@@ -77,6 +78,10 @@ <Projection> ApiFuture<FastOpenObjectReadSession<Projection>> fastOpenReadSessio
     checkArgument(
         openRequest.getReadRangesList().isEmpty(),
         "ranges included in the initial request are not supported");
+    checkArgument(
+        config.getType() == ProjectionType.STREAM_READ,
+        "unsupported RangeProjectionConfig: %s",
+        config.getClass().getName());
     ObjectReadSessionState state = new ObjectReadSessionState(ctx, openRequest);
 
     ZeroCopyBidiStreamingCallable<BidiReadObjectRequest, BidiReadObjectResponse> callable =
