From 7b8a918b7feed21e5f004d735f260c01c44180ed Mon Sep 17 00:00:00 2001
From: Justin <14909116+ThatGuyJustin@users.noreply.github.com>
Date: Mon, 7 Oct 2019 10:54:12 -0400
Subject: [PATCH] Adding latency checking

*Note: Latency is in seconds NOT ms*
---
 disco/gateway/client.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/disco/gateway/client.py b/disco/gateway/client.py
index a951145..5ef02f1 100644
--- a/disco/gateway/client.py
+++ b/disco/gateway/client.py
@@ -2,6 +2,7 @@ import gevent
 import zlib
 import six
 import ssl
+import time
 
 from websocket import ABNF
 
@@ -70,6 +71,11 @@ class GatewayClient(LoggingClass):
         self._heartbeat_task = None
         self._heartbeat_acknowledged = True
 
+        # Latency
+        self._last_send = time.perf_counter()
+        self._last_ack = time.perf_counter()
+        self.latency = float('inf')
+
     def send(self, op, data):
         self.limiter.check()
         return self._send(op, data)
@@ -89,11 +95,18 @@ class GatewayClient(LoggingClass):
                 self._heartbeat_acknowledged = True
                 self.ws.close(status=4000)
                 return
+            else:
+                self._last_send = time.perf_counter()
 
             self._send(OPCode.HEARTBEAT, self.seq)
             self._heartbeat_acknowledged = False
             gevent.sleep(interval / 1000)
 
+    def ack(self):
+        ack_time = time.perf_counter()
+        self._last_ack = ack_time
+        self.latency = ack_time - self._last_send
+
     def handle_dispatch(self, packet):
         obj = GatewayEvent.from_dispatch(self.client, packet)
         self.log.debug('GatewayClient.handle_dispatch %s', obj.__class__.__name__)
@@ -107,6 +120,7 @@ class GatewayClient(LoggingClass):
     def handle_heartbeat_acknowledge(self, _):
         self.log.debug('Received HEARTBEAT_ACK')
         self._heartbeat_acknowledged = True
+        self.ack()
 
     def handle_reconnect(self, _):
         self.log.warning('Received RECONNECT request, forcing a fresh reconnect')