diamondburned/arikawa
1
0
Fork 0
mirror of https://github.com/diamondburned/arikawa.git synced 2025-07-12 08:16:09 +00:00
Code Issues Releases Activity

wsutil: Reset SendLimiter after Dial

Browse source
This commit is contained in:
diamondburned 2020-11-01 10:09:41 -08:00
parent bd1696ef65
commit a3aa490024
2 changed files with 15 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
utils/wsutil/conn.go
View file

@ -182,10 +182,10 @@ func startReadLoop(conn *websocket.Conn, eventCh chan<- Event) {
return return
} }
// // Check if the error is a normal one: // Check if the error is a normal one:
// if websocket.IsCloseError(err, websocket.CloseNormalClosure) { if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
// return return
// } }
// Unusual error; log and exit: // Unusual error; log and exit:
eventCh <- Event{nil, errors.Wrap(err, "WS error")} eventCh <- Event{nil, errors.Wrap(err, "WS error")}

19
utils/wsutil/ws.go
View file

@ -41,15 +41,15 @@ type Websocket struct {
addr string addr string
closed bool closed bool
sendLimiter *rate.Limiter
dialLimiter *rate.Limiter
// Constants. These must not be changed after the Websocket instance is used // Constants. These must not be changed after the Websocket instance is used
// once, as they are not thread-safe. // once, as they are not thread-safe.
// Timeout for connecting and writing to the Websocket, uses default // Timeout for connecting and writing to the Websocket, uses default
// WSTimeout (global). // WSTimeout (global).
Timeout time.Duration Timeout time.Duration
SendLimiter *rate.Limiter
DialLimiter *rate.Limiter
} }
// New creates a default Websocket with the given address. // New creates a default Websocket with the given address.
@ -64,10 +64,10 @@ func NewCustom(conn Connection, addr string) *Websocket {
addr: addr, addr: addr,
closed: true, closed: true,
Timeout: WSTimeout, sendLimiter: NewSendLimiter(),
dialLimiter: NewDialLimiter(),
SendLimiter: NewSendLimiter(), Timeout: WSTimeout,
DialLimiter: NewDialLimiter(),
} }
} }
@ -80,7 +80,7 @@ func (ws *Websocket) Dial(ctx context.Context) error {
ctx = tctx ctx = tctx
} }
if err := ws.DialLimiter.Wait(ctx); err != nil { if err := ws.dialLimiter.Wait(ctx); err != nil {
// Expired, fatal error // Expired, fatal error
return errors.Wrap(err, "failed to wait") return errors.Wrap(err, "failed to wait")
} }
@ -99,6 +99,9 @@ func (ws *Websocket) Dial(ctx context.Context) error {
ws.closed = false ws.closed = false
// Reset the send limiter.
ws.sendLimiter = NewSendLimiter()
return nil return nil
} }
@ -125,7 +128,7 @@ func (ws *Websocket) Send(b []byte) error {
func (ws *Websocket) SendCtx(ctx context.Context, b []byte) error { func (ws *Websocket) SendCtx(ctx context.Context, b []byte) error {
WSDebug("Waiting for the send rate limiter...") WSDebug("Waiting for the send rate limiter...")
if err := ws.SendLimiter.Wait(ctx); err != nil { if err := ws.sendLimiter.Wait(ctx); err != nil {
WSDebug("Send rate limiter timed out.") WSDebug("Send rate limiter timed out.")
return errors.Wrap(err, "SendLimiter failed") return errors.Wrap(err, "SendLimiter failed")
} }