Optimize BotInitProcessor and add tests

mirai-core/src/commonMain/kotlin/network/components/BotInitProcessor.kt

@@ -10,7 +10,10 @@
 package net.mamoe.mirai.internal.network.components
 
 import kotlinx.atomicfu.atomic
-import kotlinx.coroutines.*
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.isActive
+import kotlinx.coroutines.launch
 import net.mamoe.mirai.internal.QQAndroidBot
 import net.mamoe.mirai.internal.network.component.ComponentKey
 import net.mamoe.mirai.internal.network.component.ComponentStorage
@@ -22,28 +25,41 @@ import net.mamoe.mirai.internal.network.handler.state.StateObserver
 import net.mamoe.mirai.internal.network.protocol.packet.chat.receive.MessageSvcPushForceOffline
 import net.mamoe.mirai.utils.MiraiLogger
 import net.mamoe.mirai.utils.Symbol
-import net.mamoe.mirai.utils.hierarchicalName
 
 
 /**
- * Facade of [ContactUpdater], [OtherClientUpdater], [ConfigPushSyncer].
  * Handles initialization jobs after successful logon.
  *
+ * The initialization includes:
+ * - Downloading contact list, which might read from local cache
+ * - Synchronizing message sequence id
+ * - Synchronizing BDH session for resource uploading
+ *
+ * Calls [ContactUpdater], [OtherClientUpdater], [ConfigPushSyncer], ... (see [BotInitProcessorImpl])
+ *
  * Attached to handler state [NetworkHandler.State.LOADING] [as state observer][asObserver] in [QQAndroidBot.stateObserverChain].
  */
 internal interface BotInitProcessor {
     /**
-     * Called when login was potentially halted. see [MessageSvcPushForceOffline]
+     * Do initialization. Implementor must ensure initialization runs exactly single time.
+     */
+    suspend fun init()
+
+    /**
+     * Called when login was potentially halted, meaning the data might not have been loaded,
+     * so we need to set the flag that helps keep single-initialization to UNINITIALIZED.
+     *
+     * This is called in [MessageSvcPushForceOffline], which is in case connection is closed by server during the [NetworkHandler.State.LOADING] state.
+     *
+     * See [BotInitProcessorImpl.state].
      */
     fun setLoginHalted()
 
-    suspend fun init(scope: CoroutineScope)
-
     companion object : ComponentKey<BotInitProcessor>
 }
 
 internal fun BotInitProcessor.asObserver(targetState: State = State.LOADING): StateObserver {
-    return JobAttachStateObserver("BotInitProcessor.init", targetState) { init(it) }
+    return JobAttachStateObserver("BotInitProcessor.init", targetState) { init() }
 }
 
 
@@ -64,10 +80,10 @@ internal class BotInitProcessorImpl(
         state.compareAndSet(expect = INITIALIZING, update = UNINITIALIZED)
     }
 
-    override suspend fun init(scope: CoroutineScope) {
+    override suspend fun init() {
         if (!state.compareAndSet(expect = UNINITIALIZED, update = INITIALIZING)) return
 
-        scope.launch(scope.hierarchicalName("BotInitProcessorImpl.init")) {
+        try {
             check(bot.isActive) { "bot is dead therefore network can't init." }
             context[ContactUpdater].closeAllContacts(CancellationException("re-init"))
 
@@ -86,11 +102,10 @@ internal class BotInitProcessorImpl(
 
             state.value = INITIALIZED
             bot.components[SsoProcessor].firstLoginSucceed = true
-        }.apply {
-            invokeOnCompletion { e ->
-                if (e != null) setLoginHalted()
-            }
-        }.join()
+        } catch (e: Throwable) {
+            setLoginHalted()
+            throw e
+        }
     }
 
     private inline fun runWithCoverage(block: () -> Unit) {

mirai-core/src/commonMain/kotlin/network/protocol/packet/chat/receive/MessageSvc.PushForceOffline.kt

@@ -31,7 +31,7 @@ internal object MessageSvcPushForceOffline :
 
     override suspend fun QQAndroidBot.handle(packet: RequestPushForceOffline) {
         components[AccountSecretsManager].invalidate() // otherwise you receive `MessageSvc.PushForceOffline` again just after logging in.
-        components[BotInitProcessor].setLoginHalted()
+        components[BotInitProcessor].setLoginHalted() // so that BotInitProcessor will be run on successful reconnection.
         network.close(ForceOfflineException(packet.title, "Closed by MessageSvc.PushForceOffline: $packet"))
     }
 }

mirai-core/src/commonTest/kotlin/network/component/BotInitProcessorTest.kt

@@ -0,0 +1,102 @@
+/*
+ * Copyright 2019-2021 Mamoe Technologies and contributors.
+ *
+ *  此源代码的使用受 GNU AFFERO GENERAL PUBLIC LICENSE version 3 许可证的约束, 可以在以下链接找到该许可证.
+ *  Use of this source code is governed by the GNU AGPLv3 license that can be found through the following link.
+ *
+ *  https://github.com/mamoe/mirai/blob/master/LICENSE
+ */
+
+package net.mamoe.mirai.internal.network.component
+
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.isActive
+import net.mamoe.mirai.internal.contact.uin
+import net.mamoe.mirai.internal.network.components.BotInitProcessor
+import net.mamoe.mirai.internal.network.framework.AbstractNettyNHTest
+import net.mamoe.mirai.internal.network.framework.AbstractNettyNHTestWithSelector
+import net.mamoe.mirai.internal.network.handler.NetworkHandler
+import net.mamoe.mirai.internal.network.protocol.data.jce.RequestPushForceOffline
+import net.mamoe.mirai.internal.network.protocol.packet.IncomingPacket
+import net.mamoe.mirai.internal.network.protocol.packet.chat.receive.MessageSvcPushForceOffline
+import net.mamoe.mirai.internal.test.runBlockingUnit
+import org.junit.jupiter.api.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+
+internal class BotInitProcessorTest {
+    class WithoutSelector : AbstractNettyNHTest() {
+        @Test
+        fun `BotInitProcessor halted`() = runBlockingUnit {
+            val p = setComponent(BotInitProcessor, object : BotInitProcessor {
+                var ranTimes = 0
+                var haltedTimes = 0
+                var def = CompletableDeferred<Unit>()
+                override suspend fun init() {
+                    ranTimes++
+                    def.await()
+                }
+
+                override fun setLoginHalted() {
+                    haltedTimes += 1
+                }
+            })
+            assertTrue { network.isActive }
+            network.setStateLoading(channel)
+            assertEquals(1, p.ranTimes)
+            assertEquals(0, p.haltedTimes)
+            assertState(NetworkHandler.State.LOADING)
+            network.collectReceived(IncomingPacket(
+                MessageSvcPushForceOffline.commandName,
+                RequestPushForceOffline(bot.uin)
+            ))
+            assertEquals(1, p.ranTimes)
+            assertEquals(1, p.haltedTimes)
+            eventDispatcher.joinBroadcast()
+            assertFalse { network.isActive }
+            network.assertState(NetworkHandler.State.CLOSED) // we do not use selector in this test so it will be CLOSED. It will recover (reconnect) instead in real.
+        }
+    }
+
+    class WithSelector : AbstractNettyNHTestWithSelector() {
+        @Test
+        fun `BotInitProcessor halted`() = runBlockingUnit {
+            bot.configuration.autoReconnectOnForceOffline = true
+            val p = setComponent(BotInitProcessor, object : BotInitProcessor {
+                var ranTimes = 0
+                var haltedTimes = 0
+                var def = CompletableDeferred<Unit>()
+                override suspend fun init() {
+                    ranTimes++
+                    def.await()
+                }
+
+                override fun setLoginHalted() {
+                    haltedTimes += 1
+                }
+            })
+            assertTrue { network.isActive }
+            network.setStateLoading(channel)
+            assertEquals(1, p.ranTimes)
+            assertEquals(0, p.haltedTimes)
+            assertState(NetworkHandler.State.LOADING)
+
+            network.currentInstance().collectReceived(IncomingPacket(
+                MessageSvcPushForceOffline.commandName,
+                RequestPushForceOffline(bot.uin)
+            ))
+            // all jobs launched during `collectReceived` are UNDISPATCHED, `collectReceived` returns on `def.await()` (suspension point)
+            // first run is CANCELLED.
+
+
+            assertEquals(1, p.ranTimes)
+            assertEquals(1, p.haltedTimes)
+
+            p.def.complete(Unit) // then BotInitProcessor.init finishes async.
+            network.resumeConnection() // join async
+            assertEquals(2, p.ranTimes) // we expect selector has run `init`
+            assertEquals(1, p.haltedTimes)
+        }
+    }
+}