1
0
Fork 0
mirror of https://github.com/diamondburned/arikawa.git synced 2024-11-27 17:23:00 +00:00

Gateway: Add automatic BotData connection

This commit modifies Gateway constructors to allow the user to easily
feed existing Identifier instances as well as updating those instances
to adhere to the Discord-returned gateway rate limits.

These changes should make it easier for typical bot sharding, although
automatic sharding is not implemented.
This commit is contained in:
diamondburned 2021-04-02 18:05:27 -07:00
parent 050c1b6497
commit 8e02b1652c

View file

@ -11,6 +11,7 @@ import (
"context" "context"
"net/http" "net/http"
"net/url" "net/url"
"strings"
"sync" "sync"
"time" "time"
@ -50,6 +51,7 @@ type SessionStartLimit struct {
Total int `json:"total"` Total int `json:"total"`
Remaining int `json:"remaining"` Remaining int `json:"remaining"`
ResetAfter discord.Milliseconds `json:"reset_after"` ResetAfter discord.Milliseconds `json:"reset_after"`
MaxConcurrency int `json:"max_concurrency"`
} }
// URL asks Discord for a Websocket URL to the Gateway. // URL asks Discord for a Websocket URL to the Gateway.
@ -139,13 +141,32 @@ func NewGatewayWithIntents(token string, intents ...Intents) (*Gateway, error) {
return g, nil return g, nil
} }
// NewGateway creates a new Gateway with the default stdlib JSON driver. For // NewGateway creates a new Gateway to the default Discord server.
// more information, refer to NewGatewayWithDriver.
func NewGateway(token string) (*Gateway, error) { func NewGateway(token string) (*Gateway, error) {
URL, err := URL() return NewIdentifiedGateway(DefaultIdentifier(token))
}
// NewIdentifiedGateway creates a new Gateway with the given gateway identifier
// and the default everything. Sharded bots should prefer this function for the
// shared identifier.
func NewIdentifiedGateway(id *Identifier) (*Gateway, error) {
var gatewayURL string
var botData *BotData
var err error
if strings.HasPrefix(id.Token, "Bot ") {
botData, err = BotURL(id.Token)
if err != nil {
return nil, errors.Wrap(err, "failed to get bot data")
}
gatewayURL = botData.URL
} else {
gatewayURL, err = URL()
if err != nil { if err != nil {
return nil, errors.Wrap(err, "failed to get gateway endpoint") return nil, errors.Wrap(err, "failed to get gateway endpoint")
} }
}
// Parameters for the gateway // Parameters for the gateway
param := url.Values{ param := url.Values{
@ -154,18 +175,42 @@ func NewGateway(token string) (*Gateway, error) {
} }
// Append the form to the URL // Append the form to the URL
URL += "?" + param.Encode() gatewayURL += "?" + param.Encode()
gateway := NewCustomIdentifiedGateway(gatewayURL, id)
return NewCustomGateway(URL, token), nil // Use the supplied connect rate limit, if any.
if botData != nil && botData.StartLimit != nil {
resetAt := time.Now().Add(botData.StartLimit.ResetAfter.Duration())
limiter := gateway.Identifier.IdentifyGlobalLimit
// Update the burst to be the current given time and reset it back to
// the default when the given time is reached.
limiter.SetBurst(botData.StartLimit.Remaining)
limiter.SetBurstAt(resetAt, botData.StartLimit.Total)
// Update the maximum number of identify requests allowed per 5s.
gateway.Identifier.IdentifyShortLimit.SetBurst(botData.StartLimit.MaxConcurrency)
} }
return gateway, nil
}
// NewCustomGateway creates a new Gateway with a custom gateway URL and a new
// Identifier. Most bots connecting to the official server should not use these
// custom functions.
func NewCustomGateway(gatewayURL, token string) *Gateway { func NewCustomGateway(gatewayURL, token string) *Gateway {
return NewCustomIdentifiedGateway(gatewayURL, DefaultIdentifier(token))
}
// NewCustomIdentifiedGateway creates a new Gateway with a custom gateway URL
// and a pre-existing Identifier. Refer to NewCustomGateway.
func NewCustomIdentifiedGateway(gatewayURL string, id *Identifier) *Gateway {
return &Gateway{ return &Gateway{
WS: wsutil.NewCustom(wsutil.NewConn(), gatewayURL), WS: wsutil.NewCustom(wsutil.NewConn(), gatewayURL),
WSTimeout: wsutil.WSTimeout, WSTimeout: wsutil.WSTimeout,
Events: make(chan Event, wsutil.WSBuffer), Events: make(chan Event, wsutil.WSBuffer),
Identifier: DefaultIdentifier(token), Identifier: id,
Sequence: moreatomic.NewInt64(0), Sequence: moreatomic.NewInt64(0),
ErrorLog: wsutil.WSError, ErrorLog: wsutil.WSError,