From 798503f041a30adbf5e877e6858ee1569bd00c95 Mon Sep 17 00:00:00 2001
From: Linsen Zhou <i@lin.moe>
Date: Thu, 17 Sep 2026 17:19:44 +0800
Subject: [PATCH 09/10] limiter/limiters: reimplement concurrency limiter using
 golang.org/x/sync/semaphore

---
 internal/limits/limiters/concurrency.go | 50 ++++++++++++++++++-------
 internal/limits/limiters/limiters.go    |  7 +++-
 internal/limits/limiters/rate.go        |  5 +--
 3 files changed, 44 insertions(+), 18 deletions(-)

diff --git a/internal/limits/limiters/concurrency.go b/internal/limits/limiters/concurrency.go
index 18923e8..2955b12 100644
--- a/internal/limits/limiters/concurrency.go
+++ b/internal/limits/limiters/concurrency.go
@@ -18,7 +18,11 @@ along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 package limiters
 
-import "context"
+import (
+	"context"
+
+	"golang.org/x/sync/semaphore"
+)
 
 // Semaphore is a convenience wrapper for a channel that implements
 // semaphore-kind synchronization.
@@ -26,43 +30,63 @@ import "context"
 // If the argument given to the NewSemaphore is negative or zero,
 // all methods are no-op.
 type Semaphore struct {
-	c chan struct{}
+	weighted *semaphore.Weighted
+	ctx      context.Context
+	cancel   context.CancelFunc
 }
 
 func NewSemaphore(max int) Semaphore {
-	return Semaphore{c: make(chan struct{}, max)}
+	ctx, cancel := context.WithCancel(context.TODO())
+	s := Semaphore{weighted: nil, ctx: ctx, cancel: cancel}
+	if max > 0 {
+		s.weighted = semaphore.NewWeighted(int64(max))
+	}
+	return s
 }
 
 func (s Semaphore) Take() bool {
-	if cap(s.c) <= 0 {
+	if s.weighted == nil {
 		return true
 	}
-	s.c <- struct{}{}
+
+	if err := s.weighted.Acquire(s.ctx, 1); err != nil {
+		return false
+	}
 	return true
 }
 
 func (s Semaphore) TakeContext(ctx context.Context) error {
-	if cap(s.c) <= 0 {
+	if s.weighted == nil {
 		return nil
 	}
 	select {
-	case s.c <- struct{}{}:
-		return nil
-	case <-ctx.Done():
-		return ctx.Err()
+	case <-s.ctx.Done():
+		return ErrClosed
+	default:
 	}
+	reqCtx, reqCancel := context.WithCancel(ctx)
+	defer reqCancel()
+
+	stop := context.AfterFunc(s.ctx, func() {
+		reqCancel()
+	})
+	defer stop()
+
+	return s.weighted.Acquire(reqCtx, 1)
 }
 
 func (s Semaphore) Release() {
-	if cap(s.c) <= 0 {
+	if s.weighted == nil {
 		return
 	}
 	select {
-	case <-s.c:
+	case <-s.ctx.Done():
+		return
 	default:
-		panic("limiters: mismatched Release call")
+		s.weighted.Release(1)
 	}
 }
 
 func (s Semaphore) Close() {
+	s.cancel()
 }
diff --git a/internal/limits/limiters/limiters.go b/internal/limits/limiters/limiters.go
index 9b11761..691580b 100644
--- a/internal/limits/limiters/limiters.go
+++ b/internal/limits/limiters/limiters.go
@@ -20,7 +20,10 @@ along with this program.  If not, see <https://www.gnu.org/licenses/>.
 // of resources consumed by the server.
 package limiters
 
-import "context"
+import (
+	"context"
+	"errors"
+)
 
 // The L interface represents a blocking limiter that has some upper bound of
 // resource use and blocks when it is exceeded until enough resources are
@@ -33,3 +36,5 @@ type L interface {
 	// Close frees any resources used internally by Limiter for book-keeping.
 	Close()
 }
+
+var ErrClosed = errors.New("limiters: Bucket is closed")
diff --git a/internal/limits/limiters/rate.go b/internal/limits/limiters/rate.go
index 3500995..ee0c44a 100644
--- a/internal/limits/limiters/rate.go
+++ b/internal/limits/limiters/rate.go
@@ -20,14 +20,11 @@ package limiters
 
 import (
 	"context"
-	"errors"
+	"time"
 
 	"golang.org/x/time/rate"
-	"time"
 )
 
-var ErrClosed = errors.New("limiters: Rate bucket is closed")
-
 // Rate structure implements a basic rate-limiter for requests using the token
 // bucket approach.
 //
-- 
2.54.0