From 1d2b2377cfa680ae9a3c21269a539721487216f9 Mon Sep 17 00:00:00 2001
From: Him188 <Him188@mamoe.net>
Date: Mon, 10 May 2021 13:13:58 +0800
Subject: [PATCH] Add ConcurrentComponentStorage.creationStacktrace

---
 .../component/ConcurrentComponentStorage.kt       | 15 ++++++++++++---
 .../framework/AbstractRealNetworkHandlerTest.kt   |  1 +
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/mirai-core/src/commonMain/kotlin/network/component/ConcurrentComponentStorage.kt b/mirai-core/src/commonMain/kotlin/network/component/ConcurrentComponentStorage.kt
index 5ebe63919..43fd6d3e3 100644
--- a/mirai-core/src/commonMain/kotlin/network/component/ConcurrentComponentStorage.kt
+++ b/mirai-core/src/commonMain/kotlin/network/component/ConcurrentComponentStorage.kt
@@ -17,7 +17,9 @@ import kotlin.LazyThreadSafetyMode.NONE
  * A thread-safe implementation of [MutableComponentStorage]
  */
 internal class ConcurrentComponentStorage(
-    private val showAllComponents: Boolean = SHOW_ALL_COMPONENTS
+    private val showAllComponents: Boolean = SHOW_ALL_COMPONENTS,
+    private val creationStacktrace: Exception? =
+        if (SHOW_COMPONENTS_CREATION_STACKTRACE) Exception("Creation stacktrace") else null,
 ) : ComponentStorage, MutableComponentStorage {
     private val map = ConcurrentHashMap<ComponentKey<*>, Any?>()
 
@@ -25,7 +27,8 @@ internal class ConcurrentComponentStorage(
     override val size: Int get() = map.size
 
     override operator fun <T : Any> get(key: ComponentKey<T>): T {
-        return getOrNull(key) ?: throw NoSuchComponentException(key, this)
+        return getOrNull(key)
+            ?: throw NoSuchComponentException(key, this).apply { creationStacktrace?.let(::initCause) }
     }
 
     override fun <T : Any> getOrNull(key: ComponentKey<T>): T? {
@@ -56,4 +59,10 @@ internal class ConcurrentComponentStorage(
     }
 }
 
-private val SHOW_ALL_COMPONENTS: Boolean by lazy(NONE) { systemProp("mirai.debug.network.show.all.components", false) }
\ No newline at end of file
+private val SHOW_ALL_COMPONENTS: Boolean by lazy(NONE) { systemProp("mirai.debug.network.show.all.components", false) }
+private val SHOW_COMPONENTS_CREATION_STACKTRACE: Boolean by lazy(NONE) {
+    systemProp(
+        "mirai.debug.network.show.components.creation.stacktrace",
+        false
+    )
+}
\ No newline at end of file
diff --git a/mirai-core/src/commonTest/kotlin/network/framework/AbstractRealNetworkHandlerTest.kt b/mirai-core/src/commonTest/kotlin/network/framework/AbstractRealNetworkHandlerTest.kt
index bf42752f7..aaff662f1 100644
--- a/mirai-core/src/commonTest/kotlin/network/framework/AbstractRealNetworkHandlerTest.kt
+++ b/mirai-core/src/commonTest/kotlin/network/framework/AbstractRealNetworkHandlerTest.kt
@@ -44,6 +44,7 @@ internal abstract class AbstractRealNetworkHandlerTest<H : NetworkHandler> : Abs
     init {
         System.setProperty("mirai.debug.network.state.observer.logging", "full")
         System.setProperty("mirai.debug.network.show.all.components", "true")
+        System.setProperty("mirai.debug.network.show.components.creation.stacktrace", "true")
     }
 
     protected abstract val factory: NetworkHandlerFactory<H>