Add more locking in various vanilla model loading paths

Author: embeddedt

common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/ModelBakerImplMixin.java

@@ -1,11 +1,17 @@
 package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
 
 import com.llamalad7.mixinextras.injector.ModifyReturnValue;
+import com.llamalad7.mixinextras.injector.wrapmethod.WrapMethod;
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import net.minecraft.client.resources.model.BakedModel;
 import net.minecraft.client.resources.model.ModelBakery;
+import net.minecraft.client.resources.model.ModelState;
 import net.minecraft.client.resources.model.UnbakedModel;
 import net.minecraft.resources.ResourceLocation;
 import org.embeddedt.modernfix.ModernFix;
 import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.embeddedt.modernfix.duck.IExtendedModelBakery;
+import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.Unique;
@@ -16,6 +22,7 @@
 public abstract class ModelBakerImplMixin {
     @Shadow public abstract UnbakedModel getModel(ResourceLocation location);
 
+    @Shadow @Final private ModelBakery field_40571;
     @Unique
     private int mfix$getDepth = 0;
 
@@ -36,4 +43,15 @@ private UnbakedModel resolveParents(UnbakedModel model) {
         mfix$getDepth--;
         return model;
     }
+
+    @WrapMethod(method = "bake(Lnet/minecraft/resources/ResourceLocation;Lnet/minecraft/client/resources/model/ModelState;)Lnet/minecraft/client/resources/model/BakedModel;")
+    private BakedModel mfix$lockWhenBaking(ResourceLocation location, ModelState transform, Operation<BakedModel> original) {
+        var lock = ((IExtendedModelBakery)this.field_40571).mfix$getLock();
+        lock.lock();
+        try {
+            return original.call(location, transform);
+        } finally {
+            lock.unlock();
+        }
+    }
 }

common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/ModelBakeryMixin.java

@@ -1,6 +1,7 @@
 package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
 
 import com.llamalad7.mixinextras.injector.ModifyExpressionValue;
+import com.llamalad7.mixinextras.injector.wrapmethod.WrapMethod;
 import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
 import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
 import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
@@ -108,6 +109,16 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
         }
     }
 
+    @WrapMethod(method = "getModel")
+    private UnbakedModel mfix$lockWhenGettingModel(ResourceLocation modelLocation, Operation<UnbakedModel> original) {
+        modelBakeryLock.lock();
+        try {
+            return original.call(modelLocation);
+        } finally {
+            modelBakeryLock.unlock();
+        }
+    }
+
     @Override
     public UnbakedModel mfix$getMissingModel() {
         return missingModel;
@@ -199,10 +210,26 @@ private void onInitialLoadFinish(BlockColors blockColors, ProfilerFiller profile
     private static final int MAXIMUM_CACHE_SIZE = 1000;
 
     private void runCleanup() {
-        ((LRUMap<?, ?>)this.unbakedCache).dropEntriesToMeetSize(MAXIMUM_CACHE_SIZE);
-        ((LRUMap<?, ?>)this.bakedCache).dropEntriesToMeetSize(MAXIMUM_CACHE_SIZE);
-        ((LRUMap<?, ?>)this.topLevelModels).dropEntriesToMeetSize(MAXIMUM_CACHE_SIZE);
-        ((LRUMap<?, ?>)this.bakedTopLevelModels).dropEntriesToMeetSize(MAXIMUM_CACHE_SIZE);
+        try {
+            ((LRUMap<?, ?>)this.unbakedCache).dropEntriesToMeetSize(MAXIMUM_CACHE_SIZE);
+        } catch(RuntimeException e) {
+            throw new IllegalStateException("Exception dropping entries in unbaked cache", e);
+        }
+        try {
+            ((LRUMap<?, ?>)this.bakedCache).dropEntriesToMeetSize(MAXIMUM_CACHE_SIZE);
+        } catch(RuntimeException e) {
+            throw new IllegalStateException("Exception dropping entries in baked cache", e);
+        }
+        try {
+            ((LRUMap<?, ?>)this.topLevelModels).dropEntriesToMeetSize(MAXIMUM_CACHE_SIZE);
+        } catch(RuntimeException e) {
+            throw new IllegalStateException("Exception dropping entries in top level models", e);
+        }
+        try {
+            ((LRUMap<?, ?>)this.bakedTopLevelModels).dropEntriesToMeetSize(MAXIMUM_CACHE_SIZE);
+        } catch(RuntimeException e) {
+            throw new IllegalStateException("Exception dropping entries in baked top level models", e);
+        }
     }
 
     @Override
@@ -235,4 +262,9 @@ private void runCleanup() {
     public Map<ModelResourceLocation, BakedModel> getBakedTopLevelModels() {
         return this.mfix$emulatedBakedRegistry;
     }
+
+    @Override
+    public ReentrantLock mfix$getLock() {
+        return this.modelBakeryLock;
+    }
 }

common/src/main/java/org/embeddedt/modernfix/duck/IExtendedModelBakery.java

@@ -3,9 +3,12 @@
 import net.minecraft.client.resources.model.ModelResourceLocation;
 import net.minecraft.client.resources.model.UnbakedModel;
 
+import java.util.concurrent.locks.ReentrantLock;
+
 public interface IExtendedModelBakery {
     void mfix$tick();
     void mfix$finishLoading();
     UnbakedModel mfix$loadUnbakedModelDynamic(ModelResourceLocation location);
     UnbakedModel mfix$getMissingModel();
+    ReentrantLock mfix$getLock();
 }

gradle.properties

@@ -7,7 +7,7 @@ mixinextras_version=0.4.1
 mod_id=modernfix
 minecraft_version=1.21.1
 enabled_platforms=fabric,neoforge
-forge_version=21.1.15
+forge_version=21.1.111
 parchment_version=2024.07.07
 parchment_mc_version=1.21
 refined_storage_version=4392788

neoforge/src/main/java/org/embeddedt/modernfix/neoforge/mixin/perf/dynamic_resources/ModelBakerImplMixin.java

@@ -1,16 +1,25 @@
 package org.embeddedt.modernfix.neoforge.mixin.perf.dynamic_resources;
 
+import com.llamalad7.mixinextras.injector.wrapmethod.WrapMethod;
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import net.minecraft.client.renderer.texture.TextureAtlasSprite;
+import net.minecraft.client.resources.model.BakedModel;
+import net.minecraft.client.resources.model.Material;
 import net.minecraft.client.resources.model.ModelBakery;
 import net.minecraft.client.resources.model.ModelResourceLocation;
+import net.minecraft.client.resources.model.ModelState;
 import net.minecraft.client.resources.model.UnbakedModel;
+import net.minecraft.resources.ResourceLocation;
 import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
 import org.embeddedt.modernfix.duck.IExtendedModelBakery;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Overwrite;
 import org.spongepowered.asm.mixin.Shadow;
 
-@Mixin(targets = {"net/minecraft/client/resources/model/ModelBakery$ModelBakerImpl"})
+import java.util.function.Function;
+
+@Mixin(ModelBakery.ModelBakerImpl.class)
 @ClientOnlyMixin
 public class ModelBakerImplMixin {
     @Shadow @Final private ModelBakery field_40571;
@@ -25,4 +34,15 @@ public UnbakedModel getTopLevelModel(ModelResourceLocation location) {
         UnbakedModel model = bakery.mfix$loadUnbakedModelDynamic(location);
         return model == bakery.mfix$getMissingModel() ? null : model;
     }
+
+    @WrapMethod(method = "bake(Lnet/minecraft/resources/ResourceLocation;Lnet/minecraft/client/resources/model/ModelState;Ljava/util/function/Function;)Lnet/minecraft/client/resources/model/BakedModel;", remap = false)
+    private BakedModel mfix$lockWhenBaking(ResourceLocation location, ModelState transform, Function<Material, TextureAtlasSprite> textureGetter, Operation<BakedModel> original) {
+        var lock = ((IExtendedModelBakery)this.field_40571).mfix$getLock();
+        lock.lock();
+        try {
+            return original.call(location, transform, textureGetter);
+        } finally {
+            lock.unlock();
+        }
+    }
 }