From d96c1f84d494f8c43026d2546b2223392fb921e9 Mon Sep 17 00:00:00 2001
From: HopeBaron <kirito70k@gmail.com>
Date: Mon, 23 Nov 2020 11:41:49 +0200
Subject: [PATCH] Move reconnect retry to handshake handler

---
 CHANGELOG.md                                  |  4 ++
 .../gitlab/kordlib/gateway/DefaultGateway.kt  | 40 +++++++++----------
 .../gateway/handler/HandshakeHandler.kt       |  5 ++-
 3 files changed, 26 insertions(+), 23 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ea35396813cb..c8ec30b95323 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,10 @@
 * Re-added `values`  in `Intent` and `Intents`.
 * Added support for inline replies and `MessageBehavior#reply` to quickly create a reply to a message. #110
 
+## Changes
+
+* Gateway reconnect retries reset on handshake. #68
+
 # 0.7.0-M1
 
 ## Changes
diff --git a/gateway/src/main/kotlin/com/gitlab/kordlib/gateway/DefaultGateway.kt b/gateway/src/main/kotlin/com/gitlab/kordlib/gateway/DefaultGateway.kt
index bb035d3f0088..1b96aa4b2abd 100644
--- a/gateway/src/main/kotlin/com/gitlab/kordlib/gateway/DefaultGateway.kt
+++ b/gateway/src/main/kotlin/com/gitlab/kordlib/gateway/DefaultGateway.kt
@@ -6,15 +6,12 @@ import com.gitlab.kordlib.common.ratelimit.RateLimiter
 import com.gitlab.kordlib.gateway.GatewayCloseCode.*
 import com.gitlab.kordlib.gateway.handler.*
 import com.gitlab.kordlib.gateway.retry.Retry
-import io.ktor.client.HttpClient
-import io.ktor.client.features.websocket.DefaultClientWebSocketSession
-import io.ktor.client.features.websocket.webSocketSession
-import io.ktor.client.request.url
-import io.ktor.http.URLBuilder
-import io.ktor.http.cio.websocket.CloseReason
-import io.ktor.http.cio.websocket.Frame
-import io.ktor.http.cio.websocket.close
-import io.ktor.util.error
+import io.ktor.client.*
+import io.ktor.client.features.websocket.*
+import io.ktor.client.request.*
+import io.ktor.http.*
+import io.ktor.http.cio.websocket.*
+import io.ktor.util.*
 import kotlinx.atomicfu.AtomicRef
 import kotlinx.atomicfu.atomic
 import kotlinx.atomicfu.update
@@ -91,7 +88,7 @@ class DefaultGateway(private val data: DefaultGatewayData) : Gateway {
     init {
         val sequence = Sequence()
         SequenceHandler(events, sequence)
-        handshakeHandler = HandshakeHandler(events, ::trySend, sequence, data.identifyRateLimiter)
+        handshakeHandler = HandshakeHandler(events, ::trySend, sequence, data.identifyRateLimiter, data.reconnectRetry)
         HeartbeatHandler(events, ::trySend, { restart(Close.ZombieConnection) }, { _ping.value = it }, sequence)
         ReconnectHandler(events) { restart(Close.Reconnecting) }
         InvalidSessionHandler(events) { restart(it) }
@@ -122,7 +119,6 @@ class DefaultGateway(private val data: DefaultGatewayData) : Gateway {
 
             try {
                 readSocket()
-                data.reconnectRetry.reset() //connected and read without problems, resetting retry counter
             } catch (exception: Exception) {
                 defaultGatewayLogger.error(exception)
             }
@@ -277,8 +273,7 @@ class DefaultGateway(private val data: DefaultGatewayData) : Gateway {
                 val copy = command.copy(token = "token")
                 "Gateway >>> ${Json.encodeToString(Command.Companion, copy)}"
             }
-        }
-        else defaultGatewayLogger.trace { "Gateway >>> $json" }
+        } else defaultGatewayLogger.trace { "Gateway >>> $json" }
         socket.send(Frame.Text(json))
     }
 
@@ -294,15 +289,16 @@ class DefaultGateway(private val data: DefaultGatewayData) : Gateway {
     }
 }
 
-internal val GatewayConfiguration.identify get() = Identify(
-        token,
-        IdentifyProperties(os, name, name),
-        false.optional(),
-        50.optionalInt(),
-        shard.optional(),
-        presence,
-        intents
-)
+internal val GatewayConfiguration.identify
+    get() = Identify(
+            token,
+            IdentifyProperties(os, name, name),
+            false.optional(),
+            50.optionalInt(),
+            shard.optional(),
+            presence,
+            intents
+    )
 
 
 internal val os: String get() = System.getProperty("os.name")
diff --git a/gateway/src/main/kotlin/com/gitlab/kordlib/gateway/handler/HandshakeHandler.kt b/gateway/src/main/kotlin/com/gitlab/kordlib/gateway/handler/HandshakeHandler.kt
index fb006ec0f0a9..9d12b4e0d4da 100644
--- a/gateway/src/main/kotlin/com/gitlab/kordlib/gateway/handler/HandshakeHandler.kt
+++ b/gateway/src/main/kotlin/com/gitlab/kordlib/gateway/handler/HandshakeHandler.kt
@@ -3,6 +3,7 @@ package com.gitlab.kordlib.gateway.handler
 import com.gitlab.kordlib.common.ratelimit.RateLimiter
 import com.gitlab.kordlib.common.ratelimit.consume
 import com.gitlab.kordlib.gateway.*
+import com.gitlab.kordlib.gateway.retry.Retry
 import kotlinx.atomicfu.AtomicRef
 import kotlinx.atomicfu.atomic
 import kotlinx.atomicfu.update
@@ -12,7 +13,8 @@ internal class HandshakeHandler(
         flow: Flow<Event>,
         private val send: suspend (Command) -> Unit,
         private val sequence: Sequence,
-        private val identifyRateLimiter: RateLimiter
+        private val identifyRateLimiter: RateLimiter,
+        private val reconnectRetry: Retry
 ) : Handler(flow, "HandshakeHandler") {
 
     lateinit var configuration: GatewayConfiguration
@@ -33,6 +35,7 @@ internal class HandshakeHandler(
         }
 
         on<Hello> {
+            reconnectRetry.reset() //connected and read without problems, resetting retry counter
             identifyRateLimiter.consume {
                 if (sessionStart) send(identify)
                 else send(resume)