From 835a23a89d004fa434f32223f377e5f69924ecf6 Mon Sep 17 00:00:00 2001
From: andrei <b1naryth1ef@gmail.com>
Date: Mon, 13 Jun 2016 15:34:40 -0700
Subject: [PATCH] Implement guild sharding

---
 discord.go |  2 ++
 structs.go |  4 ++++
 wsapi.go   | 22 +++++++++++++++++++++-
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/discord.go b/discord.go
index cd60dd9..24114d1 100644
--- a/discord.go
+++ b/discord.go
@@ -40,6 +40,8 @@ func New(args ...interface{}) (s *Session, err error) {
 		StateEnabled:           true,
 		Compress:               true,
 		ShouldReconnectOnError: true,
+		ShardID:                0,
+		NumShards:              1,
 	}
 
 	// If no arguments are passed return the empty Session interface.
diff --git a/structs.go b/structs.go
index dcf7dac..00b0c4e 100644
--- a/structs.go
+++ b/structs.go
@@ -39,6 +39,10 @@ type Session struct {
 	// Should the session request compressed websocket data.
 	Compress bool
 
+	// Sharding
+	ShardID   int
+	NumShards int
+
 	// Should state tracking be enabled.
 	// State tracking is the best way for getting the the users
 	// active guilds and the members of the guilds.
diff --git a/wsapi.go b/wsapi.go
index 6de1784..dfe31e1 100644
--- a/wsapi.go
+++ b/wsapi.go
@@ -41,6 +41,7 @@ type handshakeData struct {
 	Properties     handshakeProperties `json:"properties"`
 	LargeThreshold int                 `json:"large_threshold"`
 	Compress       bool                `json:"compress"`
+	Shard          [2]int              `json:"shard"`
 }
 
 type handshakeOp struct {
@@ -74,6 +75,16 @@ func (s *Session) Open() (err error) {
 		return
 	}
 
+	if s.NumShards <= 0 {
+		err = errors.New("NumShards must be greater or equal to 1")
+		return
+	}
+
+	if s.ShardID >= s.NumShards {
+		err = errors.New("ShardID must be less than NumShards")
+		return
+	}
+
 	if s.VoiceConnections == nil {
 		s.log(LogInformational, "creating new VoiceConnections map")
 		s.VoiceConnections = make(map[string]*VoiceConnection)
@@ -132,6 +143,7 @@ func (s *Session) Open() (err error) {
 				},
 				250,
 				s.Compress,
+				[2]int{s.ShardID, s.NumShards},
 			},
 		}
 		s.log(LogInformational, "sending identify packet to gateway")
@@ -375,7 +387,15 @@ func (s *Session) onEvent(messageType int, message []byte) {
 
 		s.log(LogInformational, "sending identify packet to gateway in response to Op9")
 		s.wsMutex.Lock()
-		err = s.wsConn.WriteJSON(handshakeOp{2, handshakeData{s.Token, handshakeProperties{runtime.GOOS, "Discordgo v" + VERSION, "", "", ""}, 250, s.Compress}})
+		err = s.wsConn.WriteJSON(handshakeOp{
+			2,
+			handshakeData{
+				s.Token,
+				handshakeProperties{runtime.GOOS, "Discordgo v" + VERSION, "", "", ""},
+				250,
+				s.Compress,
+				[2]int{s.ShardID, s.NumShards},
+			}})
 		s.wsMutex.Unlock()
 		if err != nil {
 			s.log(LogWarning, "error sending gateway identify packet, %s, %s", s.gateway, err)