add shard reservation
Use dgrs to reserve either a predefined shard ID from the config or automatically reserve a free shard ID using the shared state.
package inits
import (
+ "math/rand"
"strings"
"sync/atomic"
+ "time"
"github.com/bwmarrin/discordgo"
"github.com/bwmarrin/snowflake"
@@ -14,9 +16,14 @@ import (
"github.com/zekroTJA/shinpuru/internal/util"
"github.com/zekroTJA/shinpuru/internal/util/snowflakenodes"
"github.com/zekroTJA/shinpuru/internal/util/static"
+ "github.com/zekrotja/dgrs"
)
-func InitDiscordBotSession(container di.Container) {
+func init() {
+ rand.Seed(time.Now().UnixNano())
+}
+
+func InitDiscordBotSession(container di.Container) (release func()) {
snowflake.Epoch = static.DefEpoche
err := snowflakenodes.Setup()
if err != nil {
@@ -38,10 +45,37 @@ func InitDiscordBotSession(container di.Container) {
session.StateEnabled = false
if shardCfg := cfg.Config().Discord.Sharding; shardCfg.Total > 1 {
- if shardCfg.ID < 0 || shardCfg.ID >= shardCfg.Total {
+ st := container.Get(static.DiState).(*dgrs.State)
+
+ var id int
+ if shardCfg.AutoID {
+ d := time.Duration(rand.Int63n(int64(5 * time.Second)))
+ logrus.
+ WithField("d", d.Round(time.Millisecond).String()).
+ Info("Sleeping before retrieving shard ID")
+ time.Sleep(d)
+ if id, err = st.ReserveShard(); err != nil {
+ logrus.WithError(err).Fatal("Failed receiving alive shards from state")
+ }
+ } else {
+ id = shardCfg.ID
+ if id < 0 || id >= shardCfg.Total {
logrus.Fatalf("Shard ID must be in range [0, %d)", shardCfg.Total)
}
- session.Identify.Shard = &[2]int{shardCfg.ID, shardCfg.Total}
+ if _, err = st.ReserveShard(id); err != nil {
+ logrus.WithError(err).Fatal("Failed reserving predefined shard")
+ }
+ }
+
+ release = func() {
+ st.ReleaseShard(id)
+ }
+
+ logrus.
+ WithField("id", id).
+ WithField("total", shardCfg.Total).
+ Info("Running in sharded mode")
+ session.Identify.Shard = &[2]int{id, shardCfg.Total}
}
listenerInviteBlock := listeners.NewListenerInviteBlock(container)
