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

GoogleCloudDataproc · Aug 11, 2022 · 0e80aed · 0e80aed
1 parent cb6998b
commit 0e80aed
Show file tree

Hide file tree

Showing 2 changed files with 74 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
@@ -418,6 +418,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
@@ -44,6 +44,15 @@ public class RetryHttpInitializer implements HttpRequestInitializer {
   /** 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);
+
   // Base impl of BackOffRequired determining the default set of cases where we'll retry on
   // unsuccessful HTTP responses; we'll mix in additional retryable response cases on top
   // of the bases cases defined by this instance.
@@ -150,7 +159,7 @@ public CredentialOrBackoffResponseHandler() {
       errorCodeHandler.setBackOffRequired(
           response ->
               BASE_HTTP_BACKOFF_REQUIRED.isRequired(response)
-                  || response.getStatusCode() == HTTP_SC_TOO_MANY_REQUESTS);
+                  || RETRYABLE_CODES.contains(response.getStatusCode()));
       if (sleeperOverride != null) {
         errorCodeHandler.setSleeper(sleeperOverride);
       }