From e28d70f5b2a00af75aa9f283e31c7da31a2913df Mon Sep 17 00:00:00 2001
From: "Jason A. Donenfeld" <Jason@zx2c4.com>
Date: Mon, 21 May 2018 03:18:56 +0200
Subject: [PATCH] ratelimiter: do not run GC with nothing to do

---
 ratelimiter/ratelimiter.go | 31 +++++++++++++++++++++++--------
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/ratelimiter/ratelimiter.go b/ratelimiter/ratelimiter.go
index 95b547e..5a844da 100644
--- a/ratelimiter/ratelimiter.go
+++ b/ratelimiter/ratelimiter.go
@@ -27,7 +27,7 @@ type RatelimiterEntry struct {
 
 type Ratelimiter struct {
 	mutex     sync.RWMutex
-	stop      chan struct{}
+	stopReset chan struct{}
 	tableIPv4 map[[net.IPv4len]byte]*RatelimiterEntry
 	tableIPv6 map[[net.IPv6len]byte]*RatelimiterEntry
 }
@@ -36,8 +36,8 @@ func (rate *Ratelimiter) Close() {
 	rate.mutex.Lock()
 	defer rate.mutex.Unlock()
 
-	if rate.stop != nil {
-		close(rate.stop)
+	if rate.stopReset != nil {
+		close(rate.stopReset)
 	}
 }
 
@@ -47,11 +47,11 @@ func (rate *Ratelimiter) Init() {
 
 	// stop any ongoing garbage collection routine
 
-	if rate.stop != nil {
-		close(rate.stop)
+	if rate.stopReset != nil {
+		close(rate.stopReset)
 	}
 
-	rate.stop = make(chan struct{})
+	rate.stopReset = make(chan struct{})
 	rate.tableIPv4 = make(map[[net.IPv4len]byte]*RatelimiterEntry)
 	rate.tableIPv6 = make(map[[net.IPv6len]byte]*RatelimiterEntry)
 
@@ -59,11 +59,16 @@ func (rate *Ratelimiter) Init() {
 
 	go func() {
 		ticker := time.NewTicker(time.Second)
+		ticker.Stop()
 		for {
 			select {
-			case <-rate.stop:
+			case _, ok := <-rate.stopReset:
 				ticker.Stop()
-				return
+				if ok {
+					ticker = time.NewTicker(time.Second)
+				} else {
+					return
+				}
 			case <-ticker.C:
 				func() {
 					rate.mutex.Lock()
@@ -84,6 +89,10 @@ func (rate *Ratelimiter) Init() {
 						}
 						entry.mutex.Unlock()
 					}
+
+					if len(rate.tableIPv4) == 0 && len(rate.tableIPv6) == 0 {
+						ticker.Stop()
+					}
 				}()
 			}
 		}
@@ -121,8 +130,14 @@ func (rate *Ratelimiter) Allow(ip net.IP) bool {
 		rate.mutex.Lock()
 		if IPv4 != nil {
 			rate.tableIPv4[keyIPv4] = entry
+			if len(rate.tableIPv4) == 1 && len(rate.tableIPv6) == 0 {
+				rate.stopReset <- struct{}{}
+			}
 		} else {
 			rate.tableIPv6[keyIPv6] = entry
+			if len(rate.tableIPv6) == 1 && len(rate.tableIPv4) == 0 {
+				rate.stopReset <- struct{}{}
+			}
 		}
 		rate.mutex.Unlock()
 		return true