Use golang.org/x/time/rate
Instead of hand-rolling our own rate-limiter based on goroutines, use golang.org/x/time/rate.
This commit is contained in:
parent
f75ee272a4
commit
b2957c05d5
53
conn.go
53
conn.go
@ -10,6 +10,7 @@ import (
|
||||
"time"
|
||||
"unicode"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
"gopkg.in/irc.v3"
|
||||
"nhooyr.io/websocket"
|
||||
)
|
||||
@ -118,46 +119,6 @@ func (wa websocketAddr) String() string {
|
||||
return string(wa)
|
||||
}
|
||||
|
||||
type rateLimiter struct {
|
||||
C <-chan struct{}
|
||||
ticker *time.Ticker
|
||||
stopped chan struct{}
|
||||
}
|
||||
|
||||
func newRateLimiter(delay time.Duration, burst int) *rateLimiter {
|
||||
ch := make(chan struct{}, burst)
|
||||
for i := 0; i < burst; i++ {
|
||||
ch <- struct{}{}
|
||||
}
|
||||
ticker := time.NewTicker(delay)
|
||||
stopped := make(chan struct{})
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
select {
|
||||
case ch <- struct{}{}:
|
||||
// This space is intentionally left blank
|
||||
case <-stopped:
|
||||
return
|
||||
}
|
||||
case <-stopped:
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
return &rateLimiter{
|
||||
C: ch,
|
||||
ticker: ticker,
|
||||
stopped: stopped,
|
||||
}
|
||||
}
|
||||
|
||||
func (rl *rateLimiter) Stop() {
|
||||
rl.ticker.Stop()
|
||||
close(rl.stopped)
|
||||
}
|
||||
|
||||
type connOptions struct {
|
||||
Logger Logger
|
||||
RateLimitDelay time.Duration
|
||||
@ -186,15 +147,13 @@ func newConn(srv *Server, ic ircConn, options *connOptions) *conn {
|
||||
}
|
||||
|
||||
go func() {
|
||||
var rl *rateLimiter
|
||||
if options.RateLimitDelay > 0 && options.RateLimitBurst > 0 {
|
||||
rl = newRateLimiter(options.RateLimitDelay, options.RateLimitBurst)
|
||||
defer rl.Stop()
|
||||
}
|
||||
ctx, cancel := c.NewContext(context.Background())
|
||||
defer cancel()
|
||||
|
||||
rl := rate.NewLimiter(rate.Every(options.RateLimitDelay), options.RateLimitBurst)
|
||||
for msg := range outgoing {
|
||||
if rl != nil {
|
||||
<-rl.C
|
||||
if err := rl.Wait(ctx); err != nil {
|
||||
break
|
||||
}
|
||||
|
||||
if c.srv.Config().Debug {
|
||||
|
1
go.mod
1
go.mod
@ -18,6 +18,7 @@ require (
|
||||
golang.org/x/crypto v0.0.0-20211115234514-b4de73f9ece8
|
||||
golang.org/x/sys v0.0.0-20211116061358-0a5406a5449c // indirect
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 // indirect
|
||||
golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11
|
||||
google.golang.org/protobuf v1.27.1 // indirect
|
||||
gopkg.in/irc.v3 v3.1.4
|
||||
nhooyr.io/websocket v1.8.7
|
||||
|
2
go.sum
2
go.sum
@ -379,6 +379,8 @@ golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11 h1:GZokNIeuVkl3aZHJchRrr13WCsols02MLUcz1U9is6M=
|
||||
golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
|
||||
|
Loading…
Reference in New Issue
Block a user