feat(storage): the session gives a fast open read to the matching readAs

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

@@ -34,7 +34,10 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Map.Entry;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.locks.ReentrantLock;
@@ -52,21 +55,29 @@ final class ObjectReadSessionImpl implements ObjectReadSession {
 
   private final ConcurrentIdentityMap<ObjectReadSessionStream, ObjectReadSessionState> children;
 
+  // Fast open reads: reads sent in the request that opened the stream, keyed by the config each was
+  // created from. The first readAs with an equal config is given the read instead of sending a new
+  // one.
+  private final ConcurrentMap<ReadProjectionConfig<?>, ObjectReadSessionStreamRead<?>>
+      fastOpenReads;
+
   private volatile boolean open;
 
   ObjectReadSessionImpl(
       ScheduledExecutorService executor,
       ZeroCopyBidiStreamingCallable<BidiReadObjectRequest, BidiReadObjectResponse> callable,
       ObjectReadSessionStream stream,
       ObjectReadSessionState state,
-      RetryContextProvider retryContextProvider) {
+      RetryContextProvider retryContextProvider,
+      Map<ReadProjectionConfig<?>, ObjectReadSessionStreamRead<?>> fastOpenReads) {
     this.executor = executor;
     this.callable = callable;
     this.stream = stream;
     this.state = state;
     this.resource = state.getMetadata();
     this.retryContextProvider = retryContextProvider;
     this.children = new ConcurrentIdentityMap<>();
+    this.fastOpenReads = new ConcurrentHashMap<>(fastOpenReads);
     this.open = true;
   }
 
@@ -80,6 +91,14 @@ public <Projection> Projection readAs(ReadProjectionConfig<Projection> config) {
     checkState(open, "Session already closed");
     switch (config.getType()) {
       case STREAM_READ:
+        ObjectReadSessionStreamRead<?> fastOpenRead = fastOpenReads.remove(config);
+        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();
+          return projection;
+        }
         long readId = state.newReadId();
         ObjectReadSessionStreamRead<Projection> read =
             config.cast().newRead(readId, retryContextProvider.create());
@@ -113,6 +132,11 @@ public void close() throws IOException {
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
       throw new InterruptedIOException();
+    } finally {
+      // a fast open read nobody asked for may already be complete and out of the state, so
+      // stream.close() above does not reach it; close it here to release its buffers
+      GrpcUtils.closeAll(fastOpenReads.values());
+      fastOpenReads.clear();
     }
   }
 

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

@@ -27,11 +27,14 @@
 import com.google.cloud.storage.ReadProjectionConfig.ProjectionType;
 import com.google.cloud.storage.RetryContext.RetryContextProvider;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.storage.v2.BidiReadObjectRequest;
 import com.google.storage.v2.BidiReadObjectResponse;
 import java.io.IOException;
 import java.time.Duration;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 
@@ -66,6 +69,9 @@ ApiFuture<ObjectReadSession> readSession(BidiReadObjectRequest req, GrpcCallCont
    * Open a new {@link ObjectReadSession} using fast open: the ranges in {@code fastOpenConfigs} are
    * sent in the same request that opens the session, saving a round trip before the first bytes.
    *
+   * <p>The first {@code readAs} with a config equal to one of these is given that fast open read
+   * instead of sending a new one.
+   *
    * @param fastOpenConfigs the reads to send with the open, for example {@code
    *     ReadProjectionConfigs.asFutureBytes().withRangeSpec(RangeSpec.of(0, 1024))}; empty to only
    *     open the session
@@ -77,6 +83,7 @@ ApiFuture<ObjectReadSession> readSession(
     checkArgument(
         req.getReadRangesList().isEmpty(),
         "ranges included in the initial request are not supported");
+    checkArgument(fastOpenConfigs.size() <= 1, "at most one fast open read is supported");
     ObjectReadSessionState state = new ObjectReadSessionState(ctx, req);
 
     ZeroCopyBidiStreamingCallable<BidiReadObjectRequest, BidiReadObjectResponse> callable =
@@ -86,18 +93,21 @@ ApiFuture<ObjectReadSession> readSession(
 
     // Add each fast open read to the session state. The open request below is built from the
     // state, so these reads are sent as part of the open.
+    Map<ReadProjectionConfig<?>, ObjectReadSessionStreamRead<?>> fastOpenReads = new HashMap<>();
     for (ReadProjectionConfig<?> config : fastOpenConfigs) {
       long readId = state.newReadId();
       ObjectReadSessionStreamRead<?> read =
           config.cast().newRead(readId, retryContextProvider.create());
       state.putOutstandingRead(readId, read);
+      fastOpenReads.put(config, read);
     }
 
     ApiFuture<ObjectReadSession> objectReadSessionFuture =
         ApiFutures.transform(
             stream,
             nowOpen ->
-                new ObjectReadSessionImpl(executor, callable, stream, state, retryContextProvider),
+                new ObjectReadSessionImpl(
+                    executor, callable, stream, state, retryContextProvider, fastOpenReads),
             executor);
     stream.send(state.getOpenArguments().getReq());
     return objectReadSessionFuture;
@@ -132,7 +142,7 @@ <Projection> ApiFuture<FastOpenObjectReadSession<Projection>> fastOpenReadSessio
             nowOpen ->
                 new FastOpenObjectReadSession<>(
                     new ObjectReadSessionImpl(
-                        executor, callable, stream, state, retryContextProvider),
+                        executor, callable, stream, state, retryContextProvider, ImmutableMap.of()),
                     read,
                     stream),
             executor);

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

@@ -18,9 +18,12 @@
 
 import static com.google.cloud.storage.ByteSizeConstants._2MiB;
 import static com.google.cloud.storage.ITObjectReadSessionFakeTest.getReadRange;
+import static com.google.cloud.storage.ITObjectReadSessionFakeTest.read;
 import static com.google.cloud.storage.PackagePrivateMethodWorkarounds.maybeGetStorageDataClient;
+import static com.google.cloud.storage.TestUtils.xxd;
 import static com.google.common.truth.Truth.assertThat;
 import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.junit.Assert.assertThrows;
 
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.cloud.storage.ITObjectReadSessionFakeTest.FakeStorage;
@@ -35,6 +38,7 @@
 import com.google.storage.v2.ObjectRangeData;
 import io.grpc.stub.StreamObserver;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -66,6 +70,8 @@ public final class ITFastOpenFakeTest {
 
   private static final ReadAsFutureBytes READ_10_TO_20 =
       ReadProjectionConfigs.asFutureBytes().withRangeSpec(RangeSpec.of(10, 10));
+  private static final ReadAsFutureBytes READ_20_TO_30 =
+      ReadProjectionConfigs.asFutureBytes().withRangeSpec(RangeSpec.of(20, 10));
 
   /** Every request the fake server received, in order. */
   private final List<BidiReadObjectRequest> requestsReceived =
@@ -104,6 +110,114 @@ public void openRequest_withoutFastOpenReads_isUnchanged() throws Exception {
     assertThat(requestsReceived).containsExactly(OPEN);
   }
 
+  @Test
+  public void readAs_withTheFastOpenConfig_doesNotSendAnotherRequest() throws Exception {
+    FakeStorage server =
+        fakeServer(ImmutableMap.of(OPEN_WITH_READ_OF_10_TO_20, metadataAndBytes(1, 10, 10)));
+
+    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)) {
+        byte[] actual = session.readAs(READ_10_TO_20).get(2, SECONDS);
+        assertThat(xxd(actual)).isEqualTo(xxd(slice(10, 10)));
+      }
+    }
+
+    assertThat(requestsReceived).containsExactly(OPEN_WITH_READ_OF_10_TO_20);
+  }
+
+  @Test
+  public void readAs_withTheFastOpenConfigAgain_sendsANewRead() 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)));
+
+    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)) {
+        byte[] fastOpen = session.readAs(READ_10_TO_20).get(2, SECONDS);
+        byte[] second = session.readAs(READ_10_TO_20).get(2, SECONDS);
+        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 readAs_withADifferentConfig_sendsANewRead_andKeepsTheFastOpenRead() throws Exception {
+    BidiReadObjectRequest readOf20To30 = read(2, 20, 10);
+    FakeStorage server =
+        fakeServer(
+            ImmutableMap.of(
+                OPEN_WITH_READ_OF_10_TO_20, metadataAndBytes(1, 10, 10),
+                readOf20To30, bytes(2, 20, 10)));
+
+    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)) {
+        byte[] other = session.readAs(READ_20_TO_30).get(2, SECONDS);
+        byte[] fastOpen = session.readAs(READ_10_TO_20).get(2, SECONDS);
+        assertThat(xxd(other)).isEqualTo(xxd(slice(20, 10)));
+        assertThat(xxd(fastOpen)).isEqualTo(xxd(slice(10, 10)));
+      }
+    }
+
+    assertThat(requestsReceived)
+        .containsExactly(OPEN_WITH_READ_OF_10_TO_20, readOf20To30)
+        .inOrder();
+  }
+
+  @Test
+  public void close_withAnUnclaimedFastOpenRead_closesTheSession() throws Exception {
+    FakeStorage server =
+        fakeServer(ImmutableMap.of(OPEN_WITH_READ_OF_10_TO_20, metadataAndBytes(1, 10, 10)));
+
+    try (FakeServer fakeServer = FakeServer.of(server);
+        Storage storage = fakeServer.getGrpcStorageOptions().getService()) {
+      StorageDataClient dataClient = maybeGetStorageDataClient(storage);
+
+      ObjectReadSession session =
+          dataClient.readSession(OPEN, CTX, ImmutableList.of(READ_10_TO_20)).get(5, SECONDS);
+      session.close();
+
+      assertThrows(IllegalStateException.class, () -> session.readAs(READ_10_TO_20));
+    }
+
+    assertThat(requestsReceived).containsExactly(OPEN_WITH_READ_OF_10_TO_20);
+  }
+
+  @Test
+  public void readSession_withMoreThanOneFastOpenConfig_isRejected() throws Exception {
+    FakeStorage server = fakeServer(ImmutableMap.of());
+
+    try (FakeServer fakeServer = FakeServer.of(server);
+        Storage storage = fakeServer.getGrpcStorageOptions().getService()) {
+      StorageDataClient dataClient = maybeGetStorageDataClient(storage);
+
+      assertThrows(
+          IllegalArgumentException.class,
+          () -> dataClient.readSession(OPEN, CTX, ImmutableList.of(READ_10_TO_20, READ_20_TO_30)));
+    }
+
+    assertThat(requestsReceived).isEmpty();
+  }
+
   /**
    * A fake server that answers each request in {@code responses} and adds it to {@link
    * #requestsReceived}. Any other request is answered with an error.
@@ -138,4 +252,13 @@ private static BidiReadObjectResponse metadataAndBytes(int readId, int offset, i
                 .setRangeEnd(true))
         .build();
   }
+
+  /** The response to a read request of {@code [offset, offset + length)}. */
+  private static BidiReadObjectResponse bytes(int readId, int offset, int length) {
+    return metadataAndBytes(readId, offset, length).toBuilder().clearMetadata().build();
+  }
+
+  private static byte[] slice(int offset, int length) {
+    return Arrays.copyOfRange(ALL_OBJECT_BYTES, offset, offset + length);
+  }
 }