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:
parent
050c1b6497
commit
8e02b1652c
|
@ -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,
|
||||||
|
|
Loading…
Reference in a new issue