Skip to content

Commit

Permalink
fix: add strict client side response validation for gRPC chunked resu…
Browse files Browse the repository at this point in the history 
…mable uploads

* Rename JsonResumableSessionFailureScenario to ResumableSessionFailureScenario. The failure scenarios themselves are not json specific, and the methods which are json specific can have grpc overloads
* Add more tests to validate GapicUnbufferedChunkedResumableWritableByteChannel is able to properly detect and handle various success responses from GCS which are not success for the client.
* Trim down ResumableWrite a bit as the logic
  • Loading branch information
@BenWhitehead
BenWhitehead committed May 6, 2024
1 parent b84654e commit 374279c
Show file tree
Hide file tree
Showing 11 changed files with 969 additions and 122 deletions.
158 changes: 90 additions & 68 deletions ...ain/java/com/google/cloud/storage/GapicUnbufferedChunkedResumableWritableByteChannel.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import com.google.api.gax.retrying.ResultRetryAlgorithm;
import com.google.api.gax.rpc.ApiStreamObserver;
import com.google.api.gax.rpc.ClientStreamingCallable;
import com.google.api.gax.rpc.OutOfRangeException;
import com.google.cloud.storage.ChunkSegmenter.ChunkSegment;
import com.google.cloud.storage.Conversions.Decoder;
import com.google.cloud.storage.Crc32cValue.Crc32cLengthKnown;
Expand All @@ -41,10 +42,9 @@
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.function.Consumer;
import java.util.function.LongConsumer;
import java.util.function.Supplier;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;

final class GapicUnbufferedChunkedResumableWritableByteChannel
implements UnbufferedWritableByteChannel {
Expand All @@ -58,30 +58,26 @@ final class GapicUnbufferedChunkedResumableWritableByteChannel
private final RetryingDependencies deps;
private final ResultRetryAlgorithm<?> alg;
private final Supplier<GrpcCallContext> baseContextSupplier;
private final LongConsumer sizeCallback;
private final Consumer<WriteObjectResponse> completeCallback;

private boolean open = true;
private volatile boolean open = true;
private boolean finished = false;

GapicUnbufferedChunkedResumableWritableByteChannel(
SettableApiFuture<WriteObjectResponse> resultFuture,
@NonNull ChunkSegmenter chunkSegmenter,
ClientStreamingCallable<WriteObjectRequest, WriteObjectResponse> write,
ResumableWrite requestFactory,
WriteCtx<ResumableWrite> writeCtx,
RetryingDependencies deps,
ResultRetryAlgorithm<?> alg,
Supplier<GrpcCallContext> baseContextSupplier) {
this.resultFuture = resultFuture;
this.chunkSegmenter = chunkSegmenter;
this.write = write;
this.bucketName = requestFactory.bucketName();
this.writeCtx = new WriteCtx<>(requestFactory);
this.bucketName = writeCtx.getRequestFactory().bucketName();
this.writeCtx = writeCtx;
this.deps = deps;
this.alg = alg;
this.baseContextSupplier = baseContextSupplier;
this.sizeCallback = writeCtx.getConfirmedBytes()::set;
this.completeCallback = resultFuture::set;
}

@Override
Expand All @@ -106,7 +102,7 @@ public void close() throws IOException {
if (open && !finished) {
WriteObjectRequest message = finishMessage(true);
try {
flush(ImmutableList.of(message));
flush(ImmutableList.of(message), null, true);
finished = true;
} catch (RuntimeException e) {
resultFuture.setException(e);
Expand All @@ -122,12 +118,13 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
throw new ClosedChannelException();
}

long begin = writeCtx.getConfirmedBytes().get();
RewindableContent content = RewindableContent.of(srcs, srcsOffset, srcsLength);
ChunkSegment[] data = chunkSegmenter.segmentBuffers(srcs, srcsOffset, srcsLength);

List<WriteObjectRequest> messages = new ArrayList<>();

boolean first = true;
int bytesConsumed = 0;
for (ChunkSegment datum : data) {
Crc32cLengthKnown crc32c = datum.getCrc32c();
ByteString b = datum.getB();
Expand All @@ -144,8 +141,13 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
WriteObjectRequest.Builder builder =
writeCtx
.newRequestBuilder()
.clearWriteObjectSpec()
.clearObjectChecksums()
.setWriteOffset(offset)
.setChecksummedData(checksummedData.build());
if (!first) {
builder.clearUploadId();
}
if (!datum.isOnlyFullBlocks()) {
builder.setFinishWrite(true);
if (cumulative != null) {
Expand All @@ -155,23 +157,25 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
finished = true;
}

WriteObjectRequest build = possiblyPairDownRequest(builder, first).build();
WriteObjectRequest build = builder.build();
first = false;
messages.add(build);
bytesConsumed += contentSize;
}
if (finalize && !finished) {
messages.add(finishMessage(first));
finished = true;
}

try {
flush(messages);
flush(messages, content, finalize);
} catch (RuntimeException e) {
resultFuture.setException(e);
throw e;
}

long end = writeCtx.getConfirmedBytes().get();

long bytesConsumed = end - begin;
return bytesConsumed;
}

Expand All @@ -182,14 +186,20 @@ private WriteObjectRequest finishMessage(boolean first) {

WriteObjectRequest.Builder b =
writeCtx.newRequestBuilder().setFinishWrite(true).setWriteOffset(offset);
if (!first) {
b.clearUploadId();
}
if (crc32cValue != null) {
b.setObjectChecksums(ObjectChecksums.newBuilder().setCrc32C(crc32cValue.getValue()).build());
}
WriteObjectRequest message = possiblyPairDownRequest(b, first).build();
WriteObjectRequest message = b.build();
return message;
}

private void flush(@NonNull List<WriteObjectRequest> segments) {
private void flush(
@NonNull List<WriteObjectRequest> segments,
@Nullable RewindableContent content,
boolean finalizing) {
GrpcCallContext internalContext = contextWithBucketName(bucketName, baseContextSupplier.get());
ClientStreamingCallable<WriteObjectRequest, WriteObjectResponse> callable =
write.withDefaultCallContext(internalContext);
Expand All @@ -198,7 +208,7 @@ private void flush(@NonNull List<WriteObjectRequest> segments) {
deps,
alg,
() -> {
Observer observer = new Observer(sizeCallback, completeCallback);
Observer observer = new Observer(content, finalizing);
ApiStreamObserver<WriteObjectRequest> write = callable.clientStreamingCall(observer);

for (WriteObjectRequest message : segments) {
Expand All @@ -211,81 +221,93 @@ private void flush(@NonNull List<WriteObjectRequest> segments) {
Decoder.identity());
}

/**
* Several fields of a WriteObjectRequest are only allowed on the "first" message sent to gcs,
* this utility method centralizes the logic necessary to clear those fields for use by subsequent
* messages.
*/
private static WriteObjectRequest.Builder possiblyPairDownRequest(
WriteObjectRequest.Builder b, boolean firstMessageOfStream) {
if (firstMessageOfStream && b.getWriteOffset() == 0) {
return b;
}

if (!firstMessageOfStream) {
b.clearUploadId();
}

if (b.getWriteOffset() > 0) {
b.clearWriteObjectSpec();
}

if (b.getWriteOffset() > 0 && !b.getFinishWrite()) {
b.clearObjectChecksums();
}
return b;
}

@VisibleForTesting
WriteCtx<?> getWriteCtx() {
return writeCtx;
}

static class Observer implements ApiStreamObserver<WriteObjectResponse> {
class Observer implements ApiStreamObserver<WriteObjectResponse> {

private final LongConsumer sizeCallback;
private final Consumer<WriteObjectResponse> completeCallback;
private final RewindableContent content;
private final boolean finalizing;

private final SettableApiFuture<Void> invocationHandle;
private volatile WriteObjectResponse last;

Observer(LongConsumer sizeCallback, Consumer<WriteObjectResponse> completeCallback) {
this.sizeCallback = sizeCallback;
this.completeCallback = completeCallback;
Observer(@Nullable RewindableContent content, boolean finalizing) {
this.content = content;
this.finalizing = finalizing;
this.invocationHandle = SettableApiFuture.create();
}

@Override
public void onNext(WriteObjectResponse value) {
// incremental update
if (value.hasPersistedSize()) {
sizeCallback.accept(value.getPersistedSize());
} else if (value.hasResource()) {
sizeCallback.accept(value.getResource().getSize());
}
last = value;
}

/**
* observed exceptions so far
*
* <ol>
* <li>{@link com.google.api.gax.rpc.OutOfRangeException}
* <li>{@link com.google.api.gax.rpc.AlreadyExistsException}
* <li>{@link io.grpc.StatusRuntimeException}
* </ol>
*/
@Override
public void onError(Throwable t) {
invocationHandle.setException(t);
if (t instanceof OutOfRangeException) {
OutOfRangeException oore = (OutOfRangeException) t;
open = false;
invocationHandle.setException(
ResumableSessionFailureScenario.SCENARIO_5.toStorageException());
} else {
invocationHandle.setException(t);
}
}

@Override
public void onCompleted() {
if (last != null && last.hasResource()) {
completeCallback.accept(last);
try {
if (last == null) {
throw new StorageException(
0, "onComplete without preceding onNext, unable to determine success.");
} else if (!finalizing && last.hasPersistedSize()) { // incremental
long totalSentBytes = writeCtx.getTotalSentBytes().get();
long persistedSize = last.getPersistedSize();

if (totalSentBytes == persistedSize) {
writeCtx.getConfirmedBytes().set(persistedSize);
} else if (persistedSize < totalSentBytes) {
long delta = totalSentBytes - persistedSize;
// rewind our content and any state that my have run ahead of the actual ack'd bytes
content.rewindTo(delta);
writeCtx.getTotalSentBytes().set(persistedSize);
writeCtx.getConfirmedBytes().set(persistedSize);
} else {
throw ResumableSessionFailureScenario.SCENARIO_7.toStorageException();
}
} else if (finalizing && last.hasResource()) {
long totalSentBytes = writeCtx.getTotalSentBytes().get();
long finalSize = last.getResource().getSize();
if (totalSentBytes == finalSize) {
writeCtx.getConfirmedBytes().set(finalSize);
resultFuture.set(last);
} else if (finalSize < totalSentBytes) {
throw ResumableSessionFailureScenario.SCENARIO_4_1.toStorageException();
} else {
throw ResumableSessionFailureScenario.SCENARIO_4_2.toStorageException();
}
} else if (!finalizing && last.hasResource()) {
throw ResumableSessionFailureScenario.SCENARIO_1.toStorageException();
} else if (finalizing && last.hasPersistedSize()) {
long totalSentBytes = writeCtx.getTotalSentBytes().get();
long persistedSize = last.getPersistedSize();
if (persistedSize < totalSentBytes) {
throw ResumableSessionFailureScenario.SCENARIO_3.toStorageException();
} else {
throw ResumableSessionFailureScenario.SCENARIO_2.toStorageException();
}
} else {
throw ResumableSessionFailureScenario.SCENARIO_0.toStorageException();
}
} catch (Throwable se) {
open = false;
invocationHandle.setException(se);
} finally {
invocationHandle.set(null);
}
invocationHandle.set(null);
}

void await() {
Expand Down
8 changes: 4 additions & 4 deletions ...torage/src/main/java/com/google/cloud/storage/GapicWritableByteChannelSessionBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -299,13 +299,13 @@ UnbufferedWritableByteChannelSession<WriteObjectResponse> build() {
result,
getChunkSegmenter(),
write,
ResumableWrite.identity(start),
new WriteCtx<>(start),
deps,
alg,
Retrying::newCallContext);
} else {
return new GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
result, getChunkSegmenter(), write, ResumableWrite.identity(start));
result, getChunkSegmenter(), write, start);
}
})
.andThen(StorageByteChannels.writable()::createSynchronized));
Expand Down Expand Up @@ -340,13 +340,13 @@ BufferedWritableByteChannelSession<WriteObjectResponse> build() {
result,
getChunkSegmenter(),
write,
ResumableWrite.identity(start),
new WriteCtx<>(start),
deps,
alg,
Retrying::newCallContext);
} else {
return new GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
result, getChunkSegmenter(), write, ResumableWrite.identity(start));
result, getChunkSegmenter(), write, start);
}
})
.andThen(c -> new DefaultBufferedWritableByteChannel(bufferHandle, c))
Expand Down

0 comments on commit 374279c

Please sign in to comment.