feat(storage): serve the fast open read through wrapped configs

java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ObjectReadSessionImpl.java

@@ -91,12 +91,13 @@ public <Projection> Projection readAs(ReadProjectionConfig<Projection> config) {
     checkState(open, "Session already closed");
     switch (config.getType()) {
       case STREAM_READ:
-        ObjectReadSessionStreamRead<?> fastOpenRead = fastOpenReads.remove(config);
+        ObjectReadSessionStreamRead<?> fastOpenRead = fastOpenReads.remove(config.unwrap());
         if (fastOpenRead != null) {
           // configs are final classes with a class-exact equals, so an equal config has the same
           // Projection type
           @SuppressWarnings("unchecked")
           Projection projection = (Projection) fastOpenRead.project();
+          config.onFastOpenReadClaimed(projection);
           return projection;
         }
         long readId = state.newReadId();

java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/OtelStorageDecorator.java

@@ -19,6 +19,7 @@
 import static java.util.Objects.requireNonNull;
 
 import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutureCallback;
 import com.google.api.core.ApiFutures;
 import com.google.api.core.BetaApi;
 import com.google.api.gax.paging.Page;
@@ -1892,6 +1893,36 @@ public ProjectionType getType() {
       return delegate.getType();
     }
 
+    @Override
+    ReadProjectionConfig<Projection> unwrap() {
+      return delegate.unwrap();
+    }
+
+    @Override
+    void onFastOpenReadClaimed(Projection projection) {
+      if (!(projection instanceof ApiFuture)) {
+        parentSpan.end();
+        return;
+      }
+      // the fast open read may still be in progress or may fail, so end the span when it's done
+      ApiFutures.addCallback(
+          (ApiFuture<?>) projection,
+          new ApiFutureCallback<Object>() {
+            @Override
+            public void onFailure(Throwable t) {
+              parentSpan.recordException(t);
+              parentSpan.setStatus(StatusCode.ERROR, t.getClass().getSimpleName());
+              parentSpan.end();
+            }
+
+            @Override
+            public void onSuccess(Object result) {
+              parentSpan.end();
+            }
+          },
+          MoreExecutors.directExecutor());
+    }
+
     @Override
     Projection project(ObjectReadSession session, IOAutoCloseable closeAlongWith) {
       try {

java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ReadProjectionConfig.java

@@ -43,6 +43,14 @@ public abstract class ReadProjectionConfig<Projection> {
 
   abstract ProjectionType getType();
 
+  /** The config this one wraps, or {@code this}. Used to match a fast open read. */
+  ReadProjectionConfig<Projection> unwrap() {
+    return this;
+  }
+
+  /** Called when {@code readAs} with this config is served from the fast open read. */
+  void onFastOpenReadClaimed(Projection projection) {}
+
   Projection project(ObjectReadSession session, IOAutoCloseable closeAlongWith) {
     throw new UnsupportedOperationException(
         String.format(Locale.US, "%s#project()", this.getClass().getName()));

java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/ITFastOpenFakeTest.java

@@ -31,6 +31,7 @@
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.rpc.OutOfRangeException;
 import com.google.cloud.storage.ITObjectReadSessionFakeTest.FakeStorage;
+import com.google.cloud.storage.ReadProjectionConfigs.BaseConfig;
 import com.google.cloud.storage.it.ChecksummedTestContent;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
@@ -198,6 +199,37 @@ OPEN_WITH_READ_OF_10_TO_20, metadataAndBytes(1, 10, 10),
         .inOrder();
   }
 
+  @Test
+  public void readAs_withAWrappedFastOpenConfig_isServedFromTheFastOpenRead() throws Exception {
+    BidiReadObjectRequest secondReadOf10To20 = read(2, 10, 10);
+    FakeStorage server =
+        fakeServer(
+            ImmutableMap.of(
+                OPEN_WITH_READ_OF_10_TO_20, metadataAndBytes(1, 10, 10),
+                secondReadOf10To20, bytes(2, 10, 10)));
+    WrappedConfig wrapped = new WrappedConfig(READ_10_TO_20);
+
+    try (FakeServer fakeServer = FakeServer.of(server);
+        Storage storage = fakeServer.getGrpcStorageOptions().getService()) {
+      StorageDataClient dataClient = maybeGetStorageDataClient(storage);
+
+      try (ObjectReadSession session =
+          dataClient.readSession(OPEN, CTX, ImmutableList.of(READ_10_TO_20)).get(5, SECONDS)) {
+        ApiFuture<byte[]> fastOpenFuture = session.readAs(wrapped);
+        assertThat(wrapped.claimed).containsExactly(fastOpenFuture);
+        byte[] fastOpen = fastOpenFuture.get(2, SECONDS);
+        byte[] second = session.readAs(wrapped).get(2, SECONDS);
+        assertThat(wrapped.claimed).containsExactly(fastOpenFuture);
+        assertThat(xxd(fastOpen)).isEqualTo(xxd(slice(10, 10)));
+        assertThat(xxd(second)).isEqualTo(xxd(slice(10, 10)));
+      }
+    }
+
+    assertThat(requestsReceived)
+        .containsExactly(OPEN_WITH_READ_OF_10_TO_20, secondReadOf10To20)
+        .inOrder();
+  }
+
   @Test
   public void close_withAnUnclaimedFastOpenRead_closesTheSession() throws Exception {
     FakeStorage server =
@@ -504,4 +536,34 @@ private static BidiReadObjectResponse bytes(int readId, int offset, int length)
   private static byte[] slice(int offset, int length) {
     return Arrays.copyOfRange(ALL_OBJECT_BYTES, offset, offset + length);
   }
+
+  /** Wraps a config the way tracing does, and records each fast open hit. */
+  private static final class WrappedConfig extends ReadProjectionConfig<ApiFuture<byte[]>> {
+    private final ReadAsFutureBytes delegate;
+    private final List<ApiFuture<byte[]>> claimed = new ArrayList<>();
+
+    private WrappedConfig(ReadAsFutureBytes delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    BaseConfig<ApiFuture<byte[]>, ?> cast() {
+      return delegate.cast();
+    }
+
+    @Override
+    ProjectionType getType() {
+      return delegate.getType();
+    }
+
+    @Override
+    ReadProjectionConfig<ApiFuture<byte[]>> unwrap() {
+      return delegate;
+    }
+
+    @Override
+    void onFastOpenReadClaimed(ApiFuture<byte[]> projection) {
+      claimed.add(projection);
+    }
+  }
 }