Skip to content

Commit

Permalink
better naming and comment
Browse files Browse the repository at this point in the history
  • Loading branch information
henrod committed May 24, 2019
1 parent 64c5eeb commit a8a9f59
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 32 deletions.
26 changes: 16 additions & 10 deletions acceptorwrapper/rate_limiter.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,15 @@ import (
"github.com/topfreegames/pitaya/metrics"
)

// RateLimiter wraps net.Conn by applying rate limiting
// and return empty if exceeded
// RateLimiter wraps net.Conn by applying rate limiting and return empty
// if exceeded. It uses the leaky bucket
// algorithm (https://en.wikipedia.org/wiki/Leaky_bucket).
// Here, "limit" is the number of requests it accepts during an "interval" duration.
// After making a request, a slot is occupied and only freed after "interval"
// duration. If a new request comes when no slots are available, the buffer from
// Read is droped and ignored by pitaya.
// On the client side, this will yield a timeout error and the client must
// be prepared to handle it.
type RateLimiter struct {
net.Conn
limit int
Expand Down Expand Up @@ -72,9 +79,8 @@ func (r *RateLimiter) Read(b []byte) (n int, err error) {
}

now := time.Now()
err = r.take(now)
if err != nil {
logger.Log.Errorf("Data=%s, Error=%s", b, err)
if r.shouldRateLimit(now) {
logger.Log.Errorf("Data=%s, Error=%s", b, constants.ErrExceededRateLimiting)
metrics.ReportExceededRateLimiting(pitaya.GetMetricsReporters())
continue
}
Expand All @@ -83,20 +89,20 @@ func (r *RateLimiter) Read(b []byte) (n int, err error) {
}
}

// take saves the now as time taken or returns an error if
// shouldRateLimit saves the now as time taken or returns an error if
// in the limit of rate limiting
func (r *RateLimiter) take(now time.Time) error {
func (r *RateLimiter) shouldRateLimit(now time.Time) bool {
if r.times.Len() < r.limit {
r.times.PushBack(now)
return nil
return false
}

front := r.times.Front()
if diff := now.Sub(front.Value.(time.Time)); diff < r.interval {
return constants.ErrExceededRateLimiting
return true
}

front.Value = now
r.times.MoveToBack(front)
return nil
return false
}
43 changes: 21 additions & 22 deletions acceptorwrapper/rate_limiter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ import (

"github.com/golang/mock/gomock"
"github.com/stretchr/testify/assert"
"github.com/topfreegames/pitaya/constants"
"github.com/topfreegames/pitaya/mocks"
)

Expand All @@ -30,7 +29,7 @@ func TestRateLimiterRead(t *testing.T) {
expected int
err error
}{
"test_can_take_on_first_time": {
"test_can_read_on_first_time": {
forceDisable: false,
mock: func() {
mockConn.EXPECT().Read(buf).Return(10, nil)
Expand Down Expand Up @@ -98,7 +97,7 @@ func TestRateLimiterRead(t *testing.T) {
}
}

func TestRateLimiterTake(t *testing.T) {
func TestRateLimiterShouldRateLimit(t *testing.T) {
t.Parallel()

var (
Expand All @@ -110,34 +109,34 @@ func TestRateLimiterTake(t *testing.T) {

tables := map[string]struct {
before func()
err error
should bool
}{
"test_can_take_on_first_time": {
"test_should_not_on_first_time": {
before: func() {},
err: nil,
should: false,
},
"test_can_take_missing_one_to_limit": {
"test_should_not_missing_one_to_limit": {
before: func() {
r.take(now)
r.take(now)
r.shouldRateLimit(now)
r.shouldRateLimit(now)
},
err: nil,
should: false,
},
"test_can_take_when_oldest_request_expired": {
"test_should_not_when_oldest_request_expired": {
before: func() {
r.take(now.Add(-2 * interval))
r.take(now)
r.take(now)
r.shouldRateLimit(now.Add(-2 * interval))
r.shouldRateLimit(now)
r.shouldRateLimit(now)
},
err: nil,
should: false,
},
"test_cant_take_when_exceeded_limit": {
"test_should_when_exceeded_limit": {
before: func() {
r.take(now)
r.take(now)
r.take(now)
r.shouldRateLimit(now)
r.shouldRateLimit(now)
r.shouldRateLimit(now)
},
err: constants.ErrExceededRateLimiting,
should: true,
},
}

Expand All @@ -146,8 +145,8 @@ func TestRateLimiterTake(t *testing.T) {
r = NewRateLimiter(nil, limit, interval, false)

table.before()
err := r.take(now)
assert.Equal(t, table.err, err)
should := r.shouldRateLimit(now)
assert.Equal(t, table.should, should)
})
}
}

0 comments on commit a8a9f59

Please sign in to comment.