arikawa/utils/wsutil/ws.go

// Package wsutil provides abstractions around the Websocket, including rate
// limits.
package wsutil

import (
	"context"
	"log"
	"net/url"
	"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.
	WSTimeout = 5 * time.Minute
	// 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) }
	// WSExtraReadTimeout is the duration to be added to Hello, as a read
	// timeout for the websocket.
	WSExtraReadTimeout = time.Second
	// WSDebug is used for extra debug logging. This is expected to behave
	// similarly to log.Println().
	WSDebug = func(v ...interface{}) {}
)

type Event struct {
	Data []byte

	// Error is non-nil if Data is nil.
	Error error
}

type Websocket struct {
	Conn Connection
	Addr string

	// Timeout for connecting and writing to the Websocket, uses default
	// WSTimeout (global).
	Timeout time.Duration

	SendLimiter *rate.Limiter
	DialLimiter *rate.Limiter
}

func New(addr string) *Websocket {
	return NewCustom(NewConn(), addr)
}

// NewCustom creates a new undialed Websocket.
func NewCustom(conn Connection, addr string) *Websocket {
	return &Websocket{
		Conn: conn,
		Addr: addr,

		Timeout: WSTimeout,

		SendLimiter: NewSendLimiter(),
		DialLimiter: NewDialLimiter(),
	}
}

func (ws *Websocket) Dial(ctx context.Context) error {
	if ws.Timeout > 0 {
		tctx, cancel := context.WithTimeout(ctx, ws.Timeout)
		defer cancel()

		ctx = tctx
	}

	if err := ws.DialLimiter.Wait(ctx); err != nil {
		// Expired, fatal error
		return errors.Wrap(err, "failed to wait")
	}

	if err := ws.Conn.Dial(ctx, ws.Addr); err != nil {
		return errors.Wrap(err, "failed to dial")
	}

	// Reset the SendLimiter:
	ws.SendLimiter = NewSendLimiter()

	return nil
}

func (ws *Websocket) Listen() <-chan Event {
	return ws.Conn.Listen()
}

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 {
		return errors.Wrap(err, "SendLimiter failed")
	}

	return ws.Conn.Send(ctx, b)
}

func (ws *Websocket) Close() error {
	return ws.Conn.Close()
}

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()
}