Use juju rate limiter for connection rate limiting

This commit is contained in:
Geoff Bourne 2019-07-13 15:40:00 -05:00
parent 433dd2233b
commit 3699931af0
3 changed files with 8 additions and 2 deletions

3
go.mod
View File

@ -12,14 +12,17 @@ require (
github.com/hashicorp/golang-lru v0.5.1 // indirect
github.com/imdario/mergo v0.3.7 // indirect
github.com/json-iterator/go v1.1.6 // indirect
github.com/juju/ratelimit v1.0.1
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.1 // indirect
github.com/pkg/errors v0.8.1
github.com/sirupsen/logrus v1.4.1
github.com/spf13/pflag v1.0.3 // indirect
github.com/stretchr/testify v1.2.2
golang.org/x/crypto v0.0.0-20190411191339-88737f569e3a // indirect
golang.org/x/net v0.0.0-20190415214537-1da14a5a36f2 // indirect
golang.org/x/oauth2 v0.0.0-20190402181905-9f3314589c9a // indirect
golang.org/x/text v0.3.0
golang.org/x/time v0.0.0-20190308202827-9d24e82272b4 // indirect
gopkg.in/inf.v0 v0.9.1 // indirect
gopkg.in/yaml.v2 v2.2.2 // indirect

2
go.sum
View File

@ -20,6 +20,8 @@ github.com/imdario/mergo v0.3.7 h1:Y+UAYTZ7gDEuOfhxKWy+dvb5dRQ6rJjFSdX2HZY1/gI=
github.com/imdario/mergo v0.3.7/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJh5FfA=
github.com/json-iterator/go v1.1.6 h1:MrUvLMLTMxbqFJ9kzlvat/rYZqZnW3u4wkLzWTaFwKs=
github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCVDaaPEHmU=
github.com/juju/ratelimit v1.0.1 h1:+7AIFJVQ0EQgq/K9+0Krm7m530Du7tIz0METWzN0RgY=
github.com/juju/ratelimit v1.0.1/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk=
github.com/kisielk/errcheck v1.1.0/go.mod h1:EZBBE59ingxPouuu3KfxchcWSUPOHkagtvWXihfKN4Q=
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
github.com/konsorten/go-windows-terminal-sequences v1.0.1 h1:mweAR1A6xJ3oS2pRaGiHgQ4OO8tzTaLawm8vnODuwDk=

View File

@ -4,6 +4,7 @@ import (
"bytes"
"context"
"github.com/itzg/mc-router/mcproto"
"github.com/juju/ratelimit"
"github.com/sirupsen/logrus"
"io"
"net"
@ -43,14 +44,14 @@ func (c *connectorImpl) acceptConnections(ctx context.Context, ln net.Listener,
//noinspection GoUnhandledErrorResult
defer ln.Close()
limiter := time.Tick(time.Second / time.Duration(connRateLimit))
bucket := ratelimit.NewBucketWithRate(float64(connRateLimit), int64(connRateLimit*2))
for {
select {
case <-ctx.Done():
return
case <-limiter:
case <-time.After(bucket.Take(1)):
conn, err := ln.Accept()
if err != nil {
logrus.WithError(err).Error("Failed to accept connection")