Changeset 398 in code for trunk/conn.go

Timestamp:

Aug 19, 2020, 5:42:33 PM (5 years ago)

Author:

contact

Message:

Implement rate limiting for upstream messages

Allow up to 10 outgoing messages in a burst, then throttle to 1 message
each 2 seconds.

Closes: ​https://todo.sr.ht/~emersion/soju/87

File:

• trunk/conn.go (modified) (3 diffs)

trunk/conn.go

@@ -107 +107 @@
 }
 
+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
+        RateLimitBurst int
+}
+
 type conn struct {
         conn   ircConn
@@ -117 +163 @@
 }
 
-func newConn(srv *Server, ic ircConn, logger Logger) *conn {
+func newConn(srv *Server, ic ircConn, options *connOptions) *conn {
         outgoing := make(chan *irc.Message, 64)
         c := &conn{
@@ -123 +169 @@
                 srv:      srv,
                 outgoing: outgoing,
-                logger:   logger,
+                logger:   options.Logger,
         }
 
         go func() {
+                var rl *rateLimiter
+                if options.RateLimitDelay > 0 && options.RateLimitBurst > 0 {
+                        rl = newRateLimiter(options.RateLimitDelay, options.RateLimitBurst)
+                        defer rl.Stop()
+                }
+
                 for msg := range outgoing {
+                        if rl != nil {
+                                <-rl.C
+                        }
+
                         if c.srv.Debug {
                                 c.logger.Printf("sent: %v", msg)