From 719c3a560152bec8309f6a0e429cc4844c5e96c2 Mon Sep 17 00:00:00 2001
From: Linsen Zhou <i@lin.moe>
Date: Thu, 17 Sep 2026 17:18:01 +0800
Subject: [PATCH 08/10] limits/limiters: reimplement rate limiter using
 golang.org/x/time/rate

---
 internal/limits/limiters/rate.go | 81 +++++++++++---------------------
 1 file changed, 28 insertions(+), 53 deletions(-)

diff --git a/internal/limits/limiters/rate.go b/internal/limits/limiters/rate.go
index a774187..3500995 100644
--- a/internal/limits/limiters/rate.go
+++ b/internal/limits/limiters/rate.go
@@ -21,6 +21,8 @@ package limiters
 import (
 	"context"
 	"errors"
+
+	"golang.org/x/time/rate"
 	"time"
 )
 
@@ -37,81 +39,54 @@ var ErrClosed = errors.New("limiters: Rate bucket is closed")
 //
 // If burstSize = 0, all methods are no-op and always succeed.
 type Rate struct {
-	bucket chan struct{}
-	stop   chan struct{}
+	limiter *rate.Limiter
+	ctx     context.Context
+	cancel  context.CancelFunc
 }
 
 func NewRate(burstSize int, interval time.Duration) Rate {
-	r := Rate{
-		bucket: make(chan struct{}, burstSize),
-		stop:   make(chan struct{}),
-	}
-
-	if burstSize == 0 {
-		return r
+	ctx, cancel := context.WithCancel(context.TODO())
+	r := Rate{limiter: nil, ctx: ctx, cancel: cancel}
+	if burstSize > 0 {
+		r.limiter = rate.NewLimiter(rate.Every(interval), burstSize)
 	}
-
-	for i := 0; i < burstSize; i++ {
-		r.bucket <- struct{}{}
-	}
-
-	go r.fill(burstSize, interval)
 	return r
 }
 
-func (r Rate) fill(burstSize int, interval time.Duration) {
-	t := time.NewTimer(interval)
-	defer t.Stop()
-	for {
-		t.Reset(interval)
-		select {
-		case <-t.C:
-		case <-r.stop:
-			close(r.bucket)
-			return
-		}
-
-	fill:
-		for i := 0; i < burstSize; i++ {
-			select {
-			case r.bucket <- struct{}{}:
-			default:
-				// If there are no Take pending and the bucket is already
-				// full - don't block.
-				break fill
-			}
-		}
-	}
-}
-
 func (r Rate) Take() bool {
-	if cap(r.bucket) == 0 {
+	if r.limiter == nil {
 		return true
 	}
 
-	_, ok := <-r.bucket
-	return ok
+	if err := r.limiter.Wait(r.ctx); err != nil {
+		return false
+	}
+	return true
 }
 
 func (r Rate) TakeContext(ctx context.Context) error {
-	if cap(r.bucket) == 0 {
+	if r.limiter == nil {
 		return nil
 	}
-
 	select {
-	case _, ok := <-r.bucket:
-		if !ok {
-			return ErrClosed
-		}
-		return nil
-	case <-ctx.Done():
-		return ctx.Err()
+	case <-r.ctx.Done():
+		return ErrClosed
+	default:
 	}
+	reqCtx, reqCancel := context.WithCancel(ctx)
+	defer reqCancel()
+
+	stop := context.AfterFunc(r.ctx, func() {
+		reqCancel()
+	})
+	defer stop()
+
+	return r.limiter.Wait(reqCtx)
 }
 
 func (r Rate) Release() {
 }
 
 func (r Rate) Close() {
-	close(r.stop)
+	r.cancel()
 }
-- 
2.54.0