1
0
Fork 0
mirror of https://github.com/diamondburned/arikawa.git synced 2024-12-12 16:35:30 +00:00
arikawa/gateway/shard/shard.go
diamondburned 7f4daccd2d
shard: Separate state's cabinet, fix return bugs
This commit makes it so that all sharded state.State instances will have
its own cabinet store on construction. This ensures that when a State is
reconnected, a Ready event won't wipe the cabinet for all other states.

This commit also fixes a bug with the Shard() and FromGuildID() getters.
2021-10-31 13:10:34 -07:00

86 lines
2.5 KiB
Go

package shard
import (
"context"
"github.com/diamondburned/arikawa/v3/gateway"
"github.com/pkg/errors"
)
// Shard defines a shard gateway interface that the shard manager can use.
type Shard interface {
Open(context.Context) error
Close() error
}
// NewShardFunc is the constructor to create a new gateway. For examples, see
// package session and state's. The constructor must manually connect the
// Manager's Rescale method appropriately.
//
// A new Gateway must not open any background resources until OpenCtx is called;
// if the gateway has never been opened, its Close method will never be called.
// During callback, the Manager is not locked, so the callback can use Manager's
// methods without deadlocking.
type NewShardFunc func(m *Manager, id *gateway.Identifier) (Shard, error)
// NewGatewayShardFunc wraps around NewGatewayShard to be compatible with
// NewShardFunc.
var NewGatewayShardFunc NewShardFunc = func(m *Manager, id *gateway.Identifier) (Shard, error) {
return NewGatewayShard(m, id), nil
}
// NewGatewayShard creates a new gateway that's plugged into the shard manager.
func NewGatewayShard(m *Manager, id *gateway.Identifier) *gateway.Gateway {
gw := gateway.NewCustomIdentifiedGateway(m.GatewayURL(), id)
gw.OnShardingRequired(m.Rescale)
return gw
}
// ShardState wraps around the Gateway interface to provide additional state.
type ShardState struct {
Shard Shard
// This is a bit wasteful: 2 constant pointers are stored here, and they
// waste GC cycles. This is unavoidable, however, since the API has to take
// in a pointer to Identifier, not IdentifyData. This is to ensure rescales
// are consistent.
ID gateway.Identifier
Opened bool
}
// ShardID returns the shard state's shard ID.
func (state ShardState) ShardID() int {
return state.ID.Shard.ShardID()
}
// OpenShards opens the gateways of the given list of shard states.
func OpenShards(ctx context.Context, shards []ShardState) error {
for i, shard := range shards {
if err := shard.Shard.Open(ctx); err != nil {
CloseShards(shards)
return errors.Wrapf(err, "failed to open shard %d/%d", i, len(shards)-1)
}
// Mark as opened so we can close them.
shards[i].Opened = true
}
return nil
}
// CloseShards closes the gateways of the given list of shard states.
func CloseShards(shards []ShardState) error {
var lastError error
for i, gw := range shards {
if gw.Opened {
if err := gw.Shard.Close(); err != nil {
lastError = err
}
shards[i].Opened = false
}
}
return lastError
}