2 files changed, 50 insertions, 44 deletions

diff --git a/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererManager.java b/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererManager.java
index b1f133a..680a6da 100644
--- a/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererManager.java
+++ b/common/src/main/java/dev/isxander/yacl3/gui/image/ImageRendererManager.java
@@ -2,75 +2,86 @@ package dev.isxander.yacl3.gui.image;
 
 import com.mojang.blaze3d.systems.RenderSystem;
 import dev.isxander.yacl3.impl.utils.YACLConstants;
+import net.minecraft.client.Minecraft;
 import net.minecraft.resources.ResourceLocation;
 
 import java.util.Map;
-import java.util.Queue;
+import java.util.Optional;
 import java.util.concurrent.*;
+import java.util.function.Supplier;
 
 public class ImageRendererManager {
     private static final ExecutorService SINGLE_THREAD_EXECUTOR = Executors.newSingleThreadExecutor(task -> new Thread(task, "YACL Image Prep"));
 
     private static final Map<ResourceLocation, CompletableFuture<ImageRenderer>> IMAGE_CACHE = new ConcurrentHashMap<>();
-    private static final Queue<FactoryIDPair<?>> CREATION_QUEUE = new ConcurrentLinkedQueue<>();
 
+    @SuppressWarnings("unchecked")
     public static <T extends ImageRenderer> CompletableFuture<T> registerImage(ResourceLocation id, ImageRendererFactory<T> factory) {
-        SINGLE_THREAD_EXECUTOR.submit(() -> {
-            try {
-                ImageRendererFactory.ImageSupplier<T> supplier = factory.prepareImage();
-                CREATION_QUEUE.add(new FactoryIDPair<>(id, supplier));
-            } catch (Exception e) {
-                YACLConstants.LOGGER.error("Failed to prepare image '{}'", id, e);
-                IMAGE_CACHE.remove(id);
-            }
-        });
+        if (IMAGE_CACHE.containsKey(id)) {
+            return (CompletableFuture<T>) IMAGE_CACHE.get(id);
+        }
 
         var future = new CompletableFuture<ImageRenderer>();
         IMAGE_CACHE.put(id, future);
+
+        SINGLE_THREAD_EXECUTOR.submit(() -> {
+            Supplier<Optional<ImageRendererFactory.ImageSupplier<T>>> supplier =
+                    factory.requiresOffThreadPreparation()
+                            ? new CompletedSupplier<>(safelyPrepareFactory(id, factory))
+                            : () -> safelyPrepareFactory(id, factory);
+
+            Minecraft.getInstance().execute(() -> completeImageFactory(id, supplier, future));
+        });
+
         return (CompletableFuture<T>) future;
     }
 
-    public static void pollImageFactories() {
+    private static <T extends ImageRenderer> void completeImageFactory(ResourceLocation id, Supplier<Optional<ImageRendererFactory.ImageSupplier<T>>> supplier, CompletableFuture<ImageRenderer> future) {
         RenderSystem.assertOnRenderThread();
 
-        while (!CREATION_QUEUE.isEmpty()) {
-            FactoryIDPair<?> pair = CREATION_QUEUE.poll();
-
-            // sanity check - this should never happen
-            if (!IMAGE_CACHE.containsKey(pair.id())) {
-                YACLConstants.LOGGER.error("Tried to finalise image '{}' but it was not found in cache.", pair.id());
-                continue;
-            }
-
-            ImageRenderer image;
-            try {
-                image = pair.supplier().completeImage();
-            } catch (Exception e) {
-                YACLConstants.LOGGER.error("Failed to create image '{}'", pair.id(), e);
-                continue;
-            }
+        ImageRendererFactory.ImageSupplier<T> completableImage = supplier.get().orElse(null);
+        if (completableImage == null) {
+            return;
+        }
 
-            CompletableFuture<ImageRenderer> future = IMAGE_CACHE.get(pair.id());
-            // another sanity check - this should never happen
-            if (future.isDone()) {
-                YACLConstants.LOGGER.error("Image '{}' was already completed", pair.id());
-                continue;
-            }
+        // sanity check - this should never happen
+        if (future.isDone()) {
+            YACLConstants.LOGGER.error("Image '{}' was already completed", id);
+            return;
+        }
 
-            future.complete(image);
+        ImageRenderer image;
+        try {
+            image = completableImage.completeImage();
+        } catch (Exception e) {
+            YACLConstants.LOGGER.error("Failed to create image '{}'", id, e);
+            return;
         }
+
+        future.complete(image);
     }
 
     public static void closeAll() {
         SINGLE_THREAD_EXECUTOR.shutdownNow();
-        CREATION_QUEUE.clear();
-        IMAGE_CACHE.values().forEach(future -> {
+        IMAGE_CACHE.values().removeIf(future -> {
             if (future.isDone()) {
                 future.join().close();
             }
+            return true;
         });
-        IMAGE_CACHE.clear();
     }
 
-    private record FactoryIDPair<T extends ImageRenderer>(ResourceLocation id, ImageRendererFactory.ImageSupplier<T> supplier) {}
+    private static <T extends ImageRenderer> Optional<ImageRendererFactory.ImageSupplier<T>> safelyPrepareFactory(ResourceLocation id, ImageRendererFactory<T> factory) {
+        try {
+            return Optional.of(factory.prepareImage());
+        } catch (Exception e) {
+            YACLConstants.LOGGER.error("Failed to prepare image '{}'", id, e);
+            IMAGE_CACHE.remove(id);
+            return Optional.empty();
+        }
+    }
+
+    private record CompletedSupplier<T>(T get) implements Supplier<T> {
+    }
+
 }
diff --git a/common/src/main/java/dev/isxander/yacl3/mixin/MinecraftMixin.java b/common/src/main/java/dev/isxander/yacl3/mixin/MinecraftMixin.java
index 0b228a1..45bc314 100644
--- a/common/src/main/java/dev/isxander/yacl3/mixin/MinecraftMixin.java
+++ b/common/src/main/java/dev/isxander/yacl3/mixin/MinecraftMixin.java
@@ -13,9 +13,4 @@ public class MinecraftMixin {
     private void closeImages(CallbackInfo ci) {
         ImageRendererManager.closeAll();
     }
-
-    @Inject(method = "runTick", at = @At(value = "HEAD"))
-    private void finaliseImages(boolean tick, CallbackInfo ci) {
-        ImageRendererManager.pollImageFactories();
-    }
 }