Retry network operations with registries on I/O errors

jib-core/src/main/java/com/google/cloud/tools/jib/blob/Blob.java

@@ -30,4 +30,13 @@ public interface Blob {
    * @throws IOException if writing the BLOB fails
    */
   BlobDescriptor writeTo(OutputStream outputStream) throws IOException;
+
+  /**
+   * Enables to notify if the underlying request can be retried (useful in the context of a
+   * retryable HTTP request for ex).
+   *
+   * @return {@code true} if {@link #writeTo(OutputStream)} can be called multiple times, {@code
+   *     false} otherwise.
+   */
+  boolean isRetryable();
 }

jib-core/src/main/java/com/google/cloud/tools/jib/blob/Blobs.java

@@ -50,7 +50,11 @@ public static Blob from(String content) {
   }
 
   public static Blob from(WritableContents writable) {
-    return new WritableContentsBlob(writable);
+    return from(writable, false);
+  }
+
+  public static Blob from(WritableContents writable, boolean retryable) {
+    return new WritableContentsBlob(writable, retryable);
   }
 
   /**

jib-core/src/main/java/com/google/cloud/tools/jib/blob/FileBlob.java

@@ -39,4 +39,9 @@ public BlobDescriptor writeTo(OutputStream outputStream) throws IOException {
       return Digests.computeDigest(fileIn, outputStream);
     }
   }
+
+  @Override
+  public boolean isRetryable() {
+    return true;
+  }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/blob/InputStreamBlob.java

@@ -46,4 +46,9 @@ public BlobDescriptor writeTo(OutputStream outputStream) throws IOException {
       isWritten = true;
     }
   }
+
+  @Override
+  public boolean isRetryable() {
+    return false;
+  }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/blob/JsonBlob.java

@@ -34,4 +34,9 @@ class JsonBlob implements Blob {
   public BlobDescriptor writeTo(OutputStream outputStream) throws IOException {
     return Digests.computeDigest(template, outputStream);
   }
+
+  @Override
+  public boolean isRetryable() {
+    return true;
+  }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/blob/StringBlob.java

@@ -39,4 +39,9 @@ public BlobDescriptor writeTo(OutputStream outputStream) throws IOException {
       return Digests.computeDigest(stringIn, outputStream);
     }
   }
+
+  @Override
+  public boolean isRetryable() {
+    return true;
+  }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/blob/WritableContentsBlob.java

@@ -25,13 +25,20 @@
 class WritableContentsBlob implements Blob {
 
   private final WritableContents writableContents;
+  private final boolean retryable;
 
-  WritableContentsBlob(WritableContents writableContents) {
+  WritableContentsBlob(WritableContents writableContents, boolean retryable) {
     this.writableContents = writableContents;
+    this.retryable = retryable;
   }
 
   @Override
   public BlobDescriptor writeTo(OutputStream outputStream) throws IOException {
     return Digests.computeDigest(writableContents, outputStream);
   }
+
+  @Override
+  public boolean isRetryable() {
+    return retryable;
+  }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/builder/steps/LocalBaseImageSteps.java

@@ -305,7 +305,8 @@ private static PreparedLayer compressAndCacheTarLayer(
                         new NotifyingOutputStream(compressorStream, throttledProgressReporter)) {
                   Blobs.from(layerFile).writeTo(notifyingOutputStream);
                 }
-              });
+              },
+              true);
       return new PreparedLayer.Builder(cache.writeTarLayer(diffId, compressedBlob)).build();
     }
   }

jib-core/src/main/java/com/google/cloud/tools/jib/http/BlobHttpContent.java

@@ -59,7 +59,7 @@ public String getType() {
 
   @Override
   public boolean retrySupported() {
-    return false;
+    return blob.isRetryable();
   }
 
   @Override

jib-core/src/main/java/com/google/cloud/tools/jib/http/FailoverHttpClient.java

@@ -17,12 +17,14 @@
 package com.google.cloud.tools.jib.http;
 
 import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpBackOffIOExceptionHandler;
 import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.http.HttpMethods;
 import com.google.api.client.http.HttpRequest;
 import com.google.api.client.http.HttpResponseException;
 import com.google.api.client.http.HttpTransport;
 import com.google.api.client.http.apache.v2.ApacheHttpTransport;
+import com.google.api.client.util.ExponentialBackOff;
 import com.google.api.client.util.SslUtils;
 import com.google.cloud.tools.jib.api.LogEvent;
 import com.google.common.annotations.VisibleForTesting;
@@ -120,6 +122,7 @@ private static HttpTransport getInsecureHttpTransport() {
     }
   }
 
+  private final boolean enableRetries;
   private final boolean enableHttpAndInsecureFailover;
   private final boolean sendAuthorizationOverHttp;
   private final Consumer<LogEvent> logger;
@@ -143,6 +146,7 @@ public FailoverHttpClient(
       boolean sendAuthorizationOverHttp,
       Consumer<LogEvent> logger) {
     this(
+        true,
         enableHttpAndInsecureFailover,
         sendAuthorizationOverHttp,
         logger,
@@ -152,11 +156,28 @@ public FailoverHttpClient(
 
   @VisibleForTesting
   FailoverHttpClient(
+      boolean enableRetries,
+      boolean enableHttpAndInsecureFailover,
+      boolean sendAuthorizationOverHttp,
+      Consumer<LogEvent> logger) {
+    this(
+        enableRetries,
+        enableHttpAndInsecureFailover,
+        sendAuthorizationOverHttp,
+        logger,
+        FailoverHttpClient::getSecureHttpTransport,
+        FailoverHttpClient::getInsecureHttpTransport);
+  }
+
+  @VisibleForTesting
+  FailoverHttpClient(
+      boolean enableRetries,
       boolean enableHttpAndInsecureFailover,
       boolean sendAuthorizationOverHttp,
       Consumer<LogEvent> logger,
       Supplier<HttpTransport> secureHttpTransportFactory,
       Supplier<HttpTransport> insecureHttpTransportFactory) {
+    this.enableRetries = enableRetries;
     this.enableHttpAndInsecureFailover = enableHttpAndInsecureFailover;
     this.sendAuthorizationOverHttp = sendAuthorizationOverHttp;
     this.logger = logger;
@@ -312,6 +333,23 @@ private Response call(String httpMethod, URL url, Request request, HttpTransport
         httpTransport
             .createRequestFactory()
             .buildRequest(httpMethod, new GenericUrl(url), request.getHttpContent())
+            .setIOExceptionHandler(
+                new HttpBackOffIOExceptionHandler(new ExponentialBackOff()) {
+                  @Override
+                  public boolean handleIOException(HttpRequest request, boolean supportsRetry)
+                      throws IOException {
+                    boolean result =
+                        enableRetries && super.handleIOException(request, supportsRetry);
+                    String requestUrl = request.getRequestMethod() + " " + request.getUrl();
+                    if (result) { // google-http-client does not log that properly so let's
+                      // compensate it
+                      logger.accept(LogEvent.warn(requestUrl + " failed and will be retried"));
+                    } else {
+                      logger.accept(LogEvent.warn(requestUrl + " failed and will NOT be retried"));
+                    }
+                    return result;
+                  }
+                })
             .setUseRawRedirectUrls(true)
             .setHeaders(requestHeaders);
     if (request.getHttpTimeout() != null) {

jib-core/src/main/java/com/google/cloud/tools/jib/image/ReproducibleLayerBuilder.java

@@ -176,6 +176,6 @@ public Blob build() throws IOException {
       tarStreamBuilder.addTarArchiveEntry(entry);
     }
 
-    return Blobs.from(tarStreamBuilder::writeAsTarArchiveTo);
+    return Blobs.from(tarStreamBuilder::writeAsTarArchiveTo, false);
   }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/registry/RegistryClient.java

@@ -504,7 +504,8 @@ public Blob pullBlob(
           } catch (RegistryException ex) {
             throw new IOException(ex);
           }
-        });
+        },
+        false);
   }
 
   /**

jib-core/src/main/java/com/google/cloud/tools/jib/tar/TarStreamBuilder.java

@@ -62,7 +62,8 @@ public void writeAsTarArchiveTo(OutputStream out) throws IOException {
    * @param entry the {@link TarArchiveEntry}
    */
   public void addTarArchiveEntry(TarArchiveEntry entry) {
-    archiveMap.put(entry, entry.isFile() ? Blobs.from(entry.getPath()) : Blobs.from(ignored -> {}));
+    archiveMap.put(
+        entry, entry.isFile() ? Blobs.from(entry.getPath()) : Blobs.from(ignored -> {}, true));
   }
 
   /**
@@ -77,7 +78,7 @@ public void addByteEntry(byte[] contents, String name, Instant modificationTime)
     TarArchiveEntry entry = new TarArchiveEntry(name);
     entry.setSize(contents.length);
     entry.setModTime(modificationTime.toEpochMilli());
-    archiveMap.put(entry, Blobs.from(outputStream -> outputStream.write(contents)));
+    archiveMap.put(entry, Blobs.from(outputStream -> outputStream.write(contents), true));
   }
 
   /**

jib-core/src/test/java/com/google/cloud/tools/jib/http/FailoverHttpClientTest.java

@@ -16,6 +16,9 @@
 
 package com.google.cloud.tools.jib.http;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
 import com.google.api.client.http.GenericUrl;
 import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.http.HttpMethods;
@@ -25,15 +28,21 @@
 import com.google.api.client.http.HttpTransport;
 import com.google.cloud.tools.jib.api.LogEvent;
 import com.google.cloud.tools.jib.blob.Blobs;
+import com.sun.net.httpserver.HttpHandler;
+import com.sun.net.httpserver.HttpServer;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.net.ConnectException;
+import java.net.InetSocketAddress;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -416,6 +425,60 @@ public void testFollowFailoverHistory_portsDifferent() throws IOException {
         "Cannot verify server at https://url:2. Attempting again with no TLS verification.");
   }
 
+  @Test
+  public void testRetries() throws IOException {
+    final HttpServer server = HttpServer.create(new InetSocketAddress("localhost", 0), 64);
+    final AtomicReference<HttpHandler> current = new AtomicReference<>();
+    server.createContext("/").setHandler(exchange -> current.get().handle(exchange));
+    final AtomicBoolean failed = new AtomicBoolean();
+    final List<LogEvent> events = new ArrayList<>();
+
+    // simulate a failure
+    current.set(
+        ex -> {
+          if (failed.compareAndSet(false, true)) {
+            // simulate a success after this (first) failure
+            current.set(
+                exch -> {
+                  exch.sendResponseHeaders(200, 0);
+                  exch.close();
+                });
+
+            // here is the failure - no response sent (IOException for the client)
+            ex.close();
+            return;
+          }
+          // make the test fail
+          ex.sendResponseHeaders(423, 0);
+          ex.close();
+        });
+
+    try {
+      server.start();
+      assertEquals(
+          200,
+          new FailoverHttpClient(true, true, events::add)
+              .post(
+                  new URL("http://localhost:" + server.getAddress().getPort() + "/test"),
+                  Request.builder()
+                      .setBody(new BlobHttpContent(Blobs.from("foo"), "text/plain"))
+                      .build())
+              .getStatusCode());
+    } finally {
+      server.stop(0);
+    }
+    assertTrue(failed.get());
+    assertEquals(1, events.size());
+
+    final LogEvent warn = events.iterator().next();
+    assertEquals(LogEvent.Level.WARN, warn.getLevel());
+    assertEquals(
+        "POST http://localhost:"
+            + server.getAddress().getPort()
+            + "/test failed and will be retried",
+        warn.getMessage());
+  }
+
   private void setUpMocks(
       HttpTransport mockHttpTransport,
       HttpRequestFactory mockHttpRequestFactory,
@@ -426,6 +489,7 @@ private void setUpMocks(
             mockHttpRequestFactory.buildRequest(Mockito.any(), urlCaptor.capture(), Mockito.any()))
         .thenReturn(mockHttpRequest);
 
+    Mockito.when(mockHttpRequest.setIOExceptionHandler(Mockito.any())).thenReturn(mockHttpRequest);
     Mockito.when(mockHttpRequest.setUseRawRedirectUrls(Mockito.anyBoolean()))
         .thenReturn(mockHttpRequest);
     Mockito.when(mockHttpRequest.setHeaders(httpHeadersCaptor.capture()))
@@ -443,7 +507,12 @@ private FailoverHttpClient newHttpClient(boolean insecure, boolean authOverHttp)
           mockInsecureHttpTransport, mockInsecureHttpRequestFactory, mockInsecureHttpRequest);
     }
     return new FailoverHttpClient(
-        insecure, authOverHttp, logger, () -> mockHttpTransport, () -> mockInsecureHttpTransport);
+        true,
+        insecure,
+        authOverHttp,
+        logger,
+        () -> mockHttpTransport,
+        () -> mockInsecureHttpTransport);
   }
 
   private Request fakeRequest(Integer httpTimeout) {