(new): v1.2 release

2026-08-19 14:58:25 +03:00
parent f03a081ed6
commit 29b208eeec
79 changed files with 7301 additions and 2060 deletions
@@ -50,6 +50,9 @@ func NewRateLimiter() *RateLimiter {
 // from a background runner (e.g. once a minute) to bound memory in long-running
 // bots that serve many distinct chats.
 func (rl *RateLimiter) Cleanup(idleThreshold time.Duration) {
+	if idleThreshold <= 0 {
+		return
+	}
 	now := time.Now()
 	rl.chatMu.Lock()
 	defer rl.chatMu.Unlock()
@@ -87,7 +90,10 @@ func (rl *RateLimiter) SetGlobalLock(retryAfter int) {
 	}
 	rl.globalMu.Lock()
 	defer rl.globalMu.Unlock()
-	rl.globalLockUntil = time.Now().Add(time.Duration(retryAfter) * time.Second)
+	until := time.Now().Add(time.Duration(retryAfter) * time.Second)
+	if until.After(rl.globalLockUntil) {
+		rl.globalLockUntil = until
+	}
 }

 // SetChatLock sets a cooldown for a specific chat (e.g., after 429 for that chat).
@@ -98,7 +104,10 @@ func (rl *RateLimiter) SetChatLock(chatID int64, retryAfter int) {
 	}
 	rl.chatMu.Lock()
 	defer rl.chatMu.Unlock()
-	rl.chatLocks[chatID] = time.Now().Add(time.Duration(retryAfter) * time.Second)
+	until := time.Now().Add(time.Duration(retryAfter) * time.Second)
+	if until.After(rl.chatLocks[chatID]) {
+		rl.chatLocks[chatID] = until
+	}
 }

 // GlobalWait blocks until a global request can be made.
@@ -115,12 +124,17 @@ func (rl *RateLimiter) GlobalWait(ctx context.Context) error {
 }

 // Wait blocks until a request for the given chat can be made.
-// Waits for: chat cooldown → global cooldown → chat token bucket.
-// Note: Global limit is checked *before* chat limit to avoid overloading upstream.
+// Waits for: chat cooldown → chat token bucket → global cooldown → global token bucket.
+// Waiting for chat capacity first prevents one busy chat from reserving global
+// capacity while it is still unable to send a request.
 func (rl *RateLimiter) Wait(ctx context.Context, chatID int64) error {
 	if err := rl.waitForChatUnlock(ctx, chatID); err != nil {
 		return err
 	}
+	chatLimiter := rl.getChatLimiter(chatID)
+	if err := chatLimiter.Wait(ctx); err != nil {
+		return err
+	}
 	if err := rl.waitForGlobalUnlock(ctx); err != nil {
 		return err
 	}
@@ -130,8 +144,7 @@ func (rl *RateLimiter) Wait(ctx context.Context, chatID int64) error {
 			return err
 		}
 	}
-	chatLimiter := rl.getChatLimiter(chatID)
-	return chatLimiter.Wait(ctx)
+	return nil
 }

 func (rl *RateLimiter) getGlobalLimiter() *rate.Limiter {
@@ -233,60 +246,67 @@ func (rl *RateLimiter) Check(ctx context.Context, dropOverflow bool, chatID int6
 }

 func (rl *RateLimiter) waitForGlobalUnlock(ctx context.Context) error {
-	rl.globalMu.RLock()
-	until := rl.globalLockUntil
-	rl.globalMu.RUnlock()
+	for {
+		rl.globalMu.RLock()
+		until := rl.globalLockUntil
+		rl.globalMu.RUnlock()

-	if until.IsZero() || time.Now().After(until) {
-		return nil
-	}
+		if until.IsZero() || !time.Now().Before(until) {
+			return nil
+		}

-	select {
-	case <-time.After(time.Until(until)):
-		return nil
-	case <-ctx.Done():
-		return ctx.Err()
+		timer := time.NewTimer(time.Until(until))
+		select {
+		case <-timer.C:
+			// A concurrent 429 may have extended the cooldown. Re-read it.
+		case <-ctx.Done():
+			if !timer.Stop() {
+				select {
+				case <-timer.C:
+				default:
+				}
+			}
+			return ctx.Err()
+		}
 	}
 }

 func (rl *RateLimiter) waitForChatUnlock(ctx context.Context, chatID int64) error {
-	rl.chatMu.RLock()
-	until, ok := rl.chatLocks[chatID]
-	rl.chatMu.RUnlock()
+	for {
+		rl.chatMu.RLock()
+		until, ok := rl.chatLocks[chatID]
+		rl.chatMu.RUnlock()

-	if !ok || until.IsZero() || time.Now().After(until) {
-		return nil
-	}
+		if !ok || until.IsZero() || !time.Now().Before(until) {
+			return nil
+		}

-	select {
-	case <-time.After(time.Until(until)):
-		return nil
-	case <-ctx.Done():
-		return ctx.Err()
+		timer := time.NewTimer(time.Until(until))
+		select {
+		case <-timer.C:
+			// A concurrent 429 may have extended the cooldown. Re-read it.
+		case <-ctx.Done():
+			if !timer.Stop() {
+				select {
+				case <-timer.C:
+				default:
+				}
+			}
+			return ctx.Err()
+		}
 	}
 }

 // Updates chatLastSeen so Cleanup can evict idle entries.
 func (rl *RateLimiter) getChatLimiter(chatID int64) *rate.Limiter {
 	now := time.Now()
-
-	rl.chatMu.RLock()
-	lim, ok := rl.chatLimiters[chatID]
-	rl.chatMu.RUnlock()
-	if ok {
-		rl.chatMu.Lock()
-		rl.chatLastSeen[chatID] = now
-		rl.chatMu.Unlock()
-		return lim
-	}
-
 	rl.chatMu.Lock()
 	defer rl.chatMu.Unlock()
 	if lim, ok := rl.chatLimiters[chatID]; ok {
 		rl.chatLastSeen[chatID] = now
 		return lim
 	}
-	lim = rate.NewLimiter(1, 1)
+	lim := rate.NewLimiter(1, 1)
 	rl.chatLimiters[chatID] = lim
 	rl.chatLastSeen[chatID] = now
 	return lim
@@ -61,6 +61,86 @@ func TestRateLimiterGlobalWaitRespectsContextCancellation(t *testing.T) {
 	}
 }

+func TestRateLimiterWaitDoesNotConsumeGlobalCapacityWhileChatIsBlocked(t *testing.T) {
+	rl := NewRateLimiter()
+	rl.SetGlobalRate(1)
+
+	if !rl.Allow(42) {
+		t.Fatal("expected initial request for chat 42 to succeed")
+	}
+	rl.globalMu.Lock()
+	rl.globalLimiter = rate.NewLimiter(1, 1)
+	rl.globalMu.Unlock()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+	if err := rl.Wait(ctx, 42); err == nil {
+		t.Fatal("expected blocked chat wait to fail")
+	}
+	if !rl.GlobalAllow() {
+		t.Fatal("blocked chat wait consumed global capacity")
+	}
+}
+
+func TestRateLimiterWaitObservesExtendedCooldowns(t *testing.T) {
+	tests := []struct {
+		name string
+		wait func(*RateLimiter, context.Context) error
+		set  func(*RateLimiter, time.Time)
+	}{
+		{
+			name: "global",
+			wait: func(rl *RateLimiter, ctx context.Context) error {
+				return rl.waitForGlobalUnlock(ctx)
+			},
+			set: func(rl *RateLimiter, until time.Time) {
+				rl.globalMu.Lock()
+				rl.globalLockUntil = until
+				rl.globalMu.Unlock()
+			},
+		},
+		{
+			name: "chat",
+			wait: func(rl *RateLimiter, ctx context.Context) error {
+				return rl.waitForChatUnlock(ctx, 42)
+			},
+			set: func(rl *RateLimiter, until time.Time) {
+				rl.chatMu.Lock()
+				rl.chatLocks[42] = until
+				rl.chatMu.Unlock()
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rl := NewRateLimiter()
+			tt.set(rl, time.Now().Add(20*time.Millisecond))
+			started := time.Now()
+			done := make(chan error, 1)
+			go func() { done <- tt.wait(rl, context.Background()) }()
+
+			time.Sleep(5 * time.Millisecond)
+			tt.set(rl, time.Now().Add(70*time.Millisecond))
+			if err := <-done; err != nil {
+				t.Fatalf("wait returned error: %v", err)
+			}
+			if elapsed := time.Since(started); elapsed < 60*time.Millisecond {
+				t.Fatalf("wait ignored extended cooldown: %v", elapsed)
+			}
+		})
+	}
+}
+
+func TestRateLimiterCleanupIgnoresNonPositiveThreshold(t *testing.T) {
+	rl := NewRateLimiter()
+	want := rl.getChatLimiter(42)
+	rl.Cleanup(0)
+	if got := rl.getChatLimiter(42); got != want {
+		t.Fatal("Cleanup(0) replaced an active limiter")
+	}
+}
+
 // TestRateLimiterCleanupEvictsIdleChats guards the memory-leak fix: per-chat
 // limiter and lastSeen state must be reclaimed by Cleanup once the entry has
 // been idle for longer than the threshold, while still-active chats and
@@ -2,11 +2,11 @@ package utils

 const (
 	// VersionString is the module version string.
-	VersionString = "1.1.0"
+	VersionString = "1.2.0"
 	// VersionMajor is the module major version.
 	VersionMajor = 1
 	// VersionMinor is the module minor version.
-	VersionMinor = 1
+	VersionMinor = 2
 	// VersionPatch is the module patch version.
 	VersionPatch = 0
 	// VersionBeta is the prerelease counter for the current version.