From b2957c05d522f6c7e19fc6d50d5e85e774f31bc8 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Thu, 2 Dec 2021 23:52:29 +0100
Subject: [PATCH] Use golang.org/x/time/rate

Instead of hand-rolling our own rate-limiter based on goroutines,
use golang.org/x/time/rate.
---
 conn.go | 53 ++++++-----------------------------------------------
 go.mod  |  1 +
 go.sum  |  2 ++
 3 files changed, 9 insertions(+), 47 deletions(-)

diff --git a/conn.go b/conn.go
index 576d308..1ffa88c 100644
--- a/conn.go
+++ b/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 {
diff --git a/go.mod b/go.mod
index 304c1f0..12cccdd 100644
--- a/go.mod
+++ b/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
diff --git a/go.sum b/go.sum
index dd4f3e7..8c11e23 100644
--- a/go.sum
+++ b/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=