Fix enqueue tasks in mod loading

KiltMC · Sep 22, 2024 · e8ecb3e · e8ecb3e
1 parent 3ecb37c
commit e8ecb3e
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 7 deletions.
diff --git a/src/main/kotlin/net/minecraftforge/fml/DeferredWorkQueue.kt b/src/main/kotlin/net/minecraftforge/fml/DeferredWorkQueue.kt
@@ -0,0 +1,76 @@
+package net.minecraftforge.fml
+
+import xyz.bluspring.kilt.Kilt
+import java.util.*
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.function.Supplier
+
+class DeferredWorkQueue(private val stage: ModLoadingStage) {
+    private val tasks = ConcurrentLinkedQueue<Task>()
+
+    init {
+        queues[stage] = this
+    }
+
+    fun runTasks() {
+        if (tasks.isEmpty())
+            return
+
+        val exception = RuntimeException()
+
+        for (task in tasks) {
+            ModLoadingContext.kiltActiveModId = task.owner.modId
+
+            try {
+                task.future.exceptionally {
+                    Kilt.logger.error("Mod \"${task.owner.modId}\" encountered an error in a deferred task!")
+                    it.printStackTrace()
+
+                    exception.addSuppressed(it)
+                    null
+                }
+                task.task.run()
+            } finally {
+                ModLoadingContext.kiltActiveModId = null
+            }
+        }
+
+        if (exception.suppressed.isNotEmpty()) {
+            throw exception
+        }
+    }
+
+    fun enqueueWork(modInfo: ModContainer, work: Runnable): CompletableFuture<Void> {
+        return addTask(modInfo) { task -> CompletableFuture.runAsync(work) { task.task = it } }
+    }
+
+    fun <T> enqueueWork(modInfo: ModContainer, work: Supplier<T>): CompletableFuture<T> {
+        return addTask(modInfo) { task -> CompletableFuture.supplyAsync(work) { task.task = it } }
+    }
+
+    private fun <T> addTask(modInfo: ModContainer, futureSupplier: (Task) -> CompletableFuture<T>): CompletableFuture<T> {
+        val task = Task(modInfo)
+        val future = futureSupplier.invoke(task)
+        task.future = future
+        tasks.add(task)
+
+        return future
+    }
+
+    companion object {
+        private val queues = mutableMapOf<ModLoadingStage, DeferredWorkQueue>()
+
+        @JvmStatic
+        fun lookup(parallelClass: Optional<ModLoadingStage>): Optional<DeferredWorkQueue> {
+            return Optional.ofNullable(queues[parallelClass.orElse(null)])
+        }
+    }
+
+    private data class Task(
+        val owner: ModContainer
+    ) {
+        lateinit var task: Runnable
+        lateinit var future: CompletableFuture<*>
+    }
+}
diff --git a/src/main/kotlin/net/minecraftforge/fml/ModLoadingStage.kt b/src/main/kotlin/net/minecraftforge/fml/ModLoadingStage.kt
@@ -1,5 +1,7 @@
 package net.minecraftforge.fml
 
 enum class ModLoadingStage {
-    ERROR, VALIDATE, CONSTRUCT, COMMON_SETUP, SIDED_SETUP, ENQUEUE_IMC, PROCESS_IMC, COMPLETE, DONE
+    ERROR, VALIDATE, CONSTRUCT, COMMON_SETUP, SIDED_SETUP, ENQUEUE_IMC, PROCESS_IMC, COMPLETE, DONE;
+
+    val deferredWorkQueue = DeferredWorkQueue(this)
 }
diff --git a/src/main/kotlin/net/minecraftforge/fml/event/lifecycle/ModLifecycleEvent.kt b/src/main/kotlin/net/minecraftforge/fml/event/lifecycle/ModLifecycleEvent.kt
@@ -10,6 +10,8 @@ import java.util.stream.Stream
 open class ModLifecycleEvent(private val mod: ForgeMod?) : Event(), IModBusEvent {
     constructor() : this(null)
 
+    val container = mod?.container
+
     fun description(): String {
         return this.javaClass.name.run {
             this.substring(this.lastIndexOf('.') + 1)

diff --git a/src/main/kotlin/net/minecraftforge/fml/event/lifecycle/ParallelDispatchEvent.kt b/src/main/kotlin/net/minecraftforge/fml/event/lifecycle/ParallelDispatchEvent.kt
@@ -1,19 +1,23 @@
 package net.minecraftforge.fml.event.lifecycle
 
+import net.minecraftforge.fml.DeferredWorkQueue
 import net.minecraftforge.fml.ModLoadingStage
 import xyz.bluspring.kilt.loader.mod.ForgeMod
+import java.util.*
 import java.util.concurrent.CompletableFuture
 import java.util.function.Supplier
 
 // TODO: this should probably be made better
 open class ParallelDispatchEvent(mod: ForgeMod?, private val stage: ModLoadingStage?) : ModLifecycleEvent(mod) {
     constructor() : this(null, null)
 
+    private val queue = DeferredWorkQueue.lookup(Optional.ofNullable(stage))
+
     fun enqueueWork(work: Runnable): CompletableFuture<Void> {
-        return CompletableFuture.runAsync(work)
+        return queue.orElseThrow().enqueueWork(this.container!!, work)
     }
 
     fun <T> enqueueWork(work: Supplier<T>): CompletableFuture<T> {
-        return CompletableFuture.supplyAsync(work)
+        return queue.orElseThrow().enqueueWork(this.container!!, work)
     }
 }
diff --git a/src/main/kotlin/xyz/bluspring/kilt/loader/KiltLoader.kt b/src/main/kotlin/xyz/bluspring/kilt/loader/KiltLoader.kt
@@ -582,6 +582,12 @@ class KiltLoader {
             exceptions.addAll(initMod(mod, mod.scanData))
         }
 
+        try {
+            ModLoadingStage.CONSTRUCT.deferredWorkQueue.runTasks()
+        } catch (e: Exception) {
+            exceptions.add(e)
+        }
+
         if (exceptions.isNotEmpty()) {
             FabricGuiEntry.displayError("Errors occurred while initializing Forge mods!", null, {
                 val tab = it.addTab("Kilt Error")
@@ -686,6 +692,7 @@ class KiltLoader {
 
             // COMMON_SETUP
             ModLoader.get().kiltPostEventWrappingModsBuildEvent { FMLCommonSetupEvent(it, ModLoadingStage.COMMON_SETUP) }
+            ModLoadingStage.COMMON_SETUP.deferredWorkQueue.runTasks()
 
             // SIDED_SETUP
             ModLoader.get().kiltPostEventWrappingModsBuildEvent {
@@ -694,15 +701,19 @@ class KiltLoader {
                 else
                     FMLDedicatedServerSetupEvent(it, ModLoadingStage.SIDED_SETUP)
             }
+            ModLoadingStage.SIDED_SETUP.deferredWorkQueue.runTasks()
 
             // ENQUEUE_IMC
             ModLoader.get().kiltPostEventWrappingModsBuildEvent { InterModEnqueueEvent(it, ModLoadingStage.ENQUEUE_IMC) }
+            ModLoadingStage.ENQUEUE_IMC.deferredWorkQueue.runTasks()
 
             // PROCESS_IMC
             ModLoader.get().kiltPostEventWrappingModsBuildEvent { InterModProcessEvent(it, ModLoadingStage.PROCESS_IMC) }
+            ModLoadingStage.PROCESS_IMC.deferredWorkQueue.runTasks()
 
             // COMPLETE
             ModLoader.get().kiltPostEventWrappingModsBuildEvent { FMLLoadCompleteEvent(it, ModLoadingStage.COMPLETE) }
+            ModLoadingStage.COMPLETE.deferredWorkQueue.runTasks()
         }
     )
 

diff --git a/src/main/kotlin/xyz/bluspring/kilt/server/dedicated/KiltDedicatedServer.kt b/src/main/kotlin/xyz/bluspring/kilt/server/dedicated/KiltDedicatedServer.kt
@@ -1,14 +1,13 @@
 package xyz.bluspring.kilt.server.dedicated
 
 import net.fabricmc.api.DedicatedServerModInitializer
-import net.minecraftforge.fml.ModLoadingStage
-import net.minecraftforge.fml.event.lifecycle.FMLDedicatedServerSetupEvent
-import xyz.bluspring.kilt.Kilt
 
 class KiltDedicatedServer : DedicatedServerModInitializer {
     override fun onInitializeServer() {
-        Kilt.loader.mods.forEach { mod ->
+        /*Kilt.loader.mods.forEach { mod ->
             mod.eventBus.post(FMLDedicatedServerSetupEvent(mod, ModLoadingStage.SIDED_SETUP))
         }
+
+        ModLoadingStage.SIDED_SETUP.deferredWorkQueue.runTasks()*/
     }
 }