Add ResumeShard method and ShardState struct to manage shard session and sequence (#460)

Author: topi314

bot/event_manager.go

@@ -35,7 +35,7 @@ type EventManager interface {
 	RemoveEventListeners(eventListeners ...EventListener)
 
 	// HandleGatewayEvent calls the correct GatewayEventHandler for the payload
-	HandleGatewayEvent(gatewayEventType gateway.EventType, sequenceNumber int, shardID int, event gateway.EventData)
+	HandleGatewayEvent(gateway gateway.Gateway, eventType gateway.EventType, sequenceNumber int, event gateway.EventData)
 
 	// HandleHTTPEvent calls the HTTPServerEventHandler for the payload
 	HandleHTTPEvent(respondFunc httpserver.RespondFunc, event httpserver.EventInteractionCreate)
@@ -91,7 +91,7 @@ type GatewayEventHandler interface {
 	HandleGatewayEvent(client *Client, sequenceNumber int, shardID int, event gateway.EventData)
 }
 
-// NewGatewayEventHandler returns a new GatewayEventHandler for the given GatewayEventType and handler func
+// NewGatewayEventHandler returns a new GatewayEventHandler for the given gateway.EventType and handler func
 func NewGatewayEventHandler[T gateway.EventData](eventType gateway.EventType, handleFunc func(client *Client, sequenceNumber int, shardID int, event T)) GatewayEventHandler {
 	return &genericGatewayEventHandler[T]{eventType: eventType, handleFunc: handleFunc}
 }
@@ -128,13 +128,13 @@ type eventManagerImpl struct {
 	httpServerHandler  HTTPServerEventHandler
 }
 
-func (e *eventManagerImpl) HandleGatewayEvent(gatewayEventType gateway.EventType, sequenceNumber int, shardID int, event gateway.EventData) {
+func (e *eventManagerImpl) HandleGatewayEvent(gateway gateway.Gateway, eventType gateway.EventType, sequenceNumber int, event gateway.EventData) {
 	e.mu.Lock()
 	defer e.mu.Unlock()
-	if handler, ok := e.gatewayHandlers[gatewayEventType]; ok {
-		handler.HandleGatewayEvent(e.client, sequenceNumber, shardID, event)
+	if handler, ok := e.gatewayHandlers[eventType]; ok {
+		handler.HandleGatewayEvent(e.client, sequenceNumber, gateway.ShardID(), event)
 	} else {
-		e.logger.Warn("no handler for Gateway event found", slog.Any("event_type", gatewayEventType))
+		e.logger.Warn("no handler for Gateway event found", slog.Any("event_type", eventType))
 	}
 }
 

gateway/gateway.go

@@ -89,7 +89,7 @@ const (
 
 type (
 	// EventHandlerFunc is a function that is called when an event is received.
-	EventHandlerFunc func(gatewayEventType EventType, sequenceNumber int, shardID int, event EventData)
+	EventHandlerFunc func(gateway Gateway, eventType EventType, sequenceNumber int, event EventData)
 
 	// CreateFunc is a type that is used to create a new Gateway(s).
 	CreateFunc func(token string, eventHandlerFunc EventHandlerFunc, closeHandlerFUnc CloseHandlerFunc, opts ...ConfigOpt) Gateway
@@ -579,7 +579,7 @@ loop:
 
 			// push message to the command manager
 			if g.config.EnableRawEvents {
-				g.eventHandlerFunc(EventTypeRaw, message.S, g.config.ShardID, EventRaw{
+				g.eventHandlerFunc(g, EventTypeRaw, message.S, EventRaw{
 					EventType: message.T,
 					Payload:   bytes.NewReader(message.RawD),
 				})
@@ -589,7 +589,7 @@ loop:
 				g.config.Logger.Debug("unknown event received", slog.String("event", string(message.T)), slog.String("data", string(unknownEvent)))
 				continue
 			}
-			g.eventHandlerFunc(message.T, message.S, g.config.ShardID, eventData)
+			g.eventHandlerFunc(g, message.T, message.S, eventData)
 
 		case OpcodeHeartbeat:
 			g.sendHeartbeat()
@@ -622,7 +622,7 @@ loop:
 
 		case OpcodeHeartbeatACK:
 			newHeartbeat := time.Now().UTC()
-			g.eventHandlerFunc(EventTypeHeartbeatAck, message.S, g.config.ShardID, EventHeartbeatAck{
+			g.eventHandlerFunc(g, EventTypeHeartbeatAck, message.S, EventHeartbeatAck{
 				LastHeartbeat: g.lastHeartbeatReceived,
 				NewHeartbeat:  newHeartbeat,
 			})

sharding/shard_manager.go

@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"iter"
 	"log/slog"
+	"maps"
+	"slices"
 	"sync"
 
 	"github.com/disgoorg/snowflake/v2"
@@ -29,6 +31,9 @@ type ShardManager interface {
 	// OpenShard opens a specific shard.
 	OpenShard(ctx context.Context, shardID int) error
 
+	// ResumeShard resumes a specific shard with the given sessionID and sequence.
+	ResumeShard(ctx context.Context, shardID int, sessionID string, sequence int) error
+
 	// CloseShard closes a specific shard.
 	CloseShard(ctx context.Context, shardID int)
 
@@ -81,10 +86,8 @@ func (m *shardManagerImpl) closeHandler(shard gateway.Gateway, err error, _ bool
 	shard.Close(context.TODO())
 
 	m.shardsMu.Lock()
-	defer m.shardsMu.Unlock()
-
 	delete(m.shards, shard.ShardID())
-	delete(m.config.ShardIDs, shard.ShardID())
+	defer m.shardsMu.Unlock()
 
 	newShardCount := shard.ShardCount() * m.config.ShardSplitCount
 	if newShardCount > m.config.ShardCount {
@@ -99,20 +102,25 @@ func (m *shardManagerImpl) closeHandler(shard gateway.Gateway, err error, _ bool
 	}
 
 	var wg sync.WaitGroup
-	for i := range newShardIDs {
-		shardID := newShardIDs[i]
+	for _, shardID := range newShardIDs {
 		wg.Add(1)
+
 		go func() {
 			defer wg.Done()
-			if err := m.config.RateLimiter.WaitBucket(context.TODO(), shardID); err != nil {
+
+			if err := m.config.RateLimiter.WaitBucket(context.Background(), shardID); err != nil {
 				m.config.Logger.Error("failed to wait shard bucket", slog.Any("err", err), slog.Int("shard_id", shardID))
 				return
 			}
 			defer m.config.RateLimiter.UnlockBucket(shardID)
 
 			newShard := m.config.GatewayCreateFunc(m.token, m.eventHandlerFunc, m.closeHandler, append(m.config.GatewayConfigOpts, gateway.WithShardID(shardID), gateway.WithShardCount(newShardCount))...)
+
+			m.shardsMu.Lock()
 			m.shards[shardID] = newShard
-			if err := newShard.Open(context.TODO()); err != nil {
+			m.shardsMu.Unlock()
+
+			if err := newShard.Open(context.Background()); err != nil {
 				m.config.Logger.Error("failed to re shard", slog.Any("err", err), slog.Int("shard_id", shardID))
 			}
 		}()
@@ -122,27 +130,41 @@ func (m *shardManagerImpl) closeHandler(shard gateway.Gateway, err error, _ bool
 }
 
 func (m *shardManagerImpl) Open(ctx context.Context) {
-	m.config.Logger.Debug("opening shards", slog.String("shard_ids", fmt.Sprint(m.config.ShardIDs)))
-	var wg sync.WaitGroup
+	m.config.Logger.Debug("opening shards", slog.String("shard_ids", fmt.Sprint(slices.Collect(maps.Keys(m.config.ShardIDs)))))
 
-	m.shardsMu.Lock()
-	defer m.shardsMu.Unlock()
-	for shardID := range m.config.ShardIDs {
-		if _, ok := m.shards[shardID]; ok {
+	var wg sync.WaitGroup
+	for shardID, shardState := range m.config.ShardIDs {
+		m.shardsMu.Lock()
+		_, ok := m.shards[shardID]
+		m.shardsMu.Unlock()
+		if ok {
 			continue
 		}
 
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
+
 			if err := m.config.RateLimiter.WaitBucket(ctx, shardID); err != nil {
 				m.config.Logger.Error("failed to wait shard bucket", slog.Any("err", err), slog.Int("shard_id", shardID))
 				return
 			}
 			defer m.config.RateLimiter.UnlockBucket(shardID)
 
-			shard := m.config.GatewayCreateFunc(m.token, m.eventHandlerFunc, m.closeHandler, append(m.config.GatewayConfigOpts, gateway.WithShardID(shardID), gateway.WithShardCount(m.config.ShardCount))...)
+			opts := append(m.config.GatewayConfigOpts, gateway.WithShardID(shardID), gateway.WithShardCount(m.config.ShardCount))
+			if shardState.SessionID != "" {
+				opts = append(opts, gateway.WithSessionID(shardState.SessionID))
+			}
+			if shardState.Sequence != 0 {
+				opts = append(opts, gateway.WithSequence(shardState.Sequence))
+			}
+
+			shard := m.config.GatewayCreateFunc(m.token, m.eventHandlerFunc, m.closeHandler, opts...)
+
+			m.shardsMu.Lock()
 			m.shards[shardID] = shard
+			m.shardsMu.Unlock()
+
 			if err := shard.Open(ctx); err != nil {
 				m.config.Logger.Error("failed to open shard", slog.Any("err", err), slog.Int("shard_id", shardID))
 			}
@@ -152,40 +174,52 @@ func (m *shardManagerImpl) Open(ctx context.Context) {
 }
 
 func (m *shardManagerImpl) Close(ctx context.Context) {
-	m.config.Logger.Debug("closing shards", slog.String("shard_ids", fmt.Sprint(m.config.ShardIDs)))
+	m.config.Logger.Debug("closing shards", slog.String("shard_ids", fmt.Sprint(slices.Collect(maps.Keys(m.shards)))))
 	var wg sync.WaitGroup
 
 	m.shardsMu.Lock()
 	defer m.shardsMu.Unlock()
-	for shardID := range m.shards {
-		shard := m.shards[shardID]
-		delete(m.shards, shardID)
+	for _, shard := range m.shards {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
 			shard.Close(ctx)
 		}()
 	}
 	wg.Wait()
+	m.shards = map[int]gateway.Gateway{}
 }
 
 func (m *shardManagerImpl) OpenShard(ctx context.Context, shardID int) error {
-	return m.openShard(ctx, shardID, m.config.ShardCount)
+	return m.openShard(ctx, shardID, m.config.ShardCount, "", 0)
+}
+
+func (m *shardManagerImpl) ResumeShard(ctx context.Context, shardID int, sessionID string, sequence int) error {
+	return m.openShard(ctx, shardID, m.config.ShardCount, sessionID, sequence)
 }
 
-func (m *shardManagerImpl) openShard(ctx context.Context, shardID int, shardCount int) error {
-	m.config.Logger.Debug("opening shard", slog.Int("shard_id", shardID))
+func (m *shardManagerImpl) openShard(ctx context.Context, shardID int, shardCount int, sessionID string, sequence int) error {
+	m.config.Logger.Debug("opening shard", slog.Int("shard_id", shardID), slog.Int("shard_count", shardCount), slog.String("session_id", sessionID), slog.Int("sequence", sequence))
 
 	if err := m.config.RateLimiter.WaitBucket(ctx, shardID); err != nil {
 		return err
 	}
 	defer m.config.RateLimiter.UnlockBucket(shardID)
-	shard := m.config.GatewayCreateFunc(m.token, m.eventHandlerFunc, m.closeHandler, append(m.config.GatewayConfigOpts, gateway.WithShardID(shardID), gateway.WithShardCount(shardCount))...)
+
+	opts := append(m.config.GatewayConfigOpts, gateway.WithShardID(shardID), gateway.WithShardCount(shardCount))
+	if sessionID != "" {
+		opts = append(opts, gateway.WithSessionID(sessionID))
+	}
+	if sequence != 0 {
+		opts = append(opts, gateway.WithSequence(sequence))
+	}
+
+	shard := m.config.GatewayCreateFunc(m.token, m.eventHandlerFunc, m.closeHandler, opts...)
 
 	m.shardsMu.Lock()
-	defer m.shardsMu.Unlock()
-	m.config.ShardIDs[shardID] = struct{}{}
 	m.shards[shardID] = shard
+	defer m.shardsMu.Unlock()
+
 	return shard.Open(ctx)
 }
 
@@ -220,6 +254,7 @@ func (m *shardManagerImpl) Shards() iter.Seq[gateway.Gateway] {
 	return func(yield func(gateway.Gateway) bool) {
 		m.shardsMu.Lock()
 		defer m.shardsMu.Unlock()
+
 		for _, shard := range m.shards {
 			if !yield(shard) {
 				return

sharding/shard_manager_config.go

@@ -14,11 +14,20 @@ func defaultConfig() config {
 	}
 }
 
+// ShardState is used to tell a [gateway.Gateway] managed by the [ShardManager] which session & sequence it should use when starting the shard.
+// This is useful for resuming shards when using the [ShardManager].
+type ShardState struct {
+	// SessionID is the session ID of the shard. This is used to resume the shard.
+	SessionID string
+	// Sequence is the sequence number of the shard. This is used to resume the shard.
+	Sequence int
+}
+
 type config struct {
 	// Logger is the logger of the ShardManager. Defaults to log.Default()
 	Logger *slog.Logger
 	// ShardIDs is a map of shardIDs the ShardManager should manage. Leave this nil to manage all shards.
-	ShardIDs map[int]struct{}
+	ShardIDs map[int]ShardState
 	// ShardCount is the total shard count of the ShardManager. Leave this at 0 to let Discord calculate the shard count for you.
 	ShardCount int
 	// ShardSplitCount is the count a shard should be split into if it is too large. This is only used if AutoScaling is enabled.
@@ -63,13 +72,20 @@ func WithLogger(logger *slog.Logger) ConfigOpt {
 // WithShardIDs sets the shardIDs the ShardManager should manage.
 func WithShardIDs(shardIDs ...int) ConfigOpt {
 	return func(config *config) {
-		config.ShardIDs = map[int]struct{}{}
+		config.ShardIDs = map[int]ShardState{}
 		for _, shardID := range shardIDs {
-			config.ShardIDs[shardID] = struct{}{}
+			config.ShardIDs[shardID] = ShardState{}
 		}
 	}
 }
 
+// WithShardIDsWithStates sets the shardIDs and their [ShardState] the ShardManager should manage.
+func WithShardIDsWithStates(shards map[int]ShardState) ConfigOpt {
+	return func(config *config) {
+		config.ShardIDs = shards
+	}
+}
+
 // WithShardCount sets the shard count of the ShardManager.
 func WithShardCount(shardCount int) ConfigOpt {
 	return func(config *config) {

voice/conn.go

@@ -194,7 +194,7 @@ func (c *connImpl) HandleVoiceServerUpdate(update botgateway.EventVoiceServerUpd
 	}()
 }
 
-func (c *connImpl) handleMessage(op Opcode, data GatewayMessageData) {
+func (c *connImpl) handleMessage(gateway Gateway, op Opcode, data GatewayMessageData) {
 	switch d := data.(type) {
 	case GatewayMessageDataReady:
 		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
@@ -238,7 +238,7 @@ func (c *connImpl) handleMessage(op Opcode, data GatewayMessageData) {
 		}
 	}
 	if c.config.EventHandlerFunc != nil {
-		c.config.EventHandlerFunc(op, data)
+		c.config.EventHandlerFunc(c.gateway, op, data)
 	}
 }
 

voice/gateway.go

@@ -46,7 +46,7 @@ const (
 
 type (
 	// EventHandlerFunc is a function that handles a voice gateway event.
-	EventHandlerFunc func(opCode Opcode, data GatewayMessageData)
+	EventHandlerFunc func(gateway Gateway, opCode Opcode, data GatewayMessageData)
 
 	// CloseHandlerFunc is a function that handles a voice gateway close.
 	CloseHandlerFunc func(gateway Gateway, err error, reconnect bool)
@@ -313,7 +313,7 @@ loop:
 			}
 			g.lastHeartbeatReceived = time.Now().UTC()
 		}
-		g.eventHandlerFunc(message.Op, message.D)
+		g.eventHandlerFunc(g, message.Op, message.D)
 	}
 }