Skip to content

Commit

Permalink
fix: update BlobReadChannelV2 handling to correctly restart for decom…
Browse files Browse the repository at this point in the history
…pressed object (#1867)

* fix: update BlobReadChannelV2 handling to correctly restart for decompressed object

When downloading bytes from gcs and decompressing them, a restart of the stream needs to pickup from the offset within the compressed bytes not the decompressed bytes.

Prior to this change http-client was automatically applying gzip decompression to the stream it returns to us thereby causing our tracking to be off.

This change updates our interaction with http client to always request the raw bytes without any transform applied to them, we then at a higher level can handle whether gzip decompression needs to be plumbed in.

Fix up a couple channel closed state subtleties when buffering is used with decompression.

Add a new test leveraging the testbench which forces a broken stream on some compressed object bytes. _NOTE_ This test depends on the next release of testbench and will be failing until we get that release.

* deps(test): update testbench to v0.33.0
  • Loading branch information
BenWhitehead authored Jan 20, 2023
1 parent 2b94567 commit 93e8ed4
Show file tree
Hide file tree
Showing 10 changed files with 150 additions and 42 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package com.google.cloud.storage;

import static com.google.cloud.storage.Utils.ifNonNull;
import static java.util.Objects.requireNonNull;

import com.google.api.client.http.HttpHeaders;
import com.google.api.client.http.HttpResponse;
Expand Down Expand Up @@ -49,6 +50,7 @@
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.ScatteringByteChannel;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.function.Function;
import javax.annotation.concurrent.Immutable;
Expand Down Expand Up @@ -82,13 +84,9 @@ class ApiaryUnbufferedReadableByteChannel implements UnbufferedReadableByteChann
this.options = options;
this.resultRetryAlgorithm = resultRetryAlgorithm;
this.open = true;
this.position =
apiaryReadRequest.getByteRangeSpec() != null
? apiaryReadRequest.getByteRangeSpec().beginOffset()
: 0;
this.position = apiaryReadRequest.getByteRangeSpec().beginOffset();
}

@SuppressWarnings("UnnecessaryContinue")
@Override
public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
do {
Expand All @@ -113,12 +111,10 @@ public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
// if our retry algorithm COULD allow a retry, continue the loop and allow trying to
// open the stream again.
sbc = null;
continue;
} else if (t instanceof IOException) {
IOException ioE = (IOException) t;
if (resultRetryAlgorithm.shouldRetry(StorageException.translate(ioE), null)) {
sbc = null;
continue;
} else {
throw ioE;
}
Expand Down Expand Up @@ -148,11 +144,8 @@ private void setXGoogGeneration(long xGoogGeneration) {

private ScatteringByteChannel open() {
try {
Boolean b =
(Boolean) apiaryReadRequest.options.get(StorageRpc.Option.RETURN_RAW_INPUT_STREAM);
boolean returnRawInputStream = b != null ? b : true;
ApiaryReadRequest request = apiaryReadRequest.withNewBeginOffset(position);
Get get = createGetRequest(request, storage.objects(), xGoogGeneration, returnRawInputStream);
Get get = createGetRequest(request, storage.objects(), xGoogGeneration);

HttpResponse media = get.executeMedia();
InputStream content = media.getContent();
Expand Down Expand Up @@ -215,10 +208,7 @@ private ScatteringByteChannel open() {

@VisibleForTesting
static Get createGetRequest(
ApiaryReadRequest apiaryReadRequest,
Objects objects,
Long xGoogGeneration,
boolean returnRawInputStream)
ApiaryReadRequest apiaryReadRequest, Objects objects, Long xGoogGeneration)
throws IOException {
StorageObject from = apiaryReadRequest.getObject();
Map<StorageRpc.Option, ?> options = apiaryReadRequest.getOptions();
Expand Down Expand Up @@ -262,7 +252,9 @@ static Get createGetRequest(
base64.encode(hashFunction.hashBytes(base64.decode(key)).asBytes()));
});

get.setReturnRawInputStream(returnRawInputStream);
// gzip handling is performed upstream of here. Ensure we always get the raw input stream from
// the request
get.setReturnRawInputStream(true);
String range = apiaryReadRequest.getByteRangeSpec().getHttpRangeHeader();
if (range != null) {
get.getRequestHeaders().setRange(range);
Expand All @@ -288,7 +280,7 @@ private static String getHeaderValue(@NonNull HttpHeaders headers, @NonNull Stri
if (list.isEmpty()) {
return null;
} else {
return list.get(0);
return list.get(0).trim().toLowerCase(Locale.ENGLISH);
}
} else if (o instanceof String) {
return (String) o;
Expand All @@ -303,27 +295,32 @@ private static String getHeaderValue(@NonNull HttpHeaders headers, @NonNull Stri
static final class ApiaryReadRequest implements Serializable {
private static final long serialVersionUID = -4059435314115374448L;
private static final Gson gson = new Gson();
private transient StorageObject object;
private final Map<StorageRpc.Option, ?> options;
private final ByteRangeSpec byteRangeSpec;
@NonNull private transient StorageObject object;
@NonNull private final Map<StorageRpc.Option, ?> options;
@NonNull private final ByteRangeSpec byteRangeSpec;

private volatile String objectJson;

ApiaryReadRequest(
StorageObject object, Map<StorageRpc.Option, ?> options, ByteRangeSpec byteRangeSpec) {
this.object = object;
this.options = options;
this.byteRangeSpec = byteRangeSpec;
@NonNull StorageObject object,
@NonNull Map<StorageRpc.Option, ?> options,
@NonNull ByteRangeSpec byteRangeSpec) {
this.object = requireNonNull(object, "object must be non null");
this.options = requireNonNull(options, "options must be non null");
this.byteRangeSpec = requireNonNull(byteRangeSpec, "byteRangeSpec must be non null");
}

@NonNull
StorageObject getObject() {
return object;
}

@NonNull
Map<StorageRpc.Option, ?> getOptions() {
return options;
}

@NonNull
ByteRangeSpec getByteRangeSpec() {
return byteRangeSpec;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ public final synchronized int read(ByteBuffer dst) throws IOException {
}

@Override
public ApiFuture<BlobInfo> getObject() {
public final ApiFuture<BlobInfo> getObject() {
return ApiFutures.transform(result, objectDecoder::decode, MoreExecutors.directExecutor());
}

Expand All @@ -136,7 +136,7 @@ protected final int getChunkSize() {
}

@Nullable
protected T getResolvedObject() {
protected final T getResolvedObject() {
if (result.isDone()) {
return StorageException.wrapFutureGet(result);
} else {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ final class BlobReadChannelV2 extends BaseStorageReadChannel<StorageObject> {
private final StorageObject storageObject;
private final Map<StorageRpc.Option, ?> opts;
private final BlobReadChannelContext blobReadChannelContext;
private final boolean autoGzipDecompression;

BlobReadChannelV2(
StorageObject storageObject,
Expand All @@ -41,10 +42,15 @@ final class BlobReadChannelV2 extends BaseStorageReadChannel<StorageObject> {
this.storageObject = storageObject;
this.opts = opts;
this.blobReadChannelContext = blobReadChannelContext;
this.autoGzipDecompression =
// RETURN_RAW_INPUT_STREAM means do not add GZIPInputStream to the pipeline. Meaning, if
// RETURN_RAW_INPUT_STREAM is false, automatically attempt to decompress if Content-Encoding
// gzip.
Boolean.FALSE.equals(opts.get(StorageRpc.Option.RETURN_RAW_INPUT_STREAM));
}

@Override
public RestorableState<ReadChannel> capture() {
public synchronized RestorableState<ReadChannel> capture() {
ApiaryReadRequest apiaryReadRequest = getApiaryReadRequest();
return new BlobReadChannelV2State(
apiaryReadRequest, blobReadChannelContext.getStorageOptions(), getChunkSize());
Expand All @@ -56,6 +62,7 @@ protected LazyReadChannel<StorageObject> newLazyReadChannel() {
ResumableMedia.http()
.read()
.byteChannel(blobReadChannelContext)
.setAutoGzipDecompression(autoGzipDecompression)
.buffered(getBufferHandle())
.setApiaryReadRequest(getApiaryReadRequest())
.build());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ public int read(ByteBuffer dst) throws IOException {
if (retEOF) {
retEOF = false;
return -1;
} else if (!channel.isOpen()) {
} else if (!enqueuedBytes() && !channel.isOpen()) {
throw new ClosedChannelException();
}

Expand Down Expand Up @@ -133,7 +133,7 @@ public int read(ByteBuffer dst) throws IOException {

@Override
public boolean isOpen() {
return !retEOF && channel.isOpen();
return enqueuedBytes() || (!retEOF && channel.isOpen());
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import com.google.api.gax.rpc.ServerStreamingCallable;
import com.google.cloud.storage.BufferedReadableByteChannelSession.BufferedReadableByteChannel;
import com.google.cloud.storage.UnbufferedReadableByteChannelSession.UnbufferedReadableByteChannel;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.storage.v2.Object;
import com.google.storage.v2.ReadObjectRequest;
import com.google.storage.v2.ReadObjectResponse;
Expand Down Expand Up @@ -99,7 +100,9 @@ public UnbufferedReadableByteChannelSessionBuilder unbuffered() {
return (object, resultFuture) -> {
if (autoGzipDecompression) {
return new GzipReadableByteChannel(
new GapicUnbufferedReadableByteChannel(resultFuture, read, object, hasher));
new GapicUnbufferedReadableByteChannel(resultFuture, read, object, hasher),
ApiFutures.transform(
resultFuture, Object::getContentEncoding, MoreExecutors.directExecutor()));
} else {
return new GapicUnbufferedReadableByteChannel(resultFuture, read, object, hasher);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,13 @@

package com.google.cloud.storage;

import com.google.api.core.ApiFuture;
import com.google.cloud.storage.UnbufferedReadableByteChannelSession.UnbufferedReadableByteChannel;
import com.google.storage.v2.Object;
import java.io.*;
import java.io.ByteArrayInputStream;
import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.SequenceInputStream;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.ReadableByteChannel;
Expand All @@ -27,13 +31,15 @@
import java.util.zip.GZIPInputStream;

final class GzipReadableByteChannel implements UnbufferedReadableByteChannel {
private final GapicUnbufferedReadableByteChannel source;
private final UnbufferedReadableByteChannel source;
private final ApiFuture<String> contentEncoding;

private boolean retEOF = false;
private ScatteringByteChannel delegate;

GzipReadableByteChannel(GapicUnbufferedReadableByteChannel source) {
GzipReadableByteChannel(UnbufferedReadableByteChannel source, ApiFuture<String> contentEncoding) {
this.source = source;
this.contentEncoding = contentEncoding;
}

@Override
Expand All @@ -54,10 +60,10 @@ public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
source.read(wrap);
try {
// Step 2: wait for the object metadata, this is populated in the first message from GCS
Object object = source.getResult().get();
String contentEncoding = this.contentEncoding.get();
// if the Content-Encoding is gzip, Step 3: wire gzip decompression into the byte path
// this will have a copy impact as we are no longer controlling all the buffers
if ("gzip".equals(object.getContentEncoding())) {
if ("gzip".equals(contentEncoding) || "x-gzip".equals(contentEncoding)) {
// to wire gzip decompression into the byte path:
// Create an input stream of the first4 bytes we already read
ByteArrayInputStream first4again = new ByteArrayInputStream(first4);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import com.google.cloud.storage.BlobReadChannelV2.BlobReadChannelContext;
import com.google.cloud.storage.BufferedReadableByteChannelSession.BufferedReadableByteChannel;
import com.google.cloud.storage.UnbufferedReadableByteChannelSession.UnbufferedReadableByteChannel;
import com.google.common.util.concurrent.MoreExecutors;
import java.nio.ByteBuffer;
import java.util.function.BiFunction;
import javax.annotation.concurrent.Immutable;
Expand All @@ -51,10 +52,18 @@ public ReadableByteChannelSessionBuilder byteChannel(
public static final class ReadableByteChannelSessionBuilder {

private final BlobReadChannelContext blobReadChannelContext;
private boolean autoGzipDecompression;
// private Hasher hasher; // TODO: wire in Hasher

private ReadableByteChannelSessionBuilder(BlobReadChannelContext blobReadChannelContext) {
this.blobReadChannelContext = blobReadChannelContext;
this.autoGzipDecompression = false;
}

public ReadableByteChannelSessionBuilder setAutoGzipDecompression(
boolean autoGzipDecompression) {
this.autoGzipDecompression = autoGzipDecompression;
return this;
}

public BufferedReadableByteChannelSessionBuilder buffered() {
Expand All @@ -77,13 +86,27 @@ public UnbufferedReadableByteChannelSessionBuilder unbuffered() {
ApiaryReadRequest, SettableApiFuture<StorageObject>, UnbufferedReadableByteChannel>
bindFunction() {
// for any non-final value, create a reference to the value at this point in time
return (request, resultFuture) ->
new ApiaryUnbufferedReadableByteChannel(
boolean autoGzipDecompression = this.autoGzipDecompression;
return (request, resultFuture) -> {
if (autoGzipDecompression) {
return new GzipReadableByteChannel(
new ApiaryUnbufferedReadableByteChannel(
request,
blobReadChannelContext.getApiaryClient(),
resultFuture,
blobReadChannelContext.getStorageOptions(),
blobReadChannelContext.getRetryAlgorithmManager().idempotent()),
ApiFutures.transform(
resultFuture, StorageObject::getContentEncoding, MoreExecutors.directExecutor()));
} else {
return new ApiaryUnbufferedReadableByteChannel(
request,
blobReadChannelContext.getApiaryClient(),
resultFuture,
blobReadChannelContext.getStorageOptions(),
blobReadChannelContext.getRetryAlgorithmManager().idempotent());
}
};
}

public static final class BufferedReadableByteChannelSessionBuilder {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -427,7 +427,7 @@ public void limitAfterReadWorks() throws IOException {
String xxd1 = xxd(bytes1);
assertThat(xxd1).isEqualTo(xxdExpected1);

// seek forward to a new offset
// change the limit
reader.limit(10);

// read again
Expand Down
Loading

0 comments on commit 93e8ed4

Please sign in to comment.