Fix #918

4 files changed, 91 insertions, 11 deletions

diff --git a/forge/src/main/java/me/shedaniel/rei/mixin/forge/MixinClientPacketListener.java b/forge/src/main/java/me/shedaniel/rei/mixin/forge/MixinClientPacketListener.java
new file mode 100644
index 000000000..252a365b6
--- /dev/null
+++ b/forge/src/main/java/me/shedaniel/rei/mixin/forge/MixinClientPacketListener.java
@@ -0,0 +1,45 @@
+/*
+ * This file is licensed under the MIT License, part of Roughly Enough Items.
+ * Copyright (c) 2018, 2019, 2020, 2021, 2022 shedaniel
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package me.shedaniel.rei.mixin.forge;
+
+import me.shedaniel.rei.RoughlyEnoughItemsCoreClient;
+import net.minecraft.client.multiplayer.ClientPacketListener;
+import net.minecraft.network.protocol.game.ClientboundUpdateRecipesPacket;
+import net.minecraft.world.item.crafting.RecipeManager;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+@Mixin(ClientPacketListener.class)
+public class MixinClientPacketListener {
+    @Shadow @Final private RecipeManager recipeManager;
+    
+    @Inject(method = "handleUpdateRecipes", at = @At("HEAD"))
+    private void handleUpdateRecipes(ClientboundUpdateRecipesPacket clientboundUpdateRecipesPacket, CallbackInfo ci) {
+        RoughlyEnoughItemsCoreClient.PRE_UPDATE_RECIPES.invoker().update(recipeManager);
+    }
+}
\ No newline at end of file
diff --git a/forge/src/main/resources/rei.mixins.json b/forge/src/main/resources/rei.mixins.json
index f7fb10ab9..07bfe64ce 100644
--- a/forge/src/main/resources/rei.mixins.json
+++ b/forge/src/main/resources/rei.mixins.json
@@ -3,6 +3,7 @@
   "package": "me.shedaniel.rei.mixin.forge",
   "compatibilityLevel": "JAVA_8",
   "client": [
+    "MixinClientPacketListener",
     "MixinEffectRenderingInventoryScreen"
   ],
   "mixins": [
diff --git a/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java b/runtime/src/main/java/me/shedaniel/rei/RoughlyEnoughItemsCoreClient.java
index 07ce4f764..7d60bc638 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;
@@ -109,9 +110,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;
@@ -132,6 +131,7 @@ public class RoughlyEnoughItemsCoreClient {
         });
         return thread;
     });
+    private static final List<Future<?>> RELOAD_TASKS = new CopyOnWriteArrayList<>();
     
     public static void attachClientInternals() {
         InternalWidgets.attach();
@@ -321,8 +321,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);
@@ -455,7 +471,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;
+    }
 }