Skip to content

Commit

Permalink
fix: update grpc finalize on close resumable uploads to validate ack'…
Browse files Browse the repository at this point in the history 
…d object size (#2572)

* fix: update grpc finalize on close resumable uploads to validate ack'd object size

Follow up to #2527

Add ITGapicUnbufferedFinalizeOnCloseResumableWritableByteChannelTest with failure scenarios

* fix: update grpc finalize on close resumable uploads to validate ack'd object size

Refactoring to clean up some request building lifecycle and variable scope

* fix: update grpc finalize on close resumable uploads to validate ack'd object size

Tests are passing
  • Loading branch information
@BenWhitehead
BenWhitehead authored Jun 13, 2024
1 parent d193243 commit 55a6d15
Show file tree
Hide file tree
Showing 4 changed files with 478 additions and 110 deletions.
246 changes: 138 additions & 108 deletions .../com/google/cloud/storage/GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,13 @@

import com.google.api.core.SettableApiFuture;
import com.google.api.gax.grpc.GrpcCallContext;
import com.google.api.gax.rpc.ApiException;
import com.google.api.gax.rpc.ApiStreamObserver;
import com.google.api.gax.rpc.ClientStreamingCallable;
import com.google.cloud.storage.ChunkSegmenter.ChunkSegment;
import com.google.cloud.storage.Crc32cValue.Crc32cLengthKnown;
import com.google.cloud.storage.UnbufferedWritableByteChannelSession.UnbufferedWritableByteChannel;
import com.google.common.collect.ImmutableList;
import com.google.protobuf.ByteString;
import com.google.storage.v2.ChecksummedData;
import com.google.storage.v2.ObjectChecksums;
Expand All @@ -33,11 +35,7 @@
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.function.Consumer;
import java.util.function.LongConsumer;
import org.checkerframework.checker.nullness.qual.NonNull;

final class GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel
Expand All @@ -55,22 +53,23 @@ final class GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel
private boolean open = true;
private boolean first = true;
private boolean finished = false;
private volatile WriteObjectRequest lastWrittenRequest;

GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
SettableApiFuture<WriteObjectResponse> resultFuture,
ChunkSegmenter chunkSegmenter,
ClientStreamingCallable<WriteObjectRequest, WriteObjectResponse> write,
ResumableWrite requestFactory) {
String bucketName = requestFactory.bucketName();
WriteCtx<ResumableWrite> writeCtx) {
String bucketName = writeCtx.getRequestFactory().bucketName();
this.resultFuture = resultFuture;
this.chunkSegmenter = chunkSegmenter;

GrpcCallContext internalContext =
contextWithBucketName(bucketName, GrpcCallContext.createDefault());
this.write = write.withDefaultCallContext(internalContext);

this.writeCtx = new WriteCtx<>(requestFactory);
this.responseObserver = new Observer(writeCtx.getConfirmedBytes()::set, resultFuture::set);
this.writeCtx = writeCtx;
this.responseObserver = new Observer(internalContext);
}

@Override
Expand All @@ -92,27 +91,24 @@ public boolean isOpen() {

@Override
public void close() throws IOException {
if (!open) {
return;
}
open = false;
ApiStreamObserver<WriteObjectRequest> openedStream = openedStream();
if (!finished) {
WriteObjectRequest message = finishMessage();
try {
try {
if (!finished) {
WriteObjectRequest message = finishMessage();
lastWrittenRequest = message;
openedStream.onNext(message);
openedStream.onCompleted();
finished = true;
} catch (RuntimeException e) {
resultFuture.setException(e);
throw e;
}
} else {
try {
openedStream.onCompleted();
} catch (RuntimeException e) {
resultFuture.setException(e);
throw e;
}
openedStream.onCompleted();
responseObserver.await();
} catch (RuntimeException e) {
resultFuture.setException(e);
throw e;
}
open = false;
responseObserver.await();
}

private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, boolean finalize)
Expand All @@ -122,51 +118,54 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
}

ChunkSegment[] data = chunkSegmenter.segmentBuffers(srcs, srcsOffset, srcsLength);

List<WriteObjectRequest> messages = new ArrayList<>();
if (data.length == 0) {
return 0;
}

ApiStreamObserver<WriteObjectRequest> openedStream = openedStream();
int bytesConsumed = 0;
for (ChunkSegment datum : data) {
Crc32cLengthKnown crc32c = datum.getCrc32c();
ByteString b = datum.getB();
int contentSize = b.size();
long offset = writeCtx.getTotalSentBytes().getAndAdd(contentSize);
Crc32cLengthKnown cumulative =
writeCtx
.getCumulativeCrc32c()
.accumulateAndGet(crc32c, chunkSegmenter.getHasher()::nullSafeConcat);
ChecksummedData.Builder checksummedData = ChecksummedData.newBuilder().setContent(b);
if (crc32c != null) {
checksummedData.setCrc32C(crc32c.getValue());
}
WriteObjectRequest.Builder builder =
writeCtx
.newRequestBuilder()
.setWriteOffset(offset)
.setChecksummedData(checksummedData.build());
if (!datum.isOnlyFullBlocks()) {
builder.setFinishWrite(true);
if (cumulative != null) {
builder.setObjectChecksums(
ObjectChecksums.newBuilder().setCrc32C(cumulative.getValue()).build());
try {
for (int i = 0; i < data.length; i++) {
ChunkSegment datum = data[i];
Crc32cLengthKnown crc32c = datum.getCrc32c();
ByteString b = datum.getB();
int contentSize = b.size();
long offset = writeCtx.getTotalSentBytes().getAndAdd(contentSize);
Crc32cLengthKnown cumulative =
writeCtx
.getCumulativeCrc32c()
.accumulateAndGet(crc32c, chunkSegmenter.getHasher()::nullSafeConcat);
ChecksummedData.Builder checksummedData = ChecksummedData.newBuilder().setContent(b);
if (crc32c != null) {
checksummedData.setCrc32C(crc32c.getValue());
}
WriteObjectRequest.Builder builder = writeCtx.newRequestBuilder();
if (!first) {
builder.clearUploadId();
builder.clearWriteObjectSpec();
builder.clearObjectChecksums();
}
builder.setWriteOffset(offset).setChecksummedData(checksummedData.build());
if (!datum.isOnlyFullBlocks() || (finalize && i + 1 == data.length)) {
builder.setFinishWrite(true);
if (cumulative != null) {
builder.setObjectChecksums(
ObjectChecksums.newBuilder().setCrc32C(cumulative.getValue()).build());
}
finished = true;
}
finished = true;
}

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

try {
for (WriteObjectRequest message : messages) {
openedStream.onNext(message);
WriteObjectRequest build = builder.build();
first = false;
lastWrittenRequest = build;
openedStream.onNext(build);
bytesConsumed += contentSize;
}
if (finalize && !finished) {
WriteObjectRequest finishMessage = finishMessage();
lastWrittenRequest = finishMessage;
openedStream.onNext(finishMessage);
finished = true;
}
} catch (RuntimeException e) {
resultFuture.setException(e);
Expand Down Expand Up @@ -201,73 +200,104 @@ private ApiStreamObserver<WriteObjectRequest> openedStream() {
return stream;
}

/**
* 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 (b.getWriteOffset() > 0) {
b.clearWriteObjectSpec();
}

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

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

private final LongConsumer sizeCallback;
private final Consumer<WriteObjectResponse> completeCallback;
private final GrpcCallContext context;

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

Observer(LongConsumer sizeCallback, Consumer<WriteObjectResponse> completeCallback) {
this.sizeCallback = sizeCallback;
this.completeCallback = completeCallback;
Observer(GrpcCallContext context) {
this.context = context;
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 ApiException) {
// use StorageExceptions logic to translate from ApiException to our status codes ensuring
// things fall in line with our retry handlers.
// This is suboptimal, as it will initialize a second exception, however this is the
// unusual case, and it should not cause a significant overhead given its rarity.
StorageException tmp = StorageException.asStorageException((ApiException) t);
StorageException storageException =
ResumableSessionFailureScenario.toStorageException(
tmp.getCode(),
tmp.getMessage(),
tmp.getReason(),
ImmutableList.of(lastWrittenRequest),
null,
context,
t);
resultFuture.setException(storageException);
invocationHandle.setException(storageException);
} else {
resultFuture.setException(t);
invocationHandle.setException(t);
}
}

@Override
public void onCompleted() {
if (last != null && last.hasResource()) {
completeCallback.accept(last);
boolean finalizing = lastWrittenRequest.getFinishWrite();
if (last == null) {
clientDetectedError(
ResumableSessionFailureScenario.toStorageException(
0,
"onComplete without preceding onNext, unable to determine success.",
"invalid",
ImmutableList.of(lastWrittenRequest),
null,
context,
null));
} else if (last.hasResource() /* && finalizing*/) {
long totalSentBytes = writeCtx.getTotalSentBytes().get();
long finalSize = last.getResource().getSize();
if (totalSentBytes == finalSize) {
ok(finalSize);
} else if (finalSize < totalSentBytes) {
clientDetectedError(
ResumableSessionFailureScenario.SCENARIO_4_1.toStorageException(
ImmutableList.of(lastWrittenRequest), last, context, null));
} else {
clientDetectedError(
ResumableSessionFailureScenario.SCENARIO_4_2.toStorageException(
ImmutableList.of(lastWrittenRequest), last, context, null));
}
} else if (!finalizing || last.hasPersistedSize()) { // unexpected incremental response
clientDetectedError(
ResumableSessionFailureScenario.toStorageException(
0,
"Unexpected incremental response for finalizing request.",
"invalid",
ImmutableList.of(lastWrittenRequest),
last,
context,
null));
} else {
clientDetectedError(
ResumableSessionFailureScenario.SCENARIO_0.toStorageException(
ImmutableList.of(lastWrittenRequest), last, context, null));
}
}

private void ok(long persistedSize) {
writeCtx.getConfirmedBytes().set(persistedSize);
resultFuture.set(last);
invocationHandle.set(null);
}

private void clientDetectedError(StorageException storageException) {
open = false;
resultFuture.setException(storageException);
invocationHandle.setException(storageException);
}

void await() {
try {
invocationHandle.get();
Expand Down
4 changes: 2 additions & 2 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 @@ -305,7 +305,7 @@ UnbufferedWritableByteChannelSession<WriteObjectResponse> build() {
Retrying::newCallContext);
} else {
return new GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
result, getChunkSegmenter(), write, start);
result, getChunkSegmenter(), write, new WriteCtx<>(start));
}
})
.andThen(StorageByteChannels.writable()::createSynchronized));
Expand Down Expand Up @@ -346,7 +346,7 @@ BufferedWritableByteChannelSession<WriteObjectResponse> build() {
Retrying::newCallContext);
} else {
return new GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
result, getChunkSegmenter(), write, start);
result, getChunkSegmenter(), write, new WriteCtx<>(start));
}
})
.andThen(c -> new DefaultBufferedWritableByteChannel(bufferHandle, c))
Expand Down
Loading

0 comments on commit 55a6d15

Please sign in to comment.