update at 2019-11-26 13:50:08 by ehlxr

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
 	}
 

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)
+	}
+}