From b90f86258c08c49e2f38fb724122dae6f5646e0d Mon Sep 17 00:00:00 2001
From: Him188 <Him188@mamoe.net>
Date: Fri, 2 Jul 2021 00:38:21 +0800
Subject: [PATCH] Add property `recoverable` to `HeartbeatFailedException`, to
 control recoverability.

Fix `SelectorRecoveryTest`
---
 .../impl/netty/HeartbeatFailedException.kt    | 10 ++++--
 .../AbstractRealNetworkHandlerTest.kt         |  7 ++++
 .../network/handler/SelectorRecoveryTest.kt   | 34 ++++++++++++++-----
 3 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/mirai-core/src/commonMain/kotlin/network/impl/netty/HeartbeatFailedException.kt b/mirai-core/src/commonMain/kotlin/network/impl/netty/HeartbeatFailedException.kt
index e1791e2bc..94e9a9802 100644
--- a/mirai-core/src/commonMain/kotlin/network/impl/netty/HeartbeatFailedException.kt
+++ b/mirai-core/src/commonMain/kotlin/network/impl/netty/HeartbeatFailedException.kt
@@ -10,11 +10,15 @@
 package net.mamoe.mirai.internal.network.impl.netty
 
 import net.mamoe.mirai.internal.network.handler.selector.NetworkException
+import net.mamoe.mirai.utils.unwrap
+import net.mamoe.mirai.utils.unwrapCancellationException
+import java.io.IOException
 
 internal class HeartbeatFailedException(
     private val name: String, // kind of HB
-    override val cause: Throwable? = null,
-) : NetworkException(true) {
+    override val cause: Throwable,
+    recoverable: Boolean = cause.unwrapCancellationException() is IOException || cause is NetworkException && cause.recoverable,
+) : NetworkException(recoverable) {
     override val message: String = "Exception in $name job"
-    override fun toString(): String = "HeartbeatFailedException: $name, cause=$cause"
+    override fun toString(): String = "HeartbeatFailedException: $name, recoverable=$recoverable, cause=$cause"
 }
\ 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 e346f11d0..417c9e1e0 100644
--- a/mirai-core/src/commonTest/kotlin/network/framework/AbstractRealNetworkHandlerTest.kt
+++ b/mirai-core/src/commonTest/kotlin/network/framework/AbstractRealNetworkHandlerTest.kt
@@ -146,6 +146,13 @@ internal sealed class AbstractRealNetworkHandlerTest<H : NetworkHandler> : Abstr
         assertEquals(state, network.state)
     }
 
+    fun assertState(vararg accepted: State) {
+        val s = network.state
+        if (s !in accepted) {
+            throw AssertionError("Expected: ${accepted.joinToString()}, actual: $s")
+        }
+    }
+
     fun NetworkHandler.assertState(state: State) {
         assertEquals(state, this.state)
     }
diff --git a/mirai-core/src/commonTest/kotlin/network/handler/SelectorRecoveryTest.kt b/mirai-core/src/commonTest/kotlin/network/handler/SelectorRecoveryTest.kt
index 86b4f7c33..96b8885c8 100644
--- a/mirai-core/src/commonTest/kotlin/network/handler/SelectorRecoveryTest.kt
+++ b/mirai-core/src/commonTest/kotlin/network/handler/SelectorRecoveryTest.kt
@@ -16,8 +16,10 @@ import net.mamoe.mirai.internal.network.components.HeartbeatFailureHandler
 import net.mamoe.mirai.internal.network.components.HeartbeatScheduler
 import net.mamoe.mirai.internal.network.framework.AbstractNettyNHTestWithSelector
 import net.mamoe.mirai.internal.network.impl.netty.HeartbeatFailedException
+import net.mamoe.mirai.internal.network.impl.netty.NettyChannelException
 import net.mamoe.mirai.internal.test.runBlockingUnit
 import org.junit.jupiter.api.Test
+import java.io.IOException
 import kotlin.test.assertFails
 
 internal class SelectorRecoveryTest : AbstractNettyNHTestWithSelector() {
@@ -33,16 +35,36 @@ internal class SelectorRecoveryTest : AbstractNettyNHTestWithSelector() {
      * @see HeartbeatFailedException
      */
     @Test
-    fun `can recover on heartbeat failure`() = runBlockingUnit {
-        testRecover { HeartbeatFailedException("test", null) } // NetworkException
+    fun `can recover on heartbeat failure with IOException`() = runBlockingUnit {
+        // We allow IOException to cause a reconnect.
+        testRecoverWhenHeartbeatFailWith { IOException("test IO ex") }
+
+        // BotOfflineMonitor immediately launches a recovery which is UNDISPATCHED, so connection is immediately recovered.
+        assertState(NetworkHandler.State.CONNECTING, NetworkHandler.State.OK)
+    }
+
+    /**
+     * Emulates system hibernation and network failure.
+     * @see HeartbeatFailedException
+     */
+    @Test
+    fun `can recover on heartbeat failure with NettyChannelException`() = runBlockingUnit {
+        // We allow IOException to cause a reconnect.
+        testRecoverWhenHeartbeatFailWith { NettyChannelException("test IO ex") }
+
+        // BotOfflineMonitor immediately launches a recovery which is UNDISPATCHED, so connection is immediately recovered.
+        assertState(NetworkHandler.State.CONNECTING, NetworkHandler.State.OK)
     }
 
     @Test
     fun `cannot recover on other failures`() = runBlockingUnit {
-        testRecover { IllegalStateException() }
+        // ISE is considered as an internal error (bug).
+        testRecoverWhenHeartbeatFailWith { IllegalStateException() }
+
+        assertState(NetworkHandler.State.CLOSED)
     }
 
-    private suspend fun testRecover(exception: () -> Exception) {
+    private suspend fun testRecoverWhenHeartbeatFailWith(exception: () -> Exception) {
         val heartbeatScheduler = object : HeartbeatScheduler {
             lateinit var onHeartFailure: HeartbeatFailureHandler
             override fun launchJobsIn(
@@ -61,9 +83,5 @@ internal class SelectorRecoveryTest : AbstractNettyNHTestWithSelector() {
         assertState(NetworkHandler.State.OK)
 
         heartbeatScheduler.onHeartFailure("Test", exception())
-        assertState(NetworkHandler.State.CLOSED)
-
-        bot.network.resumeConnection() // in real, this is called by BotOnlineWatchdog in SelectorNetworkHandler
-        assertState(NetworkHandler.State.OK)
     }
 }
\ No newline at end of file