From c06fa216dffcd77627b18bab2e3b155f999355ff Mon Sep 17 00:00:00 2001
From: modmuss50 <modmuss50@gmail.com>
Date: Mon, 25 Jul 2022 20:15:56 +0100
Subject: [PATCH] Add maxRetries to the downloader, should help when
 downloading the assets.

---
 .../loom/task/DownloadAssetsTask.java         |  1 +
 .../loom/util/download/DownloadBuilder.java   | 38 +++++++++++++++++--
 .../unit/download/DownloadStringTest.groovy   | 28 +++++++++++++-
 3 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/src/main/java/net/fabricmc/loom/task/DownloadAssetsTask.java b/src/main/java/net/fabricmc/loom/task/DownloadAssetsTask.java
index 46b4db21..f2131810 100644
--- a/src/main/java/net/fabricmc/loom/task/DownloadAssetsTask.java
+++ b/src/main/java/net/fabricmc/loom/task/DownloadAssetsTask.java
@@ -97,6 +97,7 @@ public abstract class DownloadAssetsTask extends AbstractLoomTask {
 						.download(url)
 						.sha1(sha1)
 						.progress(new GradleDownloadProgressListener(object.name(), progressGroup::createProgressLogger))
+						.maxRetries(3)
 						.downloadPathAsync(getAssetsPath(object, assetIndex), executor);
 			}
 		}
diff --git a/src/main/java/net/fabricmc/loom/util/download/DownloadBuilder.java b/src/main/java/net/fabricmc/loom/util/download/DownloadBuilder.java
index 89751d9d..5a10f413 100644
--- a/src/main/java/net/fabricmc/loom/util/download/DownloadBuilder.java
+++ b/src/main/java/net/fabricmc/loom/util/download/DownloadBuilder.java
@@ -31,6 +31,7 @@ import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.time.Duration;
+import java.util.Locale;
 
 @SuppressWarnings("UnusedReturnValue")
 public class DownloadBuilder {
@@ -43,6 +44,7 @@ public class DownloadBuilder {
 	private boolean offline = false;
 	private Duration maxAge = Duration.ZERO;
 	private DownloadProgressListener progressListener = DownloadProgressListener.NONE;
+	private int maxRetries = 1;
 
 	private DownloadBuilder(URI url) {
 		this.url = url;
@@ -82,6 +84,11 @@ public class DownloadBuilder {
 		return this;
 	}
 
+	public DownloadBuilder maxRetries(int maxRetries) {
+		this.maxRetries = maxRetries;
+		return this;
+	}
+
 	public DownloadBuilder defaultCache() {
 		etag(true);
 		return maxAge(ONE_DAY);
@@ -96,15 +103,21 @@ public class DownloadBuilder {
 	}
 
 	public void downloadPath(Path path) throws DownloadException {
-		build().downloadPath(path);
+		withRetries(() -> {
+			build().downloadPath(path);
+			return null;
+		});
 	}
 
 	public String downloadString() throws DownloadException {
-		return build().downloadString();
+		return withRetries(() -> build().downloadString());
 	}
 
 	public String downloadString(Path cache) throws DownloadException {
-		downloadPath(cache);
+		withRetries(() -> {
+			build().downloadPath(cache);
+			return null;
+		});
 
 		try {
 			return Files.readString(cache, StandardCharsets.UTF_8);
@@ -118,4 +131,23 @@ public class DownloadBuilder {
 			throw new DownloadException("Failed to download and read string", e);
 		}
 	}
+
+	private <T> T withRetries(DownloadSupplier<T> supplier) throws DownloadException {
+		for (int i = 1; i <= maxRetries; i++) {
+			try {
+				return supplier.get();
+			} catch (DownloadException e) {
+				if (i == maxRetries) {
+					throw new DownloadException(String.format(Locale.ENGLISH, "Failed download after %d attempts", maxRetries), e);
+				}
+			}
+		}
+
+		throw new IllegalStateException();
+	}
+
+	@FunctionalInterface
+	private interface DownloadSupplier<T> {
+		T get() throws DownloadException;
+	}
 }
diff --git a/src/test/groovy/net/fabricmc/loom/test/unit/download/DownloadStringTest.groovy b/src/test/groovy/net/fabricmc/loom/test/unit/download/DownloadStringTest.groovy
index 3ec15e35..66f58dd6 100644
--- a/src/test/groovy/net/fabricmc/loom/test/unit/download/DownloadStringTest.groovy
+++ b/src/test/groovy/net/fabricmc/loom/test/unit/download/DownloadStringTest.groovy
@@ -24,6 +24,7 @@
 
 package net.fabricmc.loom.test.unit.download
 
+import io.javalin.http.HttpCode
 import net.fabricmc.loom.util.download.Download
 import net.fabricmc.loom.util.download.DownloadException
 
@@ -48,7 +49,9 @@ class DownloadStringTest extends DownloadTest {
 			}
 
 		when:
-			def result = Download.create("$PATH/stringNotFound").downloadString()
+			def result = Download.create("$PATH/stringNotFound")
+				.maxRetries(3) // Ensure we still error as expected when retrying
+				.downloadString()
 
 		then:
 			thrown DownloadException
@@ -69,4 +72,27 @@ class DownloadStringTest extends DownloadTest {
 		then:
 			result == "Hello World!"
 	}
+
+	def "String: Retries"() {
+		setup:
+			int requests = 0
+			server.get("/retryString") {
+				requests ++
+
+				if (requests < 3) {
+					it.status(HttpCode.INTERNAL_SERVER_ERROR)
+					return
+				}
+
+				it.result("Hello World " + requests)
+			}
+
+		when:
+			def result = Download.create("$PATH/retryString")
+				.maxRetries(3)
+				.downloadString()
+
+		then:
+			result == "Hello World 3"
+	}
 }