Fix concurrency issues with recipe lookup

4 files changed, 40 insertions, 23 deletions

diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplayRegistryImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplayRegistryImpl.java
index 1aff7f720..f8f303408 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplayRegistryImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplayRegistryImpl.java
@@ -39,7 +39,6 @@ import me.shedaniel.rei.api.common.plugins.PluginManager;
 import me.shedaniel.rei.impl.common.InternalLogger;
 import me.shedaniel.rei.impl.common.registry.RecipeManagerContextImpl;
 import net.minecraft.world.item.crafting.Recipe;
-import org.apache.commons.lang3.mutable.MutableLong;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.*;
@@ -53,7 +52,7 @@ public class DisplayRegistryImpl extends RecipeManagerContextImpl<REIClientPlugi
     private final List<DynamicDisplayGenerator<?>> globalDisplayGenerators = new ArrayList<>();
     private final List<DisplayVisibilityPredicate> visibilityPredicates = new ArrayList<>();
     private final List<DisplayFiller<?>> fillers = new ArrayList<>();
-    private final MutableLong lastAddWarning = new MutableLong(-1);
+    private long lastAddWarning = -1;
     private DisplaysHolder displaysHolder = new DisplaysHolderImpl(false);
     
     public DisplayRegistryImpl() {
@@ -73,12 +72,11 @@ public class DisplayRegistryImpl extends RecipeManagerContextImpl<REIClientPlugi
     @Override
     public void add(Display display, @Nullable Object origin) {
         if (!PluginManager.areAnyReloading()) {
-            if (lastAddWarning != null) {
-                if (lastAddWarning.getValue() > 0 && System.currentTimeMillis() - lastAddWarning.getValue() > 5000) {
-                    InternalLogger.getInstance().warn("Detected runtime DisplayRegistry modification, this can be extremely dangerous!");
-                }
-                lastAddWarning.setValue(System.currentTimeMillis());
+            if (lastAddWarning < 0 || System.currentTimeMillis() - lastAddWarning > 5000) {
+                InternalLogger.getInstance().warn("Detected runtime DisplayRegistry modification, this can be extremely dangerous!");
+                InternalLogger.getInstance().debug("Detected runtime DisplayRegistry modification, this can be extremely dangerous!", new Throwable());
             }
+            lastAddWarning = System.currentTimeMillis();
         }
         
         this.displaysHolder.add(display, origin);
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolderImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolderImpl.java
index b2bdce14d..bf09618c7 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolderImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolderImpl.java
@@ -64,7 +64,7 @@ public class DisplaysHolderImpl implements DisplaysHolder {
             if (list == null) {
                 return null;
             } else {
-                return ((DisplaysList) list).unmodifiableList;
+                return ((DisplaysList) list).synchronizedList;
             }
         }, key -> CategoryRegistry.getInstance().tryGet(key).isPresent());
         this.displaysByInput = createSetMultimap();
@@ -183,12 +183,11 @@ public class DisplaysHolderImpl implements DisplaysHolder {
     }
     
     private static class DisplaysList extends ArrayList<Display> {
-        private final List<Display> unmodifiableList;
         private final List<Display> synchronizedList;
         
         public DisplaysList() {
-            this.synchronizedList = Collections.synchronizedList(this);
-            this.unmodifiableList = Collections.unmodifiableList(synchronizedList);
+            List<Display> unmodifiableList = Collections.unmodifiableList(this);
+            this.synchronizedList = Collections.synchronizedList(unmodifiableList);
         }
     }
     
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 bd181824d..d5140f2c8 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
@@ -231,6 +231,7 @@ public class PluginManagerImpl<P extends REIPlugin<?>> implements PluginManager<
     
     @Override
     public void pre(ReloadStage stage) {
+        this.reloading = stage;
         List<PluginWrapper<P>> plugins = new ArrayList<>(getPluginWrapped().toList());
         plugins.sort(Comparator.comparingDouble(PluginWrapper<P>::getPriority).reversed());
         Collections.reverse(plugins);
@@ -260,8 +261,22 @@ public class PluginManagerImpl<P extends REIPlugin<?>> implements PluginManager<
                 }
             });
         } catch (Throwable throwable) {
+            this.reloading = null;
             new RuntimeException("Failed to run pre registration").printStackTrace();
         }
+        try (SectionClosable preStageAll = section(stage, "pre-stage/");
+             PerformanceLogger.Plugin perfLogger = RoughlyEnoughItemsCore.PERFORMANCE_LOGGER.stage("Pre Stage " + stage.name())) {
+            for (Reloadable<P> reloadable : reloadables) {
+                Class<?> reloadableClass = reloadable.getClass();
+                try (SectionClosable preStage = section(stage, "pre-stage/" + name(reloadableClass) + "/");
+                     PerformanceLogger.Plugin.Inner inner = perfLogger.stage(name(reloadableClass))) {
+                    reloadable.preStage(stage);
+                } catch (Throwable throwable) {
+                    throwable.printStackTrace();
+                }
+            }
+        }
+        this.reloading = null;
         this.reloadStopwatch.stop();
         InternalLogger.getInstance().debug("========================================");
         InternalLogger.getInstance().debug(name(pluginClass) + " finished pre-reload for " + stage + " in " + reloadStopwatch + ".");
@@ -270,6 +285,7 @@ public class PluginManagerImpl<P extends REIPlugin<?>> implements PluginManager<
     
     @Override
     public void post(ReloadStage stage) {
+        this.reloading = stage;
         List<PluginWrapper<P>> plugins = new ArrayList<>(getPluginWrapped().toList());
         plugins.sort(Comparator.comparingDouble(PluginWrapper<P>::getPriority).reversed());
         Collections.reverse(plugins);
@@ -296,8 +312,22 @@ public class PluginManagerImpl<P extends REIPlugin<?>> implements PluginManager<
                 }
             });
         } catch (Throwable throwable) {
+            this.reloading = null;
             new RuntimeException("Failed to run post registration").printStackTrace();
         }
+        try (SectionClosable postStageAll = section(stage, "post-stage/");
+             PerformanceLogger.Plugin perfLogger = RoughlyEnoughItemsCore.PERFORMANCE_LOGGER.stage("Pre Stage " + stage.name())) {
+            for (Reloadable<P> reloadable : reloadables) {
+                Class<?> reloadableClass = reloadable.getClass();
+                try (SectionClosable postStage = section(stage, "post-stage/" + name(reloadableClass) + "/");
+                     PerformanceLogger.Plugin.Inner inner = perfLogger.stage(name(reloadableClass))) {
+                    reloadable.postStage(stage);
+                } catch (Throwable throwable) {
+                    throwable.printStackTrace();
+                }
+            }
+        }
+        this.reloading = null;
         this.reloadStopwatch.stop();
         postStopwatch.stop();
         InternalLogger.getInstance().debug("========================================");
diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/PluginStageExecutionWatcher.java b/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/PluginStageExecutionWatcher.java
index 539d7ab76..f579ebdad 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/PluginStageExecutionWatcher.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/PluginStageExecutionWatcher.java
@@ -69,13 +69,10 @@ public class PluginStageExecutionWatcher implements HintProvider {
             
             @Override
             public void startReload() {
-                for (ReloadStage stage : ReloadStage.values()) {
-                    startReload(stage);
-                }
             }
             
             @Override
-            public void startReload(ReloadStage stage) {
+            public void preStage(ReloadStage stage) {
                 synchronized (allStages) {
                     if (manager == PluginManager.getInstance() && stage.ordinal() == 0) {
                         allStages.clear();
@@ -85,14 +82,7 @@ public class PluginStageExecutionWatcher implements HintProvider {
             }
             
             @Override
-            public void endReload() {
-                for (ReloadStage stage : ReloadStage.values()) {
-                    endReload(stage);
-                }
-            }
-            
-            @Override
-            public void endReload(ReloadStage stage) {
+            public void postStage(ReloadStage stage) {
                 synchronized (allStages) {
                     data().finishedStages.add(stage);
                 }