diff --git a/gcsio/src/test/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageTest.java b/gcsio/src/test/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageTest.java index b8a1d33ee7..45a9688b74 100644 --- a/gcsio/src/test/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageTest.java +++ b/gcsio/src/test/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageTest.java @@ -26,7 +26,6 @@ import static org.mockito.Matchers.anyString; import static org.mockito.Matchers.eq; import static org.mockito.Matchers.isNull; -import static org.mockito.Matchers.startsWith; import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; @@ -38,6 +37,7 @@ import com.google.api.client.googleapis.batch.json.JsonBatchCallback; import com.google.api.client.googleapis.json.GoogleJsonError; +import com.google.api.client.googleapis.services.AbstractGoogleClientRequest; import com.google.api.client.http.AbstractInputStreamContent; import com.google.api.client.http.HttpHeaders; import com.google.api.client.http.HttpRequest; @@ -62,6 +62,7 @@ import com.google.api.services.storage.model.Objects; import com.google.api.services.storage.model.StorageObject; import com.google.cloud.hadoop.gcsio.GoogleCloudStorageImpl.BackOffFactory; +import com.google.cloud.hadoop.util.AbstractGoogleAsyncWriteChannel; import com.google.cloud.hadoop.util.ApiErrorExtractor; import com.google.cloud.hadoop.util.ClientRequestHelper; import com.google.common.collect.ImmutableList; @@ -458,7 +459,11 @@ public StorageObject answer(InvocationOnMock unused) throws Throwable { verify(mockStorageObjectsInsert, times(1)).setIfGenerationMatch(eq(0L)); verify(mockStorageObjectsInsert, times(1)).setIfGenerationMatch(eq(1L)); assertEquals(OBJECT_NAME, storageObjectCaptor.getValue().getName()); - verify(mockHeaders, times(0)).set(startsWith("X-Goog-Upload-"), anyInt()); + verify(mockHeaders, times(1)).set( + eq("X-Goog-Upload-Desired-Chunk-Granularity"), + eq(AbstractGoogleAsyncWriteChannel.GCS_UPLOAD_GRANULARITY)); + verify(mockHeaders, times(0)).set(eq("X-Goog-Upload-Max-Raw-Size"), anyInt()); + verify(mockClientRequestHelper).getRequestHeaders(any(AbstractGoogleClientRequest.class)); verify(mockClientRequestHelper).setChunkSize(any(Storage.Objects.Insert.class), anyInt()); verify(mockClientRequestHelper).setDirectUploadEnabled(eq(mockStorageObjectsInsert), eq(true)); verify(mockErrorExtractor).itemNotFound(any(IOException.class)); @@ -555,7 +560,11 @@ public StorageObject answer(InvocationOnMock unused) throws Throwable { verify(mockStorageObjectsInsert, times(1)).setIfGenerationMatch(eq(222L)); verify(mockStorageObjectsInsert, times(1)).setIfGenerationMatch(eq(1L)); assertEquals(OBJECT_NAME, storageObjectCaptor.getValue().getName()); - verify(mockHeaders, times(0)).set(startsWith("X-Goog-Upload-"), anyInt()); + verify(mockHeaders, times(1)).set( + eq("X-Goog-Upload-Desired-Chunk-Granularity"), + eq(AbstractGoogleAsyncWriteChannel.GCS_UPLOAD_GRANULARITY)); + verify(mockHeaders, times(0)).set(eq("X-Goog-Upload-Max-Raw-Size"), anyInt()); + verify(mockClientRequestHelper).getRequestHeaders(any(AbstractGoogleClientRequest.class)); verify(mockClientRequestHelper).setChunkSize(any(Storage.Objects.Insert.class), anyInt()); verify(mockClientRequestHelper).setDirectUploadEnabled(eq(mockStorageObjectsInsert), eq(true)); verify(mockBackOffFactory).newBackOff(); @@ -651,7 +660,11 @@ public void run() { eq(BUCKET_NAME), any(StorageObject.class), any(AbstractInputStreamContent.class)); verify(mockStorageObjectsInsert, times(1)).setName(eq(OBJECT_NAME)); verify(mockStorageObjectsInsert, times(2)).setDisableGZipContent(eq(true)); - verify(mockHeaders, times(0)).set(startsWith("X-Goog-Upload-"), anyInt()); + verify(mockHeaders, times(1)).set( + eq("X-Goog-Upload-Desired-Chunk-Granularity"), + eq(AbstractGoogleAsyncWriteChannel.GCS_UPLOAD_GRANULARITY)); + verify(mockHeaders, times(0)).set(eq("X-Goog-Upload-Max-Raw-Size"), anyInt()); + verify(mockClientRequestHelper).getRequestHeaders(any(AbstractGoogleClientRequest.class)); verify(mockClientRequestHelper).setChunkSize(any(Storage.Objects.Insert.class), anyInt()); verify(mockClientRequestHelper).setDirectUploadEnabled(eq(mockStorageObjectsInsert), eq(true)); verify(mockStorageObjectsInsert, times(1)).setIfGenerationMatch(eq(0L)); @@ -700,7 +713,11 @@ public void testCreateObjectApiIOException() eq(BUCKET_NAME), any(StorageObject.class), any(AbstractInputStreamContent.class)); verify(mockStorageObjectsInsert, times(1)).setName(eq(OBJECT_NAME)); verify(mockStorageObjectsInsert, times(2)).setDisableGZipContent(eq(true)); - verify(mockHeaders, times(0)).set(startsWith("X-Goog-Upload-"), anyInt()); + verify(mockHeaders, times(1)).set( + eq("X-Goog-Upload-Desired-Chunk-Granularity"), + eq(AbstractGoogleAsyncWriteChannel.GCS_UPLOAD_GRANULARITY)); + verify(mockHeaders, times(0)).set(eq("X-Goog-Upload-Max-Raw-Size"), anyInt()); + verify(mockClientRequestHelper).getRequestHeaders(any(AbstractGoogleClientRequest.class)); verify(mockClientRequestHelper).setChunkSize(any(Storage.Objects.Insert.class), anyInt()); verify(mockClientRequestHelper).setDirectUploadEnabled(eq(mockStorageObjectsInsert), eq(true)); verify(mockStorageObjectsInsert, times(1)).setIfGenerationMatch(eq(0L)); @@ -756,7 +773,11 @@ public void testCreateObjectApiRuntimeException() verify(mockStorageObjectsInsert, times(1)).setName(eq(OBJECT_NAME)); verify(mockStorageObjectsInsert, times(2)).setDisableGZipContent(eq(true)); verify(mockStorageObjects, times(1)).get(anyString(), anyString()); - verify(mockHeaders, times(0)).set(startsWith("X-Goog-Upload-"), anyInt()); + verify(mockHeaders, times(1)).set( + eq("X-Goog-Upload-Desired-Chunk-Granularity"), + eq(AbstractGoogleAsyncWriteChannel.GCS_UPLOAD_GRANULARITY)); + verify(mockHeaders, times(0)).set(eq("X-Goog-Upload-Max-Raw-Size"), anyInt()); + verify(mockClientRequestHelper).getRequestHeaders(any(AbstractGoogleClientRequest.class)); verify(mockClientRequestHelper).setChunkSize(any(Storage.Objects.Insert.class), anyInt()); verify(mockClientRequestHelper).setDirectUploadEnabled(eq(mockStorageObjectsInsert), eq(true)); verify(mockStorageObjectsInsert, times(2)).setIfGenerationMatch(anyLong()); @@ -805,7 +826,11 @@ public void testCreateObjectApiError() verify(mockErrorExtractor, times(1)).itemNotFound(any(IOException.class)); verify(mockBackOff, atLeastOnce()).nextBackOffMillis(); verify(mockBackOffFactory, atLeastOnce()).newBackOff(); - verify(mockHeaders, times(0)).set(startsWith("X-Goog-Upload-"), anyInt()); + verify(mockHeaders, times(1)).set( + eq("X-Goog-Upload-Desired-Chunk-Granularity"), + eq(AbstractGoogleAsyncWriteChannel.GCS_UPLOAD_GRANULARITY)); + verify(mockHeaders, times(0)).set(eq("X-Goog-Upload-Max-Raw-Size"), anyInt()); + verify(mockClientRequestHelper).getRequestHeaders(any(AbstractGoogleClientRequest.class)); verify(mockClientRequestHelper).setChunkSize(any(Storage.Objects.Insert.class), anyInt()); verify(mockClientRequestHelper).setDirectUploadEnabled(eq(mockStorageObjectsInsert), eq(true)); verify(mockStorageObjectsInsert, times(2)).execute(); diff --git a/util/src/main/java/com/google/cloud/hadoop/util/AbstractGoogleAsyncWriteChannel.java b/util/src/main/java/com/google/cloud/hadoop/util/AbstractGoogleAsyncWriteChannel.java index 73ea9c7332..dc8239c551 100644 --- a/util/src/main/java/com/google/cloud/hadoop/util/AbstractGoogleAsyncWriteChannel.java +++ b/util/src/main/java/com/google/cloud/hadoop/util/AbstractGoogleAsyncWriteChannel.java @@ -316,11 +316,13 @@ public void initialize() throws IOException { T request = createRequest(objectContentStream); request.setDisableGZipContent(true); + // Set a larger backend upload-chunk granularity than system defaults. + HttpHeaders headers = clientRequestHelper.getRequestHeaders(request); + headers.set("X-Goog-Upload-Desired-Chunk-Granularity", + Math.min(GCS_UPLOAD_GRANULARITY, uploadBufferSize)); + // Legacy check. Will be phased out. if (limitFileSizeTo250Gb) { - HttpHeaders headers = clientRequestHelper.getRequestHeaders(request); - headers.set("X-Goog-Upload-Desired-Chunk-Granularity", - Math.min(GCS_UPLOAD_GRANULARITY, uploadBufferSize)); headers.set("X-Goog-Upload-Max-Raw-Size", UPLOAD_MAX_SIZE); }