From 51af754089396fa655530c391254ca5adaa5cc28 Mon Sep 17 00:00:00 2001
From: Sebastian Hartte <shartte@users.noreply.github.com>
Date: Sun, 26 May 2024 01:10:40 +0200
Subject: [PATCH] Implement retry-support and handling of 429

---
 .../runtime/cli/DownloadAssetsCommand.java    | 10 ++-
 .../runtime/downloads/DownloadManager.java    | 65 +++++++++++++++----
 2 files changed, 62 insertions(+), 13 deletions(-)

diff --git a/src/main/java/net/neoforged/neoform/runtime/cli/DownloadAssetsCommand.java b/src/main/java/net/neoforged/neoform/runtime/cli/DownloadAssetsCommand.java
index ac315e0..a35ad4e 100644
--- a/src/main/java/net/neoforged/neoform/runtime/cli/DownloadAssetsCommand.java
+++ b/src/main/java/net/neoforged/neoform/runtime/cli/DownloadAssetsCommand.java
@@ -44,6 +44,9 @@ public class DownloadAssetsCommand extends NeoFormEngineCommand {
     @CommandLine.Option(names = "--asset-repository")
     public URI assetRepository = URI.create("https://resources.download.minecraft.net/");
 
+    @CommandLine.Option(names = "--concurrent-downloads")
+    public int concurrentDownloads = 25;
+
     /**
      * Properties file that will receive the metadata of the asset index.
      */
@@ -118,8 +121,11 @@ protected void runWithNeoFormEngine(NeoFormEngine engine, List<AutoCloseable> cl
                 .filter(obj -> Files.notExists(objectsFolder.resolve(getObjectPath(obj))))
                 .toList();
 
-        // At most 50 concurrent downloads
-        var semaphore = new Semaphore(50);
+        if (concurrentDownloads < 1) {
+            throw new IllegalStateException("Cannot set concurrent downloads to less than 1: " + concurrentDownloads);
+        }
+
+        var semaphore = new Semaphore(concurrentDownloads);
         try (var executor = Executors.newThreadPerTaskExecutor(DOWNLOAD_THREAD_FACTORY)) {
             for (var object : objectsToDownload) {
                 var spec = new AssetDownloadSpec(object);
diff --git a/src/main/java/net/neoforged/neoform/runtime/downloads/DownloadManager.java b/src/main/java/net/neoforged/neoform/runtime/downloads/DownloadManager.java
index 64ca3f7..9ee7ff6 100644
--- a/src/main/java/net/neoforged/neoform/runtime/downloads/DownloadManager.java
+++ b/src/main/java/net/neoforged/neoform/runtime/downloads/DownloadManager.java
@@ -13,6 +13,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.StandardCopyOption;
+import java.time.Instant;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
@@ -85,18 +86,31 @@ public boolean download(DownloadSpec spec, Path finalLocation, boolean silent) t
                         .header("User-Agent", USER_AGENT)
                         .build();
 
-                HttpResponse<Path> response;
-                try {
-                    response = httpClient.send(request, HttpResponse.BodyHandlers.ofFile(partialFile));
-                } catch (InterruptedException e) {
-                    Thread.currentThread().interrupt();
-                    throw new IOException("Download interrupted", e);
-                }
+                while (true) {
+                    HttpResponse<Path> response;
+                    try {
+                        response = httpClient.send(request, HttpResponse.BodyHandlers.ofFile(partialFile));
+                    } catch (IOException e) {
+                        // We do not have an API to get this information
+                        if ("too many concurrent streams".equals(e.getMessage())) {
+                            waitForRetry(1);
+                            continue;
+                        }
+                        throw e;
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        throw new IOException("Download interrupted", e);
+                    }
 
-                if (response.statusCode() == 404) {
-                    throw new FileNotFoundException(url.toString());
-                } else if (response.statusCode() != 200) {
-                    throw new IOException("Failed to download " + url + ": " + response.statusCode());
+                    if (response.statusCode() == 200) {
+                        break;
+                    } else if (response.statusCode() == 404) {
+                        throw new FileNotFoundException(url.toString());
+                    } else if (canRetryStatusCode(response.statusCode())) {
+                        waitForRetry(response);
+                    } else {
+                        throw new IOException("Failed to download " + url + ": " + response.statusCode());
+                    }
                 }
 
                 // Validate file
@@ -125,4 +139,33 @@ public boolean download(DownloadSpec spec, Path finalLocation, boolean silent) t
         }
         return true;
     }
+
+    private static void waitForRetry(HttpResponse<?> response) throws IOException {
+        // We only support the version of this that specifies the delay in seconds
+        var retryAfter = response.headers().firstValueAsLong("Retry-After").orElse(5);
+        // Clamp some unreasonable delays to 5 minutes
+        waitForRetry(Math.clamp(retryAfter, 0, 300));
+    }
+
+    private static void waitForRetry(int seconds) throws IOException {
+        var waitUntil = Instant.now().plusSeconds(seconds);
+
+        while (Instant.now().isBefore(waitUntil)) {
+            try {
+                Thread.sleep(1000L);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new IOException("Interrupted while waiting for retry.", e);
+            }
+        }
+    }
+
+    private static boolean canRetryStatusCode(int statusCode) {
+        return statusCode == 408 // Request timeout
+               || statusCode == 425 // Too early
+               || statusCode == 429 // Rate-limit exceeded
+               || statusCode == 502
+               || statusCode == 503
+               || statusCode == 504;
+    }
 }