lib/connections: Allow on the fly changes to rate limits (fixes #3846)

Also replaces github.com/juju/ratelimit with golang.org/x/time/rate as
the latter supports changing the rate on the fly.

GitHub-Pull-Request: https://github.com/syncthing/syncthing/pull/3862

cmd/stdiscosrv/querysrv.go

@@ -19,9 +19,9 @@ import (
 	"time"
 
 	"github.com/golang/groupcache/lru"
-	"github.com/juju/ratelimit"
 	"github.com/syncthing/syncthing/lib/protocol"
 	"golang.org/x/net/context"
+	"golang.org/x/time/rate"
 )
 
 type querysrv struct {
@@ -373,14 +373,14 @@ func (s *querysrv) limit(remote net.IP) bool {
 
 	bkt, ok := s.limiter.Get(key)
 	if ok {
-		bkt := bkt.(*ratelimit.Bucket)
-		if bkt.TakeAvailable(1) != 1 {
+		bkt := bkt.(*rate.Limiter)
+		if !bkt.Allow() {
 			// Rate limit exceeded; ignore packet
 			return true
 		}
 	} else {
-		// One packet per ten seconds average rate, burst ten packets
-		s.limiter.Add(key, ratelimit.NewBucket(10*time.Second/time.Duration(limitAvg), int64(limitBurst)))
+		// limitAvg is in packets per ten seconds.
+		s.limiter.Add(key, rate.NewLimiter(rate.Limit(limitAvg)/10, limitBurst))
 	}
 
 	return false

cmd/strelaypoolsrv/main.go

@@ -23,14 +23,12 @@ import (
 	"time"
 
 	"github.com/golang/groupcache/lru"
-	"github.com/juju/ratelimit"
-
 	"github.com/oschwald/geoip2-golang"
-
 	"github.com/syncthing/syncthing/cmd/strelaypoolsrv/auto"
 	"github.com/syncthing/syncthing/lib/relay/client"
 	"github.com/syncthing/syncthing/lib/sync"
 	"github.com/syncthing/syncthing/lib/tlsutil"
+	"golang.org/x/time/rate"
 )
 
 type location struct {
@@ -65,12 +63,12 @@ var (
 	dir            string
 	evictionTime   = time.Hour
 	debug          bool
-	getLRUSize           = 10 << 10
-	getLimitBurst  int64 = 10
-	getLimitAvg          = 1
-	postLRUSize          = 1 << 10
-	postLimitBurst int64 = 2
-	postLimitAvg         = 1
+	getLRUSize     = 10 << 10
+	getLimitBurst  = 10
+	getLimitAvg    = 1
+	postLRUSize    = 1 << 10
+	postLimitBurst = 2
+	postLimitAvg   = 1
 	getLimit       time.Duration
 	postLimit      time.Duration
 	permRelaysFile string
@@ -99,10 +97,10 @@ func main() {
 	flag.DurationVar(&evictionTime, "eviction", evictionTime, "After how long the relay is evicted")
 	flag.IntVar(&getLRUSize, "get-limit-cache", getLRUSize, "Get request limiter cache size")
 	flag.IntVar(&getLimitAvg, "get-limit-avg", 2, "Allowed average get request rate, per 10 s")
-	flag.Int64Var(&getLimitBurst, "get-limit-burst", getLimitBurst, "Allowed burst get requests")
+	flag.IntVar(&getLimitBurst, "get-limit-burst", getLimitBurst, "Allowed burst get requests")
 	flag.IntVar(&postLRUSize, "post-limit-cache", postLRUSize, "Post request limiter cache size")
 	flag.IntVar(&postLimitAvg, "post-limit-avg", 2, "Allowed average post request rate, per minute")
-	flag.Int64Var(&postLimitBurst, "post-limit-burst", postLimitBurst, "Allowed burst post requests")
+	flag.IntVar(&postLimitBurst, "post-limit-burst", postLimitBurst, "Allowed burst post requests")
 	flag.StringVar(&permRelaysFile, "perm-relays", "", "Path to list of permanent relays")
 	flag.StringVar(&ipHeader, "ip-header", "", "Name of header which holds clients ip:port. Only meaningful when running behind a reverse proxy.")
 	flag.StringVar(&geoipPath, "geoip", "GeoLite2-City.mmdb", "Path to GeoLite2-City database")
@@ -446,7 +444,7 @@ func evict(relay relay) func() {
 	}
 }
 
-func limit(addr string, cache *lru.Cache, lock sync.RWMutex, rate time.Duration, burst int64) bool {
+func limit(addr string, cache *lru.Cache, lock sync.RWMutex, intv time.Duration, burst int) bool {
 	host, _, err := net.SplitHostPort(addr)
 	if err != nil {
 		return false
@@ -456,14 +454,14 @@ func limit(addr string, cache *lru.Cache, lock sync.RWMutex, rate time.Duration,
 	bkt, ok := cache.Get(host)
 	lock.RUnlock()
 	if ok {
-		bkt := bkt.(*ratelimit.Bucket)
-		if bkt.TakeAvailable(1) != 1 {
+		bkt := bkt.(*rate.Limiter)
+		if !bkt.Allow() {
 			// Rate limit
 			return true
 		}
 	} else {
 		lock.Lock()
-		cache.Add(host, ratelimit.NewBucket(rate, burst))
+		cache.Add(host, rate.NewLimiter(rate.Every(intv), burst))
 		lock.Unlock()
 	}
 	return false

cmd/strelaysrv/main.go

@@ -20,10 +20,10 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/juju/ratelimit"
 	"github.com/syncthing/syncthing/lib/osutil"
 	"github.com/syncthing/syncthing/lib/relay/protocol"
 	"github.com/syncthing/syncthing/lib/tlsutil"
+	"golang.org/x/time/rate"
 
 	"github.com/syncthing/syncthing/lib/config"
 	"github.com/syncthing/syncthing/lib/nat"
@@ -68,8 +68,8 @@ var (
 	globalLimitBps  int
 	overLimit       int32
 	descriptorLimit int64
-	sessionLimiter  *ratelimit.Bucket
-	globalLimiter   *ratelimit.Bucket
+	sessionLimiter  *rate.Limiter
+	globalLimiter   *rate.Limiter
 
 	statusAddr       string
 	poolAddrs        string
@@ -215,10 +215,10 @@ func main() {
 	}
 
 	if sessionLimitBps > 0 {
-		sessionLimiter = ratelimit.NewBucketWithRate(float64(sessionLimitBps), int64(2*sessionLimitBps))
+		sessionLimiter = rate.NewLimiter(rate.Limit(sessionLimitBps), 2*sessionLimitBps)
 	}
 	if globalLimitBps > 0 {
-		globalLimiter = ratelimit.NewBucketWithRate(float64(globalLimitBps), int64(2*globalLimitBps))
+		globalLimiter = rate.NewLimiter(rate.Limit(globalLimitBps), 2*globalLimitBps)
 	}
 
 	if statusAddr != "" {

cmd/strelaysrv/session.go

@@ -7,15 +7,16 @@ import (
 	"encoding/hex"
 	"fmt"
 	"log"
+	"math"
 	"net"
 	"sync"
 	"sync/atomic"
 	"time"
 
-	"github.com/juju/ratelimit"
-	"github.com/syncthing/syncthing/lib/relay/protocol"
+	"golang.org/x/time/rate"
 
 	syncthingprotocol "github.com/syncthing/syncthing/lib/protocol"
+	"github.com/syncthing/syncthing/lib/relay/protocol"
 )
 
 var (
@@ -26,7 +27,7 @@ var (
 	bytesProxied    int64
 )
 
-func newSession(serverid, clientid syncthingprotocol.DeviceID, sessionRateLimit, globalRateLimit *ratelimit.Bucket) *session {
+func newSession(serverid, clientid syncthingprotocol.DeviceID, sessionRateLimit, globalRateLimit *rate.Limiter) *session {
 	serverkey := make([]byte, 32)
 	_, err := rand.Read(serverkey)
 	if err != nil {
@@ -108,7 +109,7 @@ type session struct {
 	clientkey []byte
 	clientid  syncthingprotocol.DeviceID
 
-	rateLimit func(bytes int64)
+	rateLimit func(bytes int)
 
 	connsChan chan net.Conn
 	conns     []net.Conn
@@ -268,7 +269,7 @@ func (s *session) proxy(c1, c2 net.Conn) error {
 		}
 
 		if s.rateLimit != nil {
-			s.rateLimit(int64(n))
+			s.rateLimit(n)
 		}
 
 		c2.SetWriteDeadline(time.Now().Add(networkTimeout))
@@ -283,7 +284,7 @@ func (s *session) String() string {
 	return fmt.Sprintf("<%s/%s>", hex.EncodeToString(s.clientkey)[:5], hex.EncodeToString(s.serverkey)[:5])
 }
 
-func makeRateLimitFunc(sessionRateLimit, globalRateLimit *ratelimit.Bucket) func(int64) {
+func makeRateLimitFunc(sessionRateLimit, globalRateLimit *rate.Limiter) func(int) {
 	// This may be a case of super duper premature optimization... We build an
 	// optimized function to do the rate limiting here based on what we need
 	// to do and then use it in the loop.
@@ -298,29 +299,55 @@ func makeRateLimitFunc(sessionRateLimit, globalRateLimit *ratelimit.Bucket) func
 
 	if sessionRateLimit == nil {
 		// We only have a global limiter
-		return func(bytes int64) {
-			globalRateLimit.Wait(bytes)
+		return func(bytes int) {
+			take(bytes, globalRateLimit)
 		}
 	}
 
 	if globalRateLimit == nil {
 		// We only have a session limiter
-		return func(bytes int64) {
-			sessionRateLimit.Wait(bytes)
+		return func(bytes int) {
+			take(bytes, sessionRateLimit)
 		}
 	}
 
 	// We have both. Queue the bytes on both the global and session specific
-	// rate limiters. Wait for both in parallell, so that the actual send
-	// happens when both conditions are satisfied. In practice this just means
-	// wait the longer of the two times.
-	return func(bytes int64) {
-		t0 := sessionRateLimit.Take(bytes)
-		t1 := globalRateLimit.Take(bytes)
-		if t0 > t1 {
-			time.Sleep(t0)
-		} else {
-			time.Sleep(t1)
-		}
+	// rate limiters.
+	return func(bytes int) {
+		take(bytes, sessionRateLimit, globalRateLimit)
+	}
+}
+
+// take is a utility function to consume tokens from a set of rate.Limiters.
+// Tokens are consumed in parallel on all limiters, respecting their
+// individual burst sizes.
+func take(tokens int, ls ...*rate.Limiter) {
+	// minBurst is the smallest burst size supported by all limiters.
+	minBurst := int(math.MaxInt32)
+	for _, l := range ls {
+		if burst := l.Burst(); burst < minBurst {
+			minBurst = burst
+		}
+	}
+
+	for tokens > 0 {
+		// chunk is how many tokens we can consume at a time
+		chunk := tokens
+		if chunk > minBurst {
+			chunk = minBurst
+		}
+
+		// maxDelay is the longest delay mandated by any of the limiters for
+		// the chosen chunk size.
+		var maxDelay time.Duration
+		for _, l := range ls {
+			res := l.ReserveN(time.Now(), chunk)
+			if del := res.Delay(); del > maxDelay {
+				maxDelay = del
+			}
+		}
+
+		time.Sleep(maxDelay)
+		tokens -= chunk
 	}
 }

cmd/syncthing/main.go

@@ -637,9 +637,6 @@ func syncthingMain(runtimeOptions RuntimeOptions) {
 		},
 	}
 
-	// If the read or write rate should be limited, set up a rate limiter for it.
-	// This will be used on connections created in the connect and listen routines.
-
 	opts := cfg.Options()
 
 	if !opts.SymlinksEnabled {