Skip to content

Commit

Permalink
do not use global timeNow variables (#2477)
Browse files Browse the repository at this point in the history
  • Loading branch information
aldas authored Jul 11, 2023
1 parent 44ead54 commit 1ee8e22
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 20 deletions.
21 changes: 10 additions & 11 deletions middleware/rate_limiter.go
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,8 @@ type (
burst int
expiresIn time.Duration
lastCleanup time.Time

timeNow func() time.Time
}
// Visitor signifies a unique user's limiter details
Visitor struct {
Expand Down Expand Up @@ -219,7 +221,8 @@ func NewRateLimiterMemoryStoreWithConfig(config RateLimiterMemoryStoreConfig) (s
store.burst = int(config.Rate)
}
store.visitors = make(map[string]*Visitor)
store.lastCleanup = now()
store.timeNow = time.Now
store.lastCleanup = store.timeNow()
return
}

Expand All @@ -244,12 +247,13 @@ func (store *RateLimiterMemoryStore) Allow(identifier string) (bool, error) {
limiter.Limiter = rate.NewLimiter(store.rate, store.burst)
store.visitors[identifier] = limiter
}
limiter.lastSeen = now()
if now().Sub(store.lastCleanup) > store.expiresIn {
now := store.timeNow()
limiter.lastSeen = now
if now.Sub(store.lastCleanup) > store.expiresIn {
store.cleanupStaleVisitors()
}
store.mutex.Unlock()
return limiter.AllowN(now(), 1), nil
return limiter.AllowN(store.timeNow(), 1), nil
}

/*
Expand All @@ -258,14 +262,9 @@ of users who haven't visited again after the configured expiry time has elapsed
*/
func (store *RateLimiterMemoryStore) cleanupStaleVisitors() {
for id, visitor := range store.visitors {
if now().Sub(visitor.lastSeen) > store.expiresIn {
if store.timeNow().Sub(visitor.lastSeen) > store.expiresIn {
delete(store.visitors, id)
}
}
store.lastCleanup = now()
store.lastCleanup = store.timeNow()
}

/*
actual time method which is mocked in test file
*/
var now = time.Now
13 changes: 5 additions & 8 deletions middleware/rate_limiter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package middleware

import (
"errors"
"fmt"
"math/rand"
"net/http"
"net/http/httptest"
Expand Down Expand Up @@ -340,7 +339,7 @@ func TestRateLimiterMemoryStore_Allow(t *testing.T) {

for i, tc := range testCases {
t.Logf("Running testcase #%d => %v", i, time.Duration(i)*220*time.Millisecond)
now = func() time.Time {
inMemoryStore.timeNow = func() time.Time {
return time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC).Add(time.Duration(i) * 220 * time.Millisecond)
}
allowed, _ := inMemoryStore.Allow(tc.id)
Expand All @@ -350,24 +349,22 @@ func TestRateLimiterMemoryStore_Allow(t *testing.T) {

func TestRateLimiterMemoryStore_cleanupStaleVisitors(t *testing.T) {
var inMemoryStore = NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{Rate: 1, Burst: 3})
now = time.Now
fmt.Println(now())
inMemoryStore.visitors = map[string]*Visitor{
"A": {
Limiter: rate.NewLimiter(1, 3),
lastSeen: now(),
lastSeen: time.Now(),
},
"B": {
Limiter: rate.NewLimiter(1, 3),
lastSeen: now().Add(-1 * time.Minute),
lastSeen: time.Now().Add(-1 * time.Minute),
},
"C": {
Limiter: rate.NewLimiter(1, 3),
lastSeen: now().Add(-5 * time.Minute),
lastSeen: time.Now().Add(-5 * time.Minute),
},
"D": {
Limiter: rate.NewLimiter(1, 3),
lastSeen: now().Add(-10 * time.Minute),
lastSeen: time.Now().Add(-10 * time.Minute),
},
}

Expand Down
2 changes: 1 addition & 1 deletion middleware/request_logger.go
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,7 @@ func (config RequestLoggerConfig) ToMiddleware() (echo.MiddlewareFunc, error) {
if config.Skipper == nil {
config.Skipper = DefaultSkipper
}
now = time.Now
now := time.Now
if config.timeNow != nil {
now = config.timeNow
}
Expand Down

0 comments on commit 1ee8e22

Please sign in to comment.