From a498f08217899651161b43218bf4533386f7f16d Mon Sep 17 00:00:00 2001
From: Andrei <b1naryth1ef@gmail.com>
Date: Thu, 22 Dec 2016 17:37:01 -0600
Subject: [PATCH] feature - add ability to control gateway max reconnects vi
 config

---
 disco/client.py         | 5 ++++-
 disco/gateway/client.py | 8 ++++----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/disco/client.py b/disco/client.py
index 0de6f45..a5791ee 100644
--- a/disco/client.py
+++ b/disco/client.py
@@ -26,6 +26,8 @@ class ClientConfig(Config):
         The shard ID for the current client instance.
     shard_count : int
         The total count of shards running.
+    max_reconnects : int
+        The maximum number of connection retries to make before giving up (0 = never give up).
     manhole_enable : bool
         Whether to enable the manhole (e.g. console backdoor server) utility.
     manhole_bind : tuple(str, int)
@@ -39,6 +41,7 @@ class ClientConfig(Config):
     token = ""
     shard_id = 0
     shard_count = 1
+    max_reconnects = 5
 
     manhole_enable = False
     manhole_bind = ('127.0.0.1', 8484)
@@ -86,7 +89,7 @@ class Client(LoggingClass):
         self.packets = Emitter(gevent.spawn)
 
         self.api = APIClient(self.config.token, self)
-        self.gw = GatewayClient(self, self.config.encoder)
+        self.gw = GatewayClient(self, self.config.max_reconnects, self.config.encoder)
         self.state = State(self, StateConfig(self.config.get('state', {})))
 
         if self.config.manhole_enable:
diff --git a/disco/gateway/client.py b/disco/gateway/client.py
index 9544af5..7386bb5 100644
--- a/disco/gateway/client.py
+++ b/disco/gateway/client.py
@@ -15,11 +15,11 @@ TEN_MEGABYTES = 10490000
 
 class GatewayClient(LoggingClass):
     GATEWAY_VERSION = 6
-    MAX_RECONNECTS = 5
 
-    def __init__(self, client, encoder='json', ipc=None):
+    def __init__(self, client, max_reconnects=5, encoder='json', ipc=None):
         super(GatewayClient, self).__init__()
         self.client = client
+        self.max_reconnects = max_reconnects
         self.encoder = ENCODERS[encoder]
 
         self.events = client.events
@@ -192,8 +192,8 @@ class GatewayClient(LoggingClass):
         self.reconnects += 1
         self.log.info('WS Closed: [%s] %s (%s)', code, reason, self.reconnects)
 
-        if self.MAX_RECONNECTS and self.reconnects > self.MAX_RECONNECTS:
-            raise Exception('Failed to reconnect after {} attempts, giving up'.format(self.MAX_RECONNECTS))
+        if self.max_reconnects and self.reconnects > self.max_reconnects:
+            raise Exception('Failed to reconnect after {} attempts, giving up'.format(self.max_reconnects))
 
         # Don't resume for these error codes
         if code and 4000 <= code <= 4010: