diff --git a/pkg/limiter.go b/pkg/limiter.go index e2912a7..170cd37 100644 --- a/pkg/limiter.go +++ b/pkg/limiter.go @@ -9,32 +9,30 @@ type LimiterServer struct { interval time.Duration maxCount int sync.Mutex - reqCount int - startTime time.Time + reqCount int + time time.Time } -func NewLimiterServer(i time.Duration, c int) *LimiterServer { +func NewLimiterServer(interval time.Duration, maxCount int) *LimiterServer { return &LimiterServer{ - interval: i, - maxCount: c, + interval: interval, + maxCount: maxCount, } } func (limiter *LimiterServer) IsAvailable() bool { limiter.Lock() defer limiter.Unlock() + now := time.Now() - if limiter.startTime.IsZero() || - limiter.startTime.Add(limiter.interval).Before(time.Now()) { - limiter.reqCount = 1 - limiter.startTime = time.Now() - - return true + if limiter.time.IsZero() || + limiter.time.Add(limiter.interval).Before(now) { + limiter.reqCount = 0 } if limiter.reqCount < limiter.maxCount { limiter.reqCount += 1 - + limiter.time = now return true } diff --git a/pkg/limiter_test.go b/pkg/limiter_test.go index 1a8f368..fdd7c7b 100644 --- a/pkg/limiter_test.go +++ b/pkg/limiter_test.go @@ -1,6 +1,7 @@ package pkg import ( + "fmt" "testing" "time" ) @@ -16,3 +17,16 @@ func TestLimiter(t *testing.T) { } } } + +func TestLimiter2(t *testing.T) { + limiter := NewLimiterServer(10*time.Second, 10) + + for i := 0; i < 20; i++ { + if limiter.IsAvailable() { + fmt.Println("hello...", limiter.reqCount) + } else { + fmt.Println("limited") + } + time.Sleep(1 * time.Second) + } +}