mirror of
https://github.com/diamondburned/arikawa.git
synced 2025-07-23 05:11:44 +00:00
Compare commits
2 commits
06fbcfcfd1
...
a25250d959
Author | SHA1 | Date | |
---|---|---|---|
|
a25250d959 | ||
|
33eb8ef315 |
|
@ -41,15 +41,15 @@ type Websocket struct {
|
|||
addr string
|
||||
closed bool
|
||||
|
||||
sendLimiter *rate.Limiter
|
||||
dialLimiter *rate.Limiter
|
||||
|
||||
// Constants. These must not be changed after the Websocket instance is used
|
||||
// once, as they are not thread-safe.
|
||||
|
||||
// Timeout for connecting and writing to the Websocket, uses default
|
||||
// WSTimeout (global).
|
||||
Timeout time.Duration
|
||||
|
||||
SendLimiter *rate.Limiter
|
||||
DialLimiter *rate.Limiter
|
||||
}
|
||||
|
||||
// New creates a default Websocket with the given address.
|
||||
|
@ -64,10 +64,10 @@ func NewCustom(conn Connection, addr string) *Websocket {
|
|||
addr: addr,
|
||||
closed: true,
|
||||
|
||||
Timeout: WSTimeout,
|
||||
sendLimiter: NewSendLimiter(),
|
||||
dialLimiter: NewDialLimiter(),
|
||||
|
||||
SendLimiter: NewSendLimiter(),
|
||||
DialLimiter: NewDialLimiter(),
|
||||
Timeout: WSTimeout,
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -80,7 +80,7 @@ func (ws *Websocket) Dial(ctx context.Context) error {
|
|||
ctx = tctx
|
||||
}
|
||||
|
||||
if err := ws.DialLimiter.Wait(ctx); err != nil {
|
||||
if err := ws.dialLimiter.Wait(ctx); err != nil {
|
||||
// Expired, fatal error
|
||||
return errors.Wrap(err, "failed to wait")
|
||||
}
|
||||
|
@ -99,6 +99,9 @@ func (ws *Websocket) Dial(ctx context.Context) error {
|
|||
|
||||
ws.closed = false
|
||||
|
||||
// Reset the send limiter.
|
||||
ws.sendLimiter = NewSendLimiter()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
|
@ -123,13 +126,20 @@ func (ws *Websocket) Send(b []byte) error {
|
|||
// SendCtx sends b over the Websocket with a deadline. It closes the internal
|
||||
// Websocket if the Send method errors out.
|
||||
func (ws *Websocket) SendCtx(ctx context.Context, b []byte) error {
|
||||
if err := ws.SendLimiter.Wait(ctx); err != nil {
|
||||
WSDebug("Waiting for the send rate limiter...")
|
||||
|
||||
if err := ws.sendLimiter.Wait(ctx); err != nil {
|
||||
WSDebug("Send rate limiter timed out.")
|
||||
return errors.Wrap(err, "SendLimiter failed")
|
||||
}
|
||||
|
||||
WSDebug("Send is passed the rate limiting. Waiting on mutex.")
|
||||
|
||||
ws.mutex.Lock()
|
||||
defer ws.mutex.Unlock()
|
||||
|
||||
WSDebug("Mutex lock acquired.")
|
||||
|
||||
if ws.closed {
|
||||
return ErrWebsocketClosed
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue