From a0ad185b98fbe83b686bedd533aa4e950b84d08b Mon Sep 17 00:00:00 2001
From: ehlxr <ehlxr.me@gmail.com>
Date: Tue, 26 Nov 2019 11:28:16 +0800
Subject: [PATCH] update at 2019-11-26 11:28:16 by ehlxr

---
 pkg/limiter.go      | 22 +++++++++++-----------
 pkg/limiter_test.go | 14 ++++++++++++++
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/pkg/limiter.go b/pkg/limiter.go
index e2912a7..a8e9d40 100644
--- a/pkg/limiter.go
+++ b/pkg/limiter.go
@@ -9,31 +9,31 @@ 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
+		limiter.time = now
 	}
 
 	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)
+	}
+}