Retry request on 408/timeout response from server (#681)

GoogleCloudDataproc · Aug 11, 2022 · ceffb0f · ceffb0f
1 parent 2d9bfc5
commit ceffb0f
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 1 deletion.
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
@@ -422,6 +422,70 @@ public void upload_failure_runtimeException() throws Exception {
         .inOrder();
   }
 
+  @Test
+  public void upload_retry_requestTimeout() throws Exception {
+    byte[] testData = new byte[MediaHttpUploader.MINIMUM_CHUNK_SIZE];
+    new Random().nextBytes(testData);
+
+    MockHttpTransport transport =
+        mockTransport(
+            emptyResponse(HttpStatusCodes.STATUS_CODE_NOT_FOUND),
+            resumableUploadResponse(BUCKET_NAME, OBJECT_NAME),
+            emptyResponse(408), // HTTP 408 Request Timeout
+            jsonDataResponse(
+                newStorageObject(BUCKET_NAME, OBJECT_NAME)
+                    .setSize(BigInteger.valueOf(testData.length))));
+
+    GoogleCloudStorage gcs = mockedGcs(GCS_OPTIONS, transport);
+
+    try (WritableByteChannel writeChannel = gcs.create(RESOURCE_ID)) {
+      writeChannel.write(ByteBuffer.wrap(testData));
+    }
+
+    assertThat(trackingRequestInitializerWithRetries.getAllRequestStrings())
+        .containsExactly(
+            getRequestString(BUCKET_NAME, OBJECT_NAME),
+            resumableUploadRequestString(
+                BUCKET_NAME, OBJECT_NAME, /* generationId= */ 0, /* replaceGenerationId= */ false),
+            resumableUploadChunkRequestString(BUCKET_NAME, OBJECT_NAME, /* uploadId= */ 1),
+            resumableUploadChunkRequestString(BUCKET_NAME, OBJECT_NAME, /* uploadId= */ 2))
+        .inOrder();
+  }
+
+  @Test
+  public void upload_noRetries_forbidden() throws Exception {
+    byte[] testData = new byte[MediaHttpUploader.MINIMUM_CHUNK_SIZE];
+    new Random().nextBytes(testData);
+
+    MockHttpTransport transport =
+        mockTransport(
+            emptyResponse(HttpStatusCodes.STATUS_CODE_NOT_FOUND),
+            resumableUploadResponse(BUCKET_NAME, OBJECT_NAME),
+            emptyResponse(HttpStatusCodes.STATUS_CODE_FORBIDDEN),
+            jsonDataResponse(
+                newStorageObject(BUCKET_NAME, OBJECT_NAME)
+                    .setSize(BigInteger.valueOf(testData.length))));
+
+    GoogleCloudStorage gcs = mockedGcs(GCS_OPTIONS, transport);
+
+    WritableByteChannel writeChannel = gcs.create(RESOURCE_ID);
+    writeChannel.write(ByteBuffer.wrap(testData));
+
+    IOException thrown = assertThrows(IOException.class, writeChannel::close);
+    assertThat(thrown)
+        .hasCauseThat()
+        .hasMessageThat()
+        .contains(String.valueOf(HttpStatusCodes.STATUS_CODE_FORBIDDEN));
+
+    assertThat(trackingRequestInitializerWithRetries.getAllRequestStrings())
+        .containsExactly(
+            getRequestString(BUCKET_NAME, OBJECT_NAME),
+            resumableUploadRequestString(
+                BUCKET_NAME, OBJECT_NAME, /* generationId= */ 0, /* replaceGenerationId= */ false),
+            resumableUploadChunkRequestString(BUCKET_NAME, OBJECT_NAME, /* uploadId= */ 1))
+        .inOrder();
+  }
+
   @Test
   public void reupload_success_singleWrite_multipleUploadChunks() throws Exception {
     byte[] testData = new byte[2 * MediaHttpUploader.MINIMUM_CHUNK_SIZE];

diff --git a/util/src/main/java/com/google/cloud/hadoop/util/RetryHttpInitializer.java b/util/src/main/java/com/google/cloud/hadoop/util/RetryHttpInitializer.java
@@ -35,6 +35,7 @@
 import com.google.common.flogger.GoogleLogger;
 import com.google.common.flogger.LogContext;
 import java.io.IOException;
+import java.util.Set;
 
 /** An implementation of {@link HttpRequestInitializer} with retries. */
 public class RetryHttpInitializer implements HttpRequestInitializer {
@@ -116,6 +117,15 @@ private static class UnsuccessfulResponseHandler implements HttpUnsuccessfulResp
     /** HTTP status code indicating too many requests in a given amount of time. */
     private static final int HTTP_SC_TOO_MANY_REQUESTS = 429;
 
+    /**
+     * HTTP status code indicating that the server has decided to close the connection rather than
+     * continue waiting
+     */
+    private static final int HTTP_REQUEST_TIMEOUT = 408;
+
+    private static final Set<Integer> RETRYABLE_CODES =
+        ImmutableSet.of(HTTP_SC_TOO_MANY_REQUESTS, HTTP_REQUEST_TIMEOUT);
+
     // The set of response codes to log URLs for with a rate limit.
     private static final ImmutableSet<Integer> RESPONSE_CODES_TO_LOG_WITH_RATE_LIMIT =
         ImmutableSet.of(HTTP_SC_TOO_MANY_REQUESTS);
@@ -132,7 +142,7 @@ private static class UnsuccessfulResponseHandler implements HttpUnsuccessfulResp
     // of the bases cases defined by this instance.
     private static final HttpBackOffUnsuccessfulResponseHandler.BackOffRequired BACK_OFF_REQUIRED =
         response ->
-            response.getStatusCode() == HTTP_SC_TOO_MANY_REQUESTS
+            RETRYABLE_CODES.contains(response.getStatusCode())
                 || HttpBackOffUnsuccessfulResponseHandler.BackOffRequired.ON_SERVER_ERROR
                     .isRequired(response);