Skip to content

Commit ab0228c

Browse files
BenWhiteheadBenWhitehead
authored
fix: correct lastChunk retry logic in BlobWriteChannel (#918)
Add new method StorageRpc#queryResumableUpload which allows getting a shallow StorageObject for a resumable upload session which is complete. Update BlobWriteChannel to use StoageRpc#queryResumableUpload instead of StorageRpc#get when attempting to validate the upload size of an object when it determines the upload is complete and is on the last chunk. If a BlobWriteChannel is opened with a conditional like IfGenerationMatch it is not possible to simply get the object, as the object can drift generationally while the resumable upload is being performed. Related to #839
1 parent 8b05867 commit ab0228c

7 files changed

Lines changed: 204 additions & 25 deletions

File tree

‎google-cloud-storage/clirr-ignored-differences.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,4 +31,9 @@
3131
<method>long getCurrentUploadOffset(java.lang.String)</method>
3232
<differenceType>7012</differenceType>
3333
</difference>
34+
<difference>
35+
<className>com/google/cloud/storage/spi/v1/StorageRpc</className>
36+
<method>com.google.api.services.storage.model.StorageObject queryCompletedResumableUpload(java.lang.String, long)</method>
37+
<differenceType>7012</differenceType>
38+
</difference>
3439
</differences>

‎google-cloud-storage/src/main/java/com/google/cloud/storage/BlobWriteChannel.java‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import com.google.cloud.RetryHelper;
2626
import com.google.cloud.WriteChannel;
2727
import com.google.cloud.storage.spi.v1.StorageRpc;
28-
import com.google.common.collect.Maps;
2928
import java.math.BigInteger;
3029
import java.net.URL;
3130
import java.util.Map;
@@ -78,12 +77,6 @@ private long getRemotePosition() {
7877
return getOptions().getStorageRpcV1().getCurrentUploadOffset(getUploadId());
7978
}
8079

81-
private StorageObject getRemoteStorageObject() {
82-
return getOptions()
83-
.getStorageRpcV1()
84-
.get(getEntity().toPb(), Maps.newEnumMap(StorageRpc.Option.class));
85-
}
86-
8780
private static StorageException unrecoverableState(
8881
String uploadId,
8982
int chunkOffset,
@@ -212,8 +205,12 @@ public void run() {
212205
if (uploadAlreadyComplete && lastChunk) {
213206
// Case 6
214207
// Request object metadata if not available
208+
long totalBytes = getPosition() + length;
215209
if (storageObject == null) {
216-
storageObject = getRemoteStorageObject();
210+
storageObject =
211+
getOptions()
212+
.getStorageRpcV1()
213+
.queryCompletedResumableUpload(getUploadId(), totalBytes);
217214
}
218215
// the following checks are defined here explicitly to provide a more
219216
// informative if either storageObject is unable to be resolved or it's size is
@@ -239,7 +236,7 @@ public void run() {
239236
remotePosition,
240237
lastChunk);
241238
}
242-
if (size.longValue() != getPosition() + length) {
239+
if (size.longValue() != totalBytes) {
243240
throw unrecoverableState(
244241
getUploadId(),
245242
chunkOffset,

‎google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java‎

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -810,6 +810,25 @@ public long getCurrentUploadOffset(String uploadId) {
810810
}
811811
}
812812

813+
@Override
814+
public StorageObject queryCompletedResumableUpload(String uploadId, long totalBytes) {
815+
try {
816+
GenericUrl url = new GenericUrl(uploadId);
817+
HttpRequest req = storage.getRequestFactory().buildPutRequest(url, new EmptyContent());
818+
req.getHeaders().setContentRange(String.format("bytes */%s", totalBytes));
819+
req.setParser(storage.getObjectParser());
820+
HttpResponse response = req.execute();
821+
// If the response is 200
822+
if (response.getStatusCode() == 200) {
823+
return response.parseAs(StorageObject.class);
824+
} else {
825+
throw buildStorageException(response.getStatusCode(), response.getStatusMessage());
826+
}
827+
} catch (IOException ex) {
828+
throw translate(ex);
829+
}
830+
}
831+
813832
@Override
814833
public StorageObject writeWithResponse(
815834
String uploadId,
@@ -875,10 +894,7 @@ public StorageObject writeWithResponse(
875894
if (exception != null) {
876895
throw exception;
877896
}
878-
GoogleJsonError error = new GoogleJsonError();
879-
error.setCode(code);
880-
error.setMessage(message);
881-
throw translate(error);
897+
throw buildStorageException(code, message);
882898
}
883899
} catch (IOException ex) {
884900
span.setStatus(Status.UNKNOWN.withDescription(ex.getMessage()));
@@ -925,10 +941,7 @@ public String open(StorageObject object, Map<Option, ?> options) {
925941
setEncryptionHeaders(requestHeaders, "x-goog-encryption-", options);
926942
HttpResponse response = httpRequest.execute();
927943
if (response.getStatusCode() != 200) {
928-
GoogleJsonError error = new GoogleJsonError();
929-
error.setCode(response.getStatusCode());
930-
error.setMessage(response.getStatusMessage());
931-
throw translate(error);
944+
throw buildStorageException(response.getStatusCode(), response.getStatusMessage());
932945
}
933946
return response.getHeaders().getLocation();
934947
} catch (IOException ex) {
@@ -962,10 +975,7 @@ public String open(String signedURL) {
962975

963976
HttpResponse response = httpRequest.execute();
964977
if (response.getStatusCode() != 201) {
965-
GoogleJsonError error = new GoogleJsonError();
966-
error.setCode(response.getStatusCode());
967-
error.setMessage(response.getStatusMessage());
968-
throw translate(error);
978+
throw buildStorageException(response.getStatusCode(), response.getStatusMessage());
969979
}
970980
return response.getHeaders().getLocation();
971981
} catch (IOException ex) {
@@ -1625,4 +1635,11 @@ public ServiceAccount getServiceAccount(String projectId) {
16251635
span.end(HttpStorageRpcSpans.END_SPAN_OPTIONS);
16261636
}
16271637
}
1638+
1639+
private static StorageException buildStorageException(int statusCode, String statusMessage) {
1640+
GoogleJsonError error = new GoogleJsonError();
1641+
error.setCode(statusCode);
1642+
error.setMessage(statusMessage);
1643+
return translate(error);
1644+
}
16281645
}

‎google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/StorageRpc.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,24 @@ void write(
338338
*/
339339
long getCurrentUploadOffset(String uploadId);
340340

341+
/**
342+
* Attempts to retrieve the StorageObject from a completed resumable upload. When a resumable
343+
* upload completes, the response will be the up-to-date StorageObject metadata. This up-to-date
344+
* metadata can then be used to validate the total size of the object along with new generation
345+
* and other information.
346+
*
347+
* <p>If for any reason, the response to the final PUT to a resumable upload is not received, this
348+
* method can be used to query for the up-to-date StorageObject. If the upload is complete, this
349+
* method can be used to access the StorageObject independently from any other liveness or
350+
* conditional criteria requirements that are otherwise applicable when using {@link
351+
* #get(StorageObject, Map)}.
352+
*
353+
* @param uploadId resumable upload ID URL
354+
* @param totalBytes the total number of bytes that should have been written.
355+
* @throws StorageException if the upload is incomplete or does not exist
356+
*/
357+
StorageObject queryCompletedResumableUpload(String uploadId, long totalBytes);
358+
341359
/**
342360
* Writes the provided bytes to a storage object at the provided location. If {@code last=true}
343361
* returns metadata of the updated object, otherwise returns null.

‎google-cloud-storage/src/main/java/com/google/cloud/storage/testing/StorageRpcTestBase.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,11 @@ public long getCurrentUploadOffset(String uploadId) {
144144
throw new UnsupportedOperationException("Not implemented yet");
145145
}
146146

147+
@Override
148+
public StorageObject queryCompletedResumableUpload(String uploadId, long totalBytes) {
149+
throw new UnsupportedOperationException("Not implemented yet");
150+
}
151+
147152
@Override
148153
public StorageObject writeWithResponse(
149154
String uploadId,

‎google-cloud-storage/src/test/java/com/google/cloud/storage/BlobWriteChannelTest.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,6 @@
4040
import com.google.cloud.storage.spi.StorageRpcFactory;
4141
import com.google.cloud.storage.spi.v1.StorageRpc;
4242
import com.google.common.collect.ImmutableMap;
43-
import com.google.common.collect.Maps;
4443
import java.io.IOException;
4544
import java.math.BigInteger;
4645
import java.net.MalformedURLException;
@@ -334,10 +333,10 @@ public void testWriteWithRetryAndObjectMetadata() throws IOException {
334333
.andThrow(socketClosedException);
335334
expect(storageRpcMock.getCurrentUploadOffset(eq(UPLOAD_ID))).andReturn(-1L);
336335
expect(storageRpcMock.getCurrentUploadOffset(eq(UPLOAD_ID))).andReturn(-1L);
337-
expect(storageRpcMock.get(BLOB_INFO.toPb(), Maps.newEnumMap(StorageRpc.Option.class)))
336+
expect(storageRpcMock.queryCompletedResumableUpload(eq(UPLOAD_ID), eq((long) MIN_CHUNK_SIZE)))
338337
.andThrow(socketClosedException);
339338
expect(storageRpcMock.getCurrentUploadOffset(eq(UPLOAD_ID))).andReturn(-1L);
340-
expect(storageRpcMock.get(BLOB_INFO.toPb(), Maps.newEnumMap(StorageRpc.Option.class)))
339+
expect(storageRpcMock.queryCompletedResumableUpload(eq(UPLOAD_ID), eq((long) MIN_CHUNK_SIZE)))
341340
.andReturn(BLOB_INFO.toPb().setSize(BigInteger.valueOf(MIN_CHUNK_SIZE)));
342341
replay(storageRpcMock);
343342
writer = new BlobWriteChannel(options, BLOB_INFO, EMPTY_RPC_OPTIONS);
@@ -487,7 +486,7 @@ public void testWriteWithLastFlushRetryChunkButCompleted() throws IOException {
487486
eq(true)))
488487
.andThrow(socketClosedException);
489488
expect(storageRpcMock.getCurrentUploadOffset(eq(UPLOAD_ID))).andReturn(-1L);
490-
expect(storageRpcMock.get(BLOB_INFO.toPb(), Maps.newEnumMap(StorageRpc.Option.class)))
489+
expect(storageRpcMock.queryCompletedResumableUpload(eq(UPLOAD_ID), eq((long) MIN_CHUNK_SIZE)))
491490
.andReturn(BLOB_INFO.toPb().setSize(BigInteger.valueOf(MIN_CHUNK_SIZE)));
492491
replay(storageRpcMock);
493492
writer = new BlobWriteChannel(options, BLOB_INFO, EMPTY_RPC_OPTIONS);

‎google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITStorageTest.java‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,9 @@
8282
import com.google.cloud.storage.StorageException;
8383
import com.google.cloud.storage.StorageOptions;
8484
import com.google.cloud.storage.StorageRoles;
85+
import com.google.cloud.storage.spi.StorageRpcFactory;
86+
import com.google.cloud.storage.spi.v1.StorageRpc;
87+
import com.google.cloud.storage.spi.v1.StorageRpc.Option;
8588
import com.google.cloud.storage.testing.RemoteStorageHelper;
8689
import com.google.common.collect.ImmutableList;
8790
import com.google.common.collect.ImmutableMap;
@@ -90,6 +93,8 @@
9093
import com.google.common.collect.Lists;
9194
import com.google.common.io.BaseEncoding;
9295
import com.google.common.io.ByteStreams;
96+
import com.google.common.reflect.AbstractInvocationHandler;
97+
import com.google.common.reflect.Reflection;
9398
import com.google.iam.v1.Binding;
9499
import com.google.iam.v1.IAMPolicyGrpc;
95100
import com.google.iam.v1.SetIamPolicyRequest;
@@ -107,9 +112,11 @@
107112
import java.io.FileOutputStream;
108113
import java.io.IOException;
109114
import java.io.InputStream;
115+
import java.lang.reflect.Method;
110116
import java.net.URL;
111117
import java.net.URLConnection;
112118
import java.nio.ByteBuffer;
119+
import java.nio.charset.StandardCharsets;
113120
import java.nio.file.Files;
114121
import java.nio.file.Path;
115122
import java.security.Key;
@@ -125,6 +132,7 @@
125132
import java.util.Set;
126133
import java.util.concurrent.ExecutionException;
127134
import java.util.concurrent.TimeUnit;
135+
import java.util.concurrent.atomic.AtomicBoolean;
128136
import java.util.logging.Level;
129137
import java.util.logging.Logger;
130138
import java.util.zip.GZIPInputStream;
@@ -139,7 +147,14 @@
139147
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
140148
import org.junit.AfterClass;
141149
import org.junit.BeforeClass;
150+
import org.junit.Rule;
142151
import org.junit.Test;
152+
import org.junit.rules.TestName;
153+
import org.threeten.bp.Clock;
154+
import org.threeten.bp.Instant;
155+
import org.threeten.bp.ZoneId;
156+
import org.threeten.bp.ZoneOffset;
157+
import org.threeten.bp.format.DateTimeFormatter;
143158

144159
public class ITStorageTest {
145160

@@ -200,6 +215,8 @@ public class ITStorageTest {
200215
private static final ImmutableList<LifecycleRule> LIFECYCLE_RULES =
201216
ImmutableList.of(LIFECYCLE_RULE_1, LIFECYCLE_RULE_2);
202217

218+
@Rule public final TestName testName = new TestName();
219+
203220
@BeforeClass
204221
public static void beforeClass() throws IOException {
205222
remoteStorageHelper = RemoteStorageHelper.create();
@@ -3813,4 +3830,125 @@ public void testWriterWithKmsKeyName() throws IOException {
38133830
assertThat(blob.getKmsKeyName()).isNotNull();
38143831
assertThat(storage.delete(BUCKET, blobName)).isTrue();
38153832
}
3833+
3834+
@Test
3835+
public void blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent_multipleChunks()
3836+
throws IOException {
3837+
int _2MiB = 256 * 1024;
3838+
int contentSize = 292_617;
3839+
3840+
blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent(_2MiB, contentSize);
3841+
}
3842+
3843+
@Test
3844+
public void blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent_singleChunk()
3845+
throws IOException {
3846+
int _4MiB = 256 * 1024 * 2;
3847+
int contentSize = 292_617;
3848+
3849+
blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent(_4MiB, contentSize);
3850+
}
3851+
3852+
private void blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent(
3853+
int chunkSize, int contentSize) throws IOException {
3854+
Instant now = Clock.systemUTC().instant();
3855+
DateTimeFormatter formatter =
3856+
DateTimeFormatter.ISO_LOCAL_DATE_TIME.withZone(ZoneId.from(ZoneOffset.UTC));
3857+
String nowString = formatter.format(now);
3858+
3859+
String blobPath = String.format("%s/%s/blob", testName.getMethodName(), nowString);
3860+
BlobId blobId = BlobId.of(BUCKET, blobPath);
3861+
BlobInfo blobInfo = BlobInfo.newBuilder(blobId).build();
3862+
3863+
Random rand = new Random(1234567890);
3864+
String randString = randString(rand, contentSize);
3865+
final byte[] randStringBytes = randString.getBytes(StandardCharsets.UTF_8);
3866+
Storage storage = StorageOptions.getDefaultInstance().getService();
3867+
WriteChannel ww = storage.writer(blobInfo);
3868+
ww.setChunkSize(chunkSize);
3869+
ww.write(ByteBuffer.wrap(randStringBytes));
3870+
ww.close();
3871+
3872+
Blob blobGen1 = storage.get(blobId);
3873+
3874+
final AtomicBoolean exceptionThrown = new AtomicBoolean(false);
3875+
3876+
Storage testStorage =
3877+
StorageOptions.newBuilder()
3878+
.setServiceRpcFactory(
3879+
new StorageRpcFactory() {
3880+
/**
3881+
* Here we're creating a proxy of StorageRpc where we can delegate all calls to
3882+
* the normal implementation, except in the case of {@link
3883+
* StorageRpc#writeWithResponse(String, byte[], int, long, int, boolean)} where
3884+
* {@code lastChunk == true}. We allow the call to execute, but instead of
3885+
* returning the result we throw an IOException to simulate a prematurely close
3886+
* connection. This behavior is to ensure appropriate handling of a completed
3887+
* upload where the ACK wasn't received. In particular, if an upload is initiated
3888+
* against an object where an {@link Option#IF_GENERATION_MATCH} simply calling
3889+
* get on an object can result in a 404 because the object that is created while
3890+
* the BlobWriteChannel is executing will be a new generation.
3891+
*/
3892+
@SuppressWarnings("UnstableApiUsage")
3893+
@Override
3894+
public StorageRpc create(final StorageOptions options) {
3895+
return Reflection.newProxy(
3896+
StorageRpc.class,
3897+
new AbstractInvocationHandler() {
3898+
final StorageRpc delegate =
3899+
(StorageRpc) StorageOptions.getDefaultInstance().getRpc();
3900+
3901+
@Override
3902+
protected Object handleInvocation(
3903+
Object proxy, Method method, Object[] args) throws Throwable {
3904+
if ("writeWithResponse".equals(method.getName())) {
3905+
Object result = method.invoke(delegate, args);
3906+
boolean lastChunk = (boolean) args[5];
3907+
// if we're on the lastChunk simulate a connection failure which
3908+
// happens after the request was processed but before response could
3909+
// be received by the client.
3910+
if (lastChunk) {
3911+
exceptionThrown.set(true);
3912+
throw StorageException.translate(
3913+
new IOException("simulated Connection closed prematurely"));
3914+
} else {
3915+
return result;
3916+
}
3917+
}
3918+
return method.invoke(delegate, args);
3919+
}
3920+
});
3921+
}
3922+
})
3923+
.build()
3924+
.getService();
3925+
3926+
try (WriteChannel w = testStorage.writer(blobGen1, BlobWriteOption.generationMatch())) {
3927+
w.setChunkSize(chunkSize);
3928+
3929+
ByteBuffer buffer = ByteBuffer.wrap(randStringBytes);
3930+
w.write(buffer);
3931+
}
3932+
3933+
assertTrue("Expected an exception to be thrown for the last chunk", exceptionThrown.get());
3934+
3935+
Blob blobGen2 = storage.get(blobId);
3936+
assertEquals(contentSize, (long) blobGen2.getSize());
3937+
assertNotEquals(blobInfo.getGeneration(), blobGen2.getGeneration());
3938+
ByteArrayOutputStream actualData = new ByteArrayOutputStream();
3939+
blobGen2.downloadTo(actualData);
3940+
assertArrayEquals(randStringBytes, actualData.toByteArray());
3941+
}
3942+
3943+
private static String randString(Random rand, int length) {
3944+
final StringBuilder sb = new StringBuilder();
3945+
while (sb.length() < length) {
3946+
int i = rand.nextInt('z');
3947+
char c = (char) i;
3948+
if (Character.isLetter(c) || Character.isDigit(c)) {
3949+
sb.append(c);
3950+
}
3951+
}
3952+
return sb.toString();
3953+
}
38163954
}

0 commit comments

Comments
 (0)