diff --git a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ObjectReadSessionImpl.java b/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ObjectReadSessionImpl.java index 085b63a14196..f33bc9674597 100644 --- a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ObjectReadSessionImpl.java +++ b/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ObjectReadSessionImpl.java @@ -91,12 +91,13 @@ public Projection readAs(ReadProjectionConfig 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(); diff --git a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/OtelStorageDecorator.java b/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/OtelStorageDecorator.java index 097bf7fbbfe8..9095a8a22cd7 100644 --- a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/OtelStorageDecorator.java +++ b/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 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() { + @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 { diff --git a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ReadProjectionConfig.java b/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ReadProjectionConfig.java index 4812adaa0ec1..49f5444ae08e 100644 --- a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ReadProjectionConfig.java +++ b/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/ReadProjectionConfig.java @@ -43,6 +43,14 @@ public abstract class ReadProjectionConfig { abstract ProjectionType getType(); + /** The config this one wraps, or {@code this}. Used to match a fast open read. */ + ReadProjectionConfig 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())); diff --git a/java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/ITFastOpenFakeTest.java b/java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/ITFastOpenFakeTest.java index 349ca1370719..73ac5780bf67 100644 --- a/java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/ITFastOpenFakeTest.java +++ b/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 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> { + private final ReadAsFutureBytes delegate; + private final List> claimed = new ArrayList<>(); + + private WrappedConfig(ReadAsFutureBytes delegate) { + this.delegate = delegate; + } + + @Override + BaseConfig, ?> cast() { + return delegate.cast(); + } + + @Override + ProjectionType getType() { + return delegate.getType(); + } + + @Override + ReadProjectionConfig> unwrap() { + return delegate; + } + + @Override + void onFastOpenReadClaimed(ApiFuture projection) { + claimed.add(projection); + } + } }