From 43b13d158fa399a3ddfbd122197e7454ed20c47a Mon Sep 17 00:00:00 2001
From: Him188 <Him188@mamoe.net>
Date: Sun, 13 Jun 2021 10:48:14 +0800
Subject: [PATCH] Implement `HeartbeatStrategy`, #1226

---
 .../network/components/HeartbeatProcessor.kt  |  8 +++++
 .../network/components/HeartbeatScheduler.kt  | 31 ++++++++++++++-----
 .../kotlin/network/components/SsoProcessor.kt |  6 ++++
 3 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/mirai-core/src/commonMain/kotlin/network/components/HeartbeatProcessor.kt b/mirai-core/src/commonMain/kotlin/network/components/HeartbeatProcessor.kt
index 2caf4ec97..e375e47bc 100644
--- a/mirai-core/src/commonMain/kotlin/network/components/HeartbeatProcessor.kt
+++ b/mirai-core/src/commonMain/kotlin/network/components/HeartbeatProcessor.kt
@@ -24,6 +24,9 @@ internal interface HeartbeatProcessor {
     @Throws(Exception::class)
     suspend fun doStatHeartbeatNow(networkHandler: NetworkHandler)
 
+    @Throws(Exception::class)
+    suspend fun doRegisterNow(networkHandler: NetworkHandler): StatSvc.Register.Response
+
     companion object : ComponentKey<HeartbeatProcessor>
 }
 
@@ -45,4 +48,9 @@ internal class HeartbeatProcessorImpl : HeartbeatProcessor {
             retry = 2
         )
     }
+
+    @Throws(Exception::class)
+    override suspend fun doRegisterNow(networkHandler: NetworkHandler): StatSvc.Register.Response {
+        return networkHandler.context[SsoProcessor].sendRegister(networkHandler)
+    }
 }
\ No newline at end of file
diff --git a/mirai-core/src/commonMain/kotlin/network/components/HeartbeatScheduler.kt b/mirai-core/src/commonMain/kotlin/network/components/HeartbeatScheduler.kt
index 2691f6e81..55bd34b5b 100644
--- a/mirai-core/src/commonMain/kotlin/network/components/HeartbeatScheduler.kt
+++ b/mirai-core/src/commonMain/kotlin/network/components/HeartbeatScheduler.kt
@@ -14,6 +14,7 @@ import net.mamoe.mirai.internal.network.component.ComponentKey
 import net.mamoe.mirai.internal.network.component.ComponentStorage
 import net.mamoe.mirai.internal.network.context.SsoProcessorContext
 import net.mamoe.mirai.internal.network.handler.NetworkHandlerSupport
+import net.mamoe.mirai.utils.BotConfiguration.HeartbeatStrategy.*
 import net.mamoe.mirai.utils.MiraiLogger
 import net.mamoe.mirai.utils.info
 
@@ -44,13 +45,29 @@ internal class TimeBasedHeartbeatSchedulerImpl(
         val heartbeatProcessor = context[HeartbeatProcessor]
 
         val list = mutableListOf<Job>()
-        list += launchHeartbeatJobAsync(
-            scope = scope,
-            name = "StatHeartbeat",
-            timeout = { context[SsoProcessorContext].configuration.statHeartbeatPeriodMillis },
-            action = { heartbeatProcessor.doStatHeartbeatNow(network) },
-            onHeartFailure = onHeartFailure
-        )
+        when (context[SsoProcessorContext].configuration.heartbeatStrategy) {
+            STAT_HB -> {
+                list += launchHeartbeatJobAsync(
+                    scope = scope,
+                    name = "StatHeartbeat",
+                    timeout = { context[SsoProcessorContext].configuration.statHeartbeatPeriodMillis },
+                    action = { heartbeatProcessor.doStatHeartbeatNow(network) },
+                    onHeartFailure = onHeartFailure
+                )
+            }
+            REGISTER -> {
+                list += launchHeartbeatJobAsync(
+                    scope = scope,
+                    name = "RegisterHeartbeat",
+                    timeout = { context[SsoProcessorContext].configuration.statHeartbeatPeriodMillis },
+                    action = { heartbeatProcessor.doRegisterNow(network) },
+                    onHeartFailure = onHeartFailure
+                )
+            }
+            NONE -> {
+            }
+        }
+
         list += launchHeartbeatJobAsync(
             scope = scope,
             name = "AliveHeartbeat",
diff --git a/mirai-core/src/commonMain/kotlin/network/components/SsoProcessor.kt b/mirai-core/src/commonMain/kotlin/network/components/SsoProcessor.kt
index bb1ef8f22..e724d340b 100644
--- a/mirai-core/src/commonMain/kotlin/network/components/SsoProcessor.kt
+++ b/mirai-core/src/commonMain/kotlin/network/components/SsoProcessor.kt
@@ -63,6 +63,8 @@ internal interface SsoProcessor {
 
     suspend fun logout(handler: NetworkHandler)
 
+    suspend fun sendRegister(handler: NetworkHandler): StatSvc.Register.Response
+
     companion object : ComponentKey<SsoProcessor>
 }
 
@@ -125,6 +127,10 @@ internal class SsoProcessorImpl(
         ssoContext.bot.logger.info { "Login successful." }
     }
 
+    override suspend fun sendRegister(handler: NetworkHandler): StatSvc.Register.Response {
+        return registerClientOnline(handler).also { registerResp = it }
+    }
+
     private suspend fun registerClientOnline(handler: NetworkHandler): StatSvc.Register.Response {
         return StatSvc.Register.online(client).sendAndExpect(handler).also { registerResp = it }
     }