Support refreshing instance in NetworkHandlerSelector

mirai-core/src/commonMain/kotlin/network/handler/selector/AbstractKeepAliveNetworkHandlerSelector.kt

@@ -46,6 +46,12 @@ internal abstract class AbstractKeepAliveNetworkHandlerSelector<H : NetworkHandl
 
     final override fun getResumedInstance(): H? = current.value
 
+    final override tailrec fun tryResumeInstanceOrCreate(): H {
+        getResumedInstance()?.let { return it }
+        refreshInstance()
+        return tryResumeInstanceOrCreate()
+    }
+
     final override suspend fun awaitResumeInstance(): H = awaitResumeInstanceImpl(0)
 
     private tailrec suspend fun awaitResumeInstanceImpl(attempted: Int): H {
@@ -73,13 +79,17 @@ internal abstract class AbstractKeepAliveNetworkHandlerSelector<H : NetworkHandl
                 }
             }
         } else {
-            synchronized(this) { // avoid concurrent `createInstance()`
+            refreshInstance()
-                if (getResumedInstance() == null) this.current.compareAndSet(null, createInstance())
-            }
             awaitResumeInstanceImpl(attempted) // directly retry, does not count for attempts.
         }
     }
 
+    protected open fun refreshInstance() {
+        synchronized(this) { // avoid concurrent `createInstance()`
+            if (getResumedInstance() == null) this.current.compareAndSet(null, createInstance())
+        }
+    }
+
     companion object {
         @JvmField
         var DEFAULT_MAX_ATTEMPTS =

mirai-core/src/commonMain/kotlin/network/handler/selector/NetworkHandlerSelector.kt

@@ -23,6 +23,11 @@ internal interface NetworkHandlerSelector<H : NetworkHandler> {
      */
     fun getResumedInstance(): H?
 
+    /**
+     * Returns the currently alive [NetworkHandler] or creates a new one.
+     */
+    fun tryResumeInstanceOrCreate(): H
+
     /**
      * Returns an alive [NetworkHandler], or suspends the coroutine until the connection has been made again.
      *

mirai-core/src/commonMain/kotlin/network/handler/selector/SelectorNetworkHandler.kt

@@ -10,7 +10,7 @@
 package net.mamoe.mirai.internal.network.handler.selector
 
 import kotlinx.coroutines.*
-import kotlinx.coroutines.selects.SelectClause1
+import kotlinx.coroutines.channels.ReceiveChannel
 import net.mamoe.mirai.internal.network.handler.NetworkHandler
 import net.mamoe.mirai.internal.network.handler.NetworkHandler.State
 import net.mamoe.mirai.internal.network.handler.NetworkHandlerContext
@@ -44,10 +44,9 @@ internal class SelectorNetworkHandler(
     }
 
     override val state: State
-        get() = selector.getResumedInstance()?.state ?: State.INITIALIZED
+        get() = selector.tryResumeInstanceOrCreate().state
-    override val onStateChanged: SelectClause1<State>
+    override val stateChannel: ReceiveChannel<State>
-        get() = selector.getResumedInstance()?.onStateChanged
+        get() = selector.tryResumeInstanceOrCreate().stateChannel
-            ?: scope.async { instance().state }.onAwait
 
     override suspend fun resumeConnection() {
         instance() // the selector will resume connection for us.

mirai-core/src/commonTest/kotlin/network/framework/TestNetworkHandler.kt

@@ -47,12 +47,16 @@ internal open class TestNetworkHandler(
                 }
             }
         }
+
+        override fun toString(): String {
+            return "TestState($correspondingState)"
+        }
     }
 
     @OptIn(TestOnly::class)
-    fun setState(correspondingState: NetworkHandler.State) {
+    fun setState(correspondingState: NetworkHandler.State): TestState? {
         // `null` means ignore checks. All test states have same type TestState.
-        setStateImpl(null) { TestState(correspondingState) }
+        return setStateImpl(null) { TestState(correspondingState) }
     }
 
     private val initialState = TestState(NetworkHandler.State.INITIALIZED)