From 2dd8fb23482a4fb4b63987e54b23e64d27c88e4d Mon Sep 17 00:00:00 2001
From: Luke Bemish <lukebemish@lukebemish.dev>
Date: Thu, 2 Jan 2025 16:18:53 -0600
Subject: [PATCH] Fix asset downloads to not try to download nearly as many
 things at once

---
 .../runtime/util/AssetsUtils.java             | 72 ++++++++++---------
 1 file changed, 40 insertions(+), 32 deletions(-)

diff --git a/src/main/java/dev/lukebemish/taskgraphrunner/runtime/util/AssetsUtils.java b/src/main/java/dev/lukebemish/taskgraphrunner/runtime/util/AssetsUtils.java
index b9f8f3a..37ade2e 100644
--- a/src/main/java/dev/lukebemish/taskgraphrunner/runtime/util/AssetsUtils.java
+++ b/src/main/java/dev/lukebemish/taskgraphrunner/runtime/util/AssetsUtils.java
@@ -6,7 +6,6 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
-import java.io.UncheckedIOException;
 import java.net.URI;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -17,6 +16,7 @@
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.concurrent.Semaphore;
 
 public final class AssetsUtils {
     private static final String INDEX_FOLDER = "indexes";
@@ -24,6 +24,8 @@ public final class AssetsUtils {
     private static final String ASSETS_BASE_URL = "https://resources.download.minecraft.net/";
     private static final Logger LOGGER = LoggerFactory.getLogger(AssetsUtils.class);
 
+    private static final String PARALLEL_DOWNLOADS_PROPERTY = "dev.lukebemish.taskgraphrunner.assets.parallel-downloads";
+
     // Sort from most to least indexes
     private static final Comparator<Target> ASSET_INDEX_COUNT_DESCENDING = Comparator.<Target>comparingInt(d -> d.indexes.size()).reversed();
 
@@ -80,40 +82,46 @@ public static Path findOrDownloadIndexAndAssets(DownloadUtils.Spec spec, String
                 try (var reader = Files.newBufferedReader(targetPath)) {
                     json = JsonUtils.GSON.fromJson(reader, JsonObject.class);
                 }
-            }
 
-            var objectsPath = assetOptions.assetRoot().resolve(OBJECT_FOLDER);
-            var objects = json.getAsJsonObject("objects");
-            var targets = objects.asMap().values().stream()
-                .distinct() // The same object can be referenced multiple times
-                .map(entry -> {
-                    var obj = entry.getAsJsonObject();
-                    var hash = obj.getAsJsonPrimitive("hash").getAsString();
-                    var size = obj.getAsJsonPrimitive("size").getAsLong();
-                    var objectPath = objectsPath.resolve(hash.substring(0, 2)).resolve(hash);
-                    var url = URI.create(ASSETS_BASE_URL + hash.substring(0, 2) + "/" + hash);
-                    var objectSpec = new DownloadUtils.Spec.ChecksumAndSize(url, hash, "SHA-1", size);
-                    return new DownloadTarget(hash, objectSpec, objectPath);
-                })
-                .toList();
-
-            try (var ignored = context.lockManager().locks(targets.stream().map(t -> "assets."+t.target().getFileName().toString()).toList())) {
-                var futures = new ArrayList<Future<?>>();
-                for (var target : targets) {
-                    futures.add(context.submit(() -> {
+                var objectsPath = assetOptions.assetRoot().resolve(OBJECT_FOLDER);
+                var objects = json.getAsJsonObject("objects");
+                var targets = objects.asMap().values().stream()
+                    .distinct() // The same object can be referenced multiple times
+                    .map(entry -> {
+                        var obj = entry.getAsJsonObject();
+                        var hash = obj.getAsJsonPrimitive("hash").getAsString();
+                        var size = obj.getAsJsonPrimitive("size").getAsLong();
+                        var objectPath = objectsPath.resolve(hash.substring(0, 2)).resolve(hash);
+                        var url = URI.create(ASSETS_BASE_URL + hash.substring(0, 2) + "/" + hash);
+                        var objectSpec = new DownloadUtils.Spec.ChecksumAndSize(url, hash, "SHA-1", size);
+                        return new DownloadTarget(hash, objectSpec, objectPath);
+                    })
+                    .toList();
+
+                try (var ignored1 = context.lockManager().locks(targets.stream().map(t -> "assets." + t.target().getFileName().toString()).toList())) {
+                    var futures = new ArrayList<Future<?>>();
+                    var semaphore = new Semaphore(Integer.getInteger(PARALLEL_DOWNLOADS_PROPERTY, 8));
+                    for (var target : targets) {
+                        futures.add(context.submit(() -> {
+                            try {
+                                semaphore.acquire();
+                                DownloadUtils.download(target.spec(), target.target());
+                            } catch (IOException | InterruptedException e) {
+                                LOGGER.error("Failed to download asset {} from {}", target.name(), target.spec().uri(), e);
+                                throw new RuntimeException(e);
+                            } finally {
+                                semaphore.release();
+                            }
+                        }));
+                    }
+                    for (var future : futures) {
                         try {
-                            DownloadUtils.download(target.spec(), target.target());
-                        } catch (IOException e) {
-                            LOGGER.error("Failed to download asset {}", target.name(), e);
-                            throw new UncheckedIOException(e);
+                            future.get();
+                        } catch (InterruptedException | ExecutionException e) {
+                            // It failed; don't keep the bad index around
+                            Files.deleteIfExists(targetPath);
+                            throw new RuntimeException(e);
                         }
-                    }));
-                }
-                for (var future : futures) {
-                    try {
-                        future.get();
-                    } catch (InterruptedException | ExecutionException e) {
-                        throw new RuntimeException(e);
                     }
                 }
             }