Fix #918

2 files changed, 45 insertions, 11 deletions

diff --git a/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java b/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java
index fbed13ecf..dab08c264 100644
--- a/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java
+++ b/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java
@@ -32,7 +32,6 @@ import dev.architectury.event.events.client.ClientGuiEvent;
 import dev.architectury.event.events.client.ClientRecipeUpdateEvent;
 import dev.architectury.event.events.client.ClientScreenInputEvent;
 import dev.architectury.networking.NetworkManager;
-import dev.architectury.platform.Platform;
 import me.shedaniel.math.Point;
 import me.shedaniel.rei.api.client.REIRuntime;
 import me.shedaniel.rei.api.client.config.ConfigObject;
@@ -51,7 +50,9 @@ import me.shedaniel.rei.api.client.registry.screen.ClickArea;
 import me.shedaniel.rei.api.client.registry.screen.OverlayDecider;
 import me.shedaniel.rei.api.client.registry.screen.ScreenRegistry;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
+import me.shedaniel.rei.api.common.plugins.PluginManager;
 import me.shedaniel.rei.api.common.plugins.PluginView;
+import me.shedaniel.rei.api.common.plugins.REIPlugin;
 import me.shedaniel.rei.api.common.registry.ReloadStage;
 import me.shedaniel.rei.api.common.util.CollectionUtils;
 import me.shedaniel.rei.api.common.util.EntryStacks;
@@ -106,9 +107,7 @@ import org.jetbrains.annotations.Nullable;
 import java.util.Collection;
 import java.util.ConcurrentModificationException;
 import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
+import java.util.concurrent.*;
 import java.util.function.BiFunction;
 import java.util.function.BooleanSupplier;
 import java.util.function.Function;
@@ -129,6 +128,7 @@ public class RoughlyEnoughItemsCoreClient {
         });
         return thread;
     });
+    private static final List<Future<?>> RELOAD_TASKS = new CopyOnWriteArrayList<>();
     
     public static void attachClientInternals() {
         InternalWidgets.attach();
@@ -318,8 +318,24 @@ public class RoughlyEnoughItemsCoreClient {
             reloadPlugins(startReload, ReloadStage.START);
         });
         ClientRecipeUpdateEvent.EVENT.register(recipeManager -> {
-            if (!Platform.isFabric()) RoughlyEnoughItemsCore.PERFORMANCE_LOGGER.clear();
-            reloadPlugins(endReload, Platform.isFabric() ? ReloadStage.END : null);
+            reloadPlugins(endReload, ReloadStage.END);
+        });
+        ClientGuiEvent.INIT_PRE.register((screen, access) -> {
+            List<ReloadStage> stages = ((PluginManagerImpl<REIPlugin<?>>) PluginManager.getInstance()).getObservedStages();
+            
+            if (Minecraft.getInstance().level != null && Minecraft.getInstance().player != null && stages.contains(ReloadStage.START)
+                && !stages.contains(ReloadStage.END) && !PluginManager.areAnyReloading() && screen instanceof AbstractContainerScreen) {
+                for (Future<?> task : RELOAD_TASKS) {
+                    if (!task.isDone()) {
+                        return EventResult.pass();
+                    }
+                }
+                
+                RoughlyEnoughItemsCore.LOGGER.error("Detected missing stage: END! This is possibly due to issues during client recipe reload! REI will force a reload of the recipes now!");
+                reloadPlugins(endReload, ReloadStage.END);
+            }
+            
+            return EventResult.pass();
         });
         ClientGuiEvent.INIT_POST.register((screen, access) -> {
             REIRuntime.getInstance().getOverlay(false, true);
@@ -452,7 +468,17 @@ public class RoughlyEnoughItemsCoreClient {
             lastReload.setValue(System.currentTimeMillis());
         }
         if (ConfigObject.getInstance().doesRegisterRecipesInAnotherThread()) {
-            CompletableFuture.runAsync(() -> RoughlyEnoughItemsCore._reloadPlugins(start), RELOAD_PLUGINS);
+            Future<?>[] futures = new Future<?>[1];
+            CompletableFuture<Void> future = CompletableFuture.runAsync(() -> RoughlyEnoughItemsCore._reloadPlugins(start), RELOAD_PLUGINS)
+                    .whenComplete((unused, throwable) -> {
+                        // Remove the future from the list of futures
+                        if (futures[0] != null) {
+                            RELOAD_TASKS.remove(futures[0]);
+                            futures[0] = null;
+                        }
+                    });
+            futures[0] = future;
+            RELOAD_TASKS.add(future);
         } else {
             RoughlyEnoughItemsCore._reloadPlugins(start);
         }
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/common/plugins/PluginManagerImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/common/plugins/PluginManagerImpl.java
index 3fa5f150d..8dbeab494 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/common/plugins/PluginManagerImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/common/plugins/PluginManagerImpl.java
@@ -60,7 +60,9 @@ public class PluginManagerImpl<P extends REIPlugin<?>> implements PluginManager<
     private final Map<Class<? extends Reloadable<P>>, Reloadable<? super P>> cache = new ConcurrentHashMap<>();
     private final Class<P> pluginClass;
     private final UnaryOperator<PluginView<P>> view;
-    private boolean reloading = false;
+    @Nullable
+    private ReloadStage reloading = null;
+    private List<ReloadStage> observedStages = new ArrayList<>();
     private final List<REIPluginProvider<P>> plugins = new ArrayList<>();
     private final LongConsumer reloadDoneListener;
     private final Stopwatch reloadStopwatch = Stopwatch.createUnstarted();
@@ -84,7 +86,7 @@ public class PluginManagerImpl<P extends REIPlugin<?>> implements PluginManager<
     
     @Override
     public boolean isReloading() {
-        return reloading;
+        return reloading != null;
     }
     
     @Override
@@ -224,6 +226,8 @@ public class PluginManagerImpl<P extends REIPlugin<?>> implements PluginManager<
         this.forcedMainThread = false;
         this.forceMainThreadStopwatch.reset();
         this.reloadStopwatch.reset().start();
+        this.observedStages.clear();
+        this.observedStages.add(stage);
         List<PluginWrapper<P>> plugins = new ArrayList<>(getPluginWrapped().toList());
         plugins.sort(Comparator.comparingDouble(PluginWrapper<P>::getPriority).reversed());
         Collections.reverse(plugins);
@@ -277,7 +281,7 @@ public class PluginManagerImpl<P extends REIPlugin<?>> implements PluginManager<
     public void startReload(ReloadStage stage) {
         try {
             this.reloadStopwatch.start();
-            reloading = true;
+            reloading = stage;
             
             // Pre Reload
             try (SectionClosable startReloadAll = section(stage, "start-reload/");
@@ -370,7 +374,11 @@ public class PluginManagerImpl<P extends REIPlugin<?>> implements PluginManager<
         } catch (Throwable throwable) {
             throwable.printStackTrace();
         } finally {
-            reloading = false;
+            reloading = null;
         }
     }
+    
+    public List<ReloadStage> getObservedStages() {
+        return observedStages;
+    }
 }