From c4dfe75a2dcc71c0049aa28a0f662135b1a64750 Mon Sep 17 00:00:00 2001
From: Liqun <wanliqunn@163.com>
Date: Wed, 22 May 2024 14:50:55 +0800
Subject: [PATCH] Improve rate limit error msg and deprecate codes (#41)

---
 rate/fixed_window.go | 21 ++++-----------------
 rate/token_bucket.go |  2 +-
 2 files changed, 5 insertions(+), 18 deletions(-)

diff --git a/rate/fixed_window.go b/rate/fixed_window.go
index 9d7d62b..db10b53 100644
--- a/rate/fixed_window.go
+++ b/rate/fixed_window.go
@@ -18,7 +18,7 @@ func NewFixedWindow(interval time.Duration, max int) Limiter {
 	return &FixedWindow{
 		interval:  interval,
 		max:       int64(max),
-		startTime: toUnixMilli(time.Now().Truncate(interval)),
+		startTime: time.Now().Truncate(interval).UnixMilli(),
 	}
 }
 
@@ -45,14 +45,14 @@ func (window *FixedWindow) LimitAt(now time.Time, n int) error {
 	atomic.AddInt64(&window.count, int64(-n))
 
 	startTime := atomic.LoadInt64(&window.startTime)
-	nextTime := fromUnixMilli(startTime).Add(window.interval)
+	nextTime := time.UnixMilli(startTime).Add(window.interval)
 	waitTime := nextTime.Sub(now)
 
 	return errRateLimited(int(window.max), waitTime)
 }
 
 func (window *FixedWindow) advance(now time.Time) {
-	truncated := toUnixMilli(now.Truncate(window.interval))
+	truncated := now.Truncate(window.interval).UnixMilli()
 
 	if startTime := atomic.LoadInt64(&window.startTime); startTime < truncated {
 		// reset
@@ -63,18 +63,5 @@ func (window *FixedWindow) advance(now time.Time) {
 
 func (window *FixedWindow) Expired() bool {
 	startTime := atomic.LoadInt64(&window.startTime)
-	return time.Since(fromUnixMilli(startTime)) > window.interval
-}
-
-// TODO: deprecate the following time utility functions if the minimum version
-// of Go required by the package is gte 1.7.
-
-// toUnixMilli returns unix timestamp in milliseconds for specific time.
-func toUnixMilli(t time.Time) int64 {
-	return t.UnixNano() / (1e6)
-}
-
-// fromUnixMilli returns time of specific unix timestamp in milliseconds.
-func fromUnixMilli(msec int64) time.Time {
-	return time.Unix(msec/1e3, (msec%1e3)*1e6)
+	return time.Since(time.UnixMilli(startTime)) > window.interval
 }
diff --git a/rate/token_bucket.go b/rate/token_bucket.go
index 577f060..b6905dc 100644
--- a/rate/token_bucket.go
+++ b/rate/token_bucket.go
@@ -43,7 +43,7 @@ func (bucket *TokenBucket) LimitAt(now time.Time, n int) error {
 
 	if waitTime := rsv.Delay(); waitTime > 0 {
 		rsv.Cancel()
-		return errRateLimited(bucket.inner.Burst(), waitTime)
+		return errRateLimited(int(bucket.inner.Limit()), waitTime)
 	}
 
 	atomic.StoreInt64(&bucket.lastSeen, now.Unix())