feat: improve bucket logic
Kim committed
Nov 13, 2025 at 18:43 UTC
fa2734a85e81e46f0de2b18f5a90ae263f4d57f5
2 files changed
+70
-34
portal/utils/ratelimit/bucket.go
+37
-34
@@ -6,51 +6,54 @@ import (
6
"time"
7
)
8
9
-// Bucket is a simple, precise byte-rate limiter (thread-safe).
10
-// All fields are integer-based (bytes, ns) to avoid float overhead.
9
+// Bucket is a very simple thread-safe rate limiter.
10
+// It uses a shared timeline (allowAt) with a fixed per-byte duration and
11
+// a maximum slack to model burst capacity.
12
type Bucket struct {
13
mu sync.Mutex
13
- rateBps int64 // bytes per second
14
- capacity int64 // max tokens (burst), typically = rateBps
15
- tokens int64 // current tokens in bytes
16
- last time.Time // last refill time
14
+ perByte time.Duration // time per byte
15
+ maxSlack time.Duration // maximum credit time (burst)
16
+ allowAt time.Time // next allowed time on the timeline
17
}
18
19
-// NewBucket creates a bucket with the given rate and burst (bytes).
20
-// If burst <= 0, it defaults to rateBps.
19
+// NewBucket creates a limiter for rateBps with burst bytes.
20
+// burst is translated to time slack = burst * perByte.
21
func NewBucket(rateBps int64, burst int64) *Bucket {
22
+ if rateBps <= 0 {
23
+ return nil
24
+ }
25
if burst <= 0 {
26
burst = rateBps
27
}
25
- return &Bucket{rateBps: rateBps, capacity: burst, tokens: burst, last: time.Now()}
28
+ perByte := time.Second / time.Duration(rateBps)
29
+ if perByte <= 0 {
30
+ perByte = time.Nanosecond
31
+ }
32
+ maxSlack := perByte * time.Duration(burst)
33
+ now := time.Now()
34
+ // Start with full burst credit available
35
+ allowAt := now.Add(-maxSlack)
36
+ return &Bucket{perByte: perByte, maxSlack: maxSlack, allowAt: allowAt}
37
}
38
28
-// Take blocks until n bytes worth of tokens are available, then consumes them.
39
+// Take blocks long enough to account for n bytes at the configured rate.
40
+// It is safe for concurrent use and coordinates consumers by a shared timeline.
41
func (b *Bucket) Take(n int64) {
30
- for {
31
- var sleep time.Duration
32
- b.mu.Lock()
33
- now := time.Now()
34
- elapsed := now.Sub(b.last)
35
- if elapsed > 0 {
36
- refill := (elapsed.Nanoseconds() * b.rateBps) / int64(time.Second)
37
- if refill > 0 {
38
- b.tokens += refill
39
- if b.tokens > b.capacity {
40
- b.tokens = b.capacity
41
- }
42
- b.last = now
43
- }
44
- }
45
- if b.tokens >= n {
46
- b.tokens -= n
47
- b.mu.Unlock()
48
- return
49
- }
50
- deficit := n - b.tokens
51
- nsNeeded := max((deficit*int64(time.Second))/b.rateBps, int64(time.Millisecond))
52
- sleep = time.Duration(nsNeeded)
53
- b.mu.Unlock()
42
+ if b == nil || n <= 0 {
43
+ return
44
+ }
45
+ b.mu.Lock()
46
+ now := time.Now()
47
+ // Refill slack over time up to maxSlack
48
+ if now.Sub(b.allowAt) > b.maxSlack {
49
+ b.allowAt = now.Add(-b.maxSlack)
50
+ }
51
+ start := b.allowAt
52
+ finish := start.Add(b.perByte * time.Duration(n))
53
+ b.allowAt = finish
54
+ b.mu.Unlock()
55
+
56
+ if sleep := finish.Sub(now); sleep > 0 {
57
time.Sleep(sleep)
58
}
59
}
portal/utils/ratelimit/bucket_test.go
new
+33
@@ -0,0 +1,33 @@
1
+package ratelimit
2
+
3
+import (
4
+ "testing"
5
+ "time"
6
+)
7
+
8
+// This test keeps expectations deliberately loose to avoid flakiness
9
+// while still catching gross misbehavior.
10
+func TestSimpleRateAndBurst(t *testing.T) {
11
+ rate := int64(1 * 1024 * 1024) // 1 MiB/s
12
+ burst := rate // allow ~1s worth of burst
13
+ b := NewBucket(rate, burst)
14
+ if b == nil {
15
+ t.Fatalf("bucket should not be nil for positive rate")
16
+ }
17
+
18
+ // First half-burst should complete quickly (use a generous threshold)
19
+ start := time.Now()
20
+ b.Take(burst / 2)
21
+ fast := time.Since(start)
22
+ if fast > 200*time.Millisecond {
23
+ t.Fatalf("half-burst took too long: %v", fast)
24
+ }
25
+
26
+ // Taking 2*rate bytes should take roughly ~1s given 1s burst credit.
27
+ start = time.Now()
28
+ b.Take(2 * rate)
29
+ elapsed := time.Since(start)
30
+ if elapsed < 700*time.Millisecond { // be tolerant to scheduling variance
31
+ t.Fatalf("expected at least ~0.7s throttling, got %v", elapsed)
32
+ }
33
+}