ratelimiter: do not run GC with nothing to do

This commit is contained in:
Jason A. Donenfeld 2018-05-21 03:18:56 +02:00
parent 84c5357cf3
commit e28d70f5b2

View File

@ -27,7 +27,7 @@ type RatelimiterEntry struct {
type Ratelimiter struct { type Ratelimiter struct {
mutex sync.RWMutex mutex sync.RWMutex
stop chan struct{} stopReset chan struct{}
tableIPv4 map[[net.IPv4len]byte]*RatelimiterEntry tableIPv4 map[[net.IPv4len]byte]*RatelimiterEntry
tableIPv6 map[[net.IPv6len]byte]*RatelimiterEntry tableIPv6 map[[net.IPv6len]byte]*RatelimiterEntry
} }
@ -36,8 +36,8 @@ func (rate *Ratelimiter) Close() {
rate.mutex.Lock() rate.mutex.Lock()
defer rate.mutex.Unlock() defer rate.mutex.Unlock()
if rate.stop != nil { if rate.stopReset != nil {
close(rate.stop) close(rate.stopReset)
} }
} }
@ -47,11 +47,11 @@ func (rate *Ratelimiter) Init() {
// stop any ongoing garbage collection routine // stop any ongoing garbage collection routine
if rate.stop != nil { if rate.stopReset != nil {
close(rate.stop) close(rate.stopReset)
} }
rate.stop = make(chan struct{}) rate.stopReset = make(chan struct{})
rate.tableIPv4 = make(map[[net.IPv4len]byte]*RatelimiterEntry) rate.tableIPv4 = make(map[[net.IPv4len]byte]*RatelimiterEntry)
rate.tableIPv6 = make(map[[net.IPv6len]byte]*RatelimiterEntry) rate.tableIPv6 = make(map[[net.IPv6len]byte]*RatelimiterEntry)
@ -59,11 +59,16 @@ func (rate *Ratelimiter) Init() {
go func() { go func() {
ticker := time.NewTicker(time.Second) ticker := time.NewTicker(time.Second)
ticker.Stop()
for { for {
select { select {
case <-rate.stop: case _, ok := <-rate.stopReset:
ticker.Stop() ticker.Stop()
return if ok {
ticker = time.NewTicker(time.Second)
} else {
return
}
case <-ticker.C: case <-ticker.C:
func() { func() {
rate.mutex.Lock() rate.mutex.Lock()
@ -84,6 +89,10 @@ func (rate *Ratelimiter) Init() {
} }
entry.mutex.Unlock() 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() rate.mutex.Lock()
if IPv4 != nil { if IPv4 != nil {
rate.tableIPv4[keyIPv4] = entry rate.tableIPv4[keyIPv4] = entry
if len(rate.tableIPv4) == 1 && len(rate.tableIPv6) == 0 {
rate.stopReset <- struct{}{}
}
} else { } else {
rate.tableIPv6[keyIPv6] = entry rate.tableIPv6[keyIPv6] = entry
if len(rate.tableIPv6) == 1 && len(rate.tableIPv4) == 0 {
rate.stopReset <- struct{}{}
}
} }
rate.mutex.Unlock() rate.mutex.Unlock()
return true return true