From e788aa09f7523a977a252ff73075f1369ed7fbc5 Mon Sep 17 00:00:00 2001
From: Him188 <Him188@mamoe.net>
Date: Sat, 5 Jun 2021 14:45:58 +0800
Subject: [PATCH] Add `awaitStateChange`

---
 .../src/commonMain/kotlin/CoroutineUtils.kt      | 11 ++++++++++-
 .../kotlin/network/handler/NetworkHandler.kt     | 10 ++++++++++
 .../commonTest/kotlin/network/AwaitStateTest.kt  | 16 ++++++++++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/mirai-core-utils/src/commonMain/kotlin/CoroutineUtils.kt b/mirai-core-utils/src/commonMain/kotlin/CoroutineUtils.kt
index 493bf8347..c02eaaf3e 100644
--- a/mirai-core-utils/src/commonMain/kotlin/CoroutineUtils.kt
+++ b/mirai-core-utils/src/commonMain/kotlin/CoroutineUtils.kt
@@ -88,4 +88,13 @@ public inline fun CoroutineContext.addNameIfAbsent(
 
 public fun CoroutineContext.addNameHierarchically(
     name: String
-): CoroutineContext = this + CoroutineName(this[CoroutineName]?.name?.plus('.')?.plus(name) ?: name)
\ No newline at end of file
+): CoroutineContext = this + CoroutineName(this[CoroutineName]?.name?.plus('.')?.plus(name) ?: name)
+
+public fun CoroutineContext.hierarchicalName(
+    name: String
+): CoroutineName = CoroutineName(this[CoroutineName]?.name?.plus('.')?.plus(name) ?: name)
+
+public fun CoroutineScope.hierarchicalName(
+    name: String
+): CoroutineName = this.coroutineContext.hierarchicalName(name)
+
diff --git a/mirai-core/src/commonMain/kotlin/network/handler/NetworkHandler.kt b/mirai-core/src/commonMain/kotlin/network/handler/NetworkHandler.kt
index 7083020f0..2423988ce 100644
--- a/mirai-core/src/commonMain/kotlin/network/handler/NetworkHandler.kt
+++ b/mirai-core/src/commonMain/kotlin/network/handler/NetworkHandler.kt
@@ -13,6 +13,7 @@ import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.ReceiveChannel
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.consumeAsFlow
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.takeWhile
 import net.mamoe.mirai.Bot
 import net.mamoe.mirai.internal.network.Packet
@@ -44,6 +45,11 @@ internal interface NetworkHandler : CoroutineScope {
      */
     val state: State
 
+    fun getLastFailure(): Throwable?
+
+    /**
+     * The channel that is sent with a [State] when changed.
+     */
     val stateChannel: ReceiveChannel<State>
 
     /**
@@ -184,4 +190,8 @@ internal val NetworkHandler.logger: MiraiLogger get() = context.logger
 internal suspend fun NetworkHandler.awaitState(suspendUntil: NetworkHandler.State) {
     if (this.state == suspendUntil) return
     stateChannel.consumeAsFlow().takeWhile { it != suspendUntil }.collect()
+}
+
+internal suspend fun NetworkHandler.awaitStateChange() {
+    stateChannel.consumeAsFlow().first()
 }
\ No newline at end of file
diff --git a/mirai-core/src/commonTest/kotlin/network/AwaitStateTest.kt b/mirai-core/src/commonTest/kotlin/network/AwaitStateTest.kt
index 79934a362..27588e9e9 100644
--- a/mirai-core/src/commonTest/kotlin/network/AwaitStateTest.kt
+++ b/mirai-core/src/commonTest/kotlin/network/AwaitStateTest.kt
@@ -17,6 +17,7 @@ import net.mamoe.mirai.internal.network.framework.AbstractMockNetworkHandlerTest
 import net.mamoe.mirai.internal.network.handler.NetworkHandler
 import net.mamoe.mirai.internal.network.handler.NetworkHandler.State.*
 import net.mamoe.mirai.internal.network.handler.awaitState
+import net.mamoe.mirai.internal.network.handler.awaitStateChange
 import net.mamoe.mirai.internal.test.runBlockingUnit
 import org.junit.jupiter.api.Test
 import java.util.concurrent.ConcurrentLinkedQueue
@@ -132,4 +133,19 @@ internal class AwaitStateTest : AbstractMockNetworkHandlerTest() {
             yield()
         }
     }
+
+    @Test
+    fun `can await change`() = runBlockingUnit(singleThreadDispatcher + Job()) {
+        createNetworkHandler().run {
+            assertState(INITIALIZED)
+            val job = launch(start = CoroutineStart.UNDISPATCHED) {
+                awaitStateChange()
+            }
+            yield()
+            assertTrue { job.isActive }
+
+            setState(CLOSED)
+            yield()
+        }
+    }
 }
\ No newline at end of file