1
0
Fork 0
mirror of https://github.com/diamondburned/arikawa.git synced 2024-12-12 08:25:10 +00:00
arikawa/utils/wsutil/ws.go

121 lines
2.5 KiB
Go
Raw Normal View History

2020-01-15 04:56:50 +00:00
// Package wsutil provides abstractions around the Websocket, including rate
// limits.
2020-01-09 05:24:45 +00:00
package wsutil
import (
"context"
"log"
"net/url"
2020-01-09 05:24:45 +00:00
"time"
"github.com/pkg/errors"
"golang.org/x/time/rate"
)
var (
// WSTimeout is the timeout for connecting and writing to the Websocket,
// before Gateway cancels and fails.
2020-08-20 21:15:52 +00:00
WSTimeout = 30 * time.Second
// WSBuffer is the size of the Event channel. This has to be at least 1 to
// make space for the first Event: Ready or Resumed.
WSBuffer = 10
// WSError is the default error handler
WSError = func(err error) { log.Println("Gateway error:", err) }
// WSDebug is used for extra debug logging. This is expected to behave
// similarly to log.Println().
WSDebug = func(v ...interface{}) {}
)
2020-01-09 05:24:45 +00:00
type Event struct {
Data []byte
// Error is non-nil if Data is nil.
Error error
}
type Websocket struct {
2020-01-15 04:43:34 +00:00
Conn Connection
Addr string
2020-01-09 05:24:45 +00:00
// Timeout for connecting and writing to the Websocket, uses default
// WSTimeout (global).
Timeout time.Duration
2020-01-15 04:43:34 +00:00
SendLimiter *rate.Limiter
DialLimiter *rate.Limiter
2020-01-09 05:24:45 +00:00
}
func New(addr string) *Websocket {
return NewCustom(NewConn(), addr)
2020-01-15 04:43:34 +00:00
}
2020-01-09 05:24:45 +00:00
2020-01-15 04:43:34 +00:00
// NewCustom creates a new undialed Websocket.
func NewCustom(conn Connection, addr string) *Websocket {
return &Websocket{
2020-01-15 04:43:34 +00:00
Conn: conn,
Addr: addr,
Timeout: WSTimeout,
2020-01-15 04:43:34 +00:00
SendLimiter: NewSendLimiter(),
DialLimiter: NewDialLimiter(),
2020-01-09 05:24:45 +00:00
}
}
func (ws *Websocket) Dial(ctx context.Context) error {
if ws.Timeout > 0 {
tctx, cancel := context.WithTimeout(ctx, ws.Timeout)
defer cancel()
ctx = tctx
}
2020-01-15 04:43:34 +00:00
if err := ws.DialLimiter.Wait(ctx); err != nil {
// Expired, fatal error
2020-05-16 21:14:49 +00:00
return errors.Wrap(err, "failed to wait")
2020-01-15 04:43:34 +00:00
}
2020-01-09 05:24:45 +00:00
2020-01-15 04:43:34 +00:00
if err := ws.Conn.Dial(ctx, ws.Addr); err != nil {
2020-05-16 21:14:49 +00:00
return errors.Wrap(err, "failed to dial")
2020-01-09 05:24:45 +00:00
}
2020-01-15 04:43:34 +00:00
return nil
2020-01-09 05:24:45 +00:00
}
func (ws *Websocket) Listen() <-chan Event {
2020-01-29 03:54:22 +00:00
return ws.Conn.Listen()
2020-01-09 05:24:45 +00:00
}
func (ws *Websocket) Send(b []byte) error {
return ws.SendCtx(context.Background(), b)
}
func (ws *Websocket) SendCtx(ctx context.Context, b []byte) error {
if err := ws.SendLimiter.Wait(ctx); err != nil {
2020-01-09 05:24:45 +00:00
return errors.Wrap(err, "SendLimiter failed")
}
return ws.Conn.Send(ctx, b)
}
func (ws *Websocket) Close() error {
return ws.Conn.Close()
2020-01-15 04:43:34 +00:00
}
func InjectValues(rawurl string, values url.Values) string {
u, err := url.Parse(rawurl)
if err != nil {
// Unknown URL, return as-is.
return rawurl
}
// Append additional parameters:
var q = u.Query()
for k, v := range values {
q[k] = append(q[k], v...)
}
u.RawQuery = q.Encode()
return u.String()
2020-01-09 05:24:45 +00:00
}