From e94e14e9d240bf0433796a74292e0673369c4923 Mon Sep 17 00:00:00 2001
From: Chris Rhodes <iopred@gmail.com>
Date: Thu, 21 Jan 2016 14:15:24 -0800
Subject: [PATCH] Support ShouldReconnectOnError with exponential backoff up to
 10 minutes.

---
 structs.go |   2 +
 wsapi.go   | 117 ++++++++++++++++++++++++++++++++---------------------
 2 files changed, 72 insertions(+), 47 deletions(-)

diff --git a/structs.go b/structs.go
index 68866ca..762c27d 100644
--- a/structs.go
+++ b/structs.go
@@ -92,6 +92,8 @@ type Session struct {
 
 	// When nil, the session is not listening.
 	listening chan interface{}
+
+	ShouldReconnectOnError bool
 }
 
 // A VoiceRegion stores data for a specific voice region server.
diff --git a/wsapi.go b/wsapi.go
index 342efdd..ff9cfc1 100644
--- a/wsapi.go
+++ b/wsapi.go
@@ -116,9 +116,32 @@ func (s *Session) listen(wsConn *websocket.Conn, listening <-chan interface{}) {
 	for {
 		messageType, message, err := wsConn.ReadMessage()
 		if err != nil {
-			// There has been an error reading, Close() the websocket so that
-			// OnDisconnect is fired.
-			s.Close()
+			// Detect if we have been closed manually. If a Close() has already
+			// happened, the websocket we are listening on will be different to the
+			// current session.
+			s.RLock()
+			sameConnection := s.wsConn == wsConn
+			s.RUnlock()
+			if sameConnection {
+				// There has been an error reading, Close() the websocket so that
+				// OnDisconnect is fired.
+				s.Close()
+
+				// Attempt to reconnect, with expenonential backoff up to 10 minutes.
+				if s.ShouldReconnectOnError {
+					wait := time.Duration(1)
+					for {
+						if s.Open() == nil {
+							return
+						}
+						<-time.After(wait * time.Second)
+						wait *= 2
+						if wait > 600 {
+							wait = 600
+						}
+					}
+				}
+			}
 			return
 		}
 
@@ -133,6 +156,50 @@ func (s *Session) listen(wsConn *websocket.Conn, listening <-chan interface{}) {
 	return
 }
 
+type heartbeatOp struct {
+	Op   int `json:"op"`
+	Data int `json:"d"`
+}
+
+func (s *Session) sendHeartbeat(wsConn *websocket.Conn) error {
+	return wsConn.WriteJSON(heartbeatOp{1, int(time.Now().Unix())})
+}
+
+// heartbeat sends regular heartbeats to Discord so it knows the client
+// is still connected.  If you do not send these heartbeats Discord will
+// disconnect the websocket connection after a few seconds.
+func (s *Session) heartbeat(wsConn *websocket.Conn, listening <-chan interface{}, i time.Duration) {
+	if listening == nil || wsConn == nil {
+		return
+	}
+
+	s.Lock()
+	s.DataReady = true
+	s.Unlock()
+
+	// Send first heartbeat immediately because lag could put the
+	// first heartbeat outside the required heartbeat interval window.
+	err := s.sendHeartbeat(wsConn)
+	if err != nil {
+		fmt.Println("Error sending initial heartbeat:", err)
+		return
+	}
+
+	ticker := time.NewTicker(i * time.Millisecond)
+	for {
+		select {
+		case <-ticker.C:
+			err := s.sendHeartbeat(wsConn)
+			if err != nil {
+				fmt.Println("Error sending heartbeat:", err)
+				return
+			}
+		case <-listening:
+			return
+		}
+	}
+}
+
 type updateStatusGame struct {
 	Name string `json:"name"`
 }
@@ -549,50 +616,6 @@ func (s *Session) event(messageType int, message []byte) (err error) {
 	return
 }
 
-type heartbeatOp struct {
-	Op   int `json:"op"`
-	Data int `json:"d"`
-}
-
-func (s *Session) sendHeartbeat(wsConn *websocket.Conn) error {
-	return wsConn.WriteJSON(heartbeatOp{1, int(time.Now().Unix())})
-}
-
-// heartbeat sends regular heartbeats to Discord so it knows the client
-// is still connected.  If you do not send these heartbeats Discord will
-// disconnect the websocket connection after a few seconds.
-func (s *Session) heartbeat(wsConn *websocket.Conn, listening <-chan interface{}, i time.Duration) {
-	if listening == nil || wsConn == nil {
-		return
-	}
-
-	s.Lock()
-	s.DataReady = true
-	s.Unlock()
-
-	// Send first heartbeat immediately because lag could put the
-	// first heartbeat outside the required heartbeat interval window.
-	err := s.sendHeartbeat(wsConn)
-	if err != nil {
-		fmt.Println("Error sending initial heartbeat:", err)
-		return
-	}
-
-	ticker := time.NewTicker(i * time.Millisecond)
-	for {
-		select {
-		case <-ticker.C:
-			err := s.sendHeartbeat(wsConn)
-			if err != nil {
-				fmt.Println("Error sending heartbeat:", err)
-				return
-			}
-		case <-listening:
-			return
-		}
-	}
-}
-
 // ------------------------------------------------------------------------------------------------
 // Code related to voice connections that initiate over the data websocket
 // ------------------------------------------------------------------------------------------------