87 lines
2 KiB
Go
87 lines
2 KiB
Go
package limiting
|
|
|
|
import (
|
|
"fmt"
|
|
)
|
|
|
|
type ReleaseFunc func()
|
|
|
|
type Limiter interface {
|
|
// TryAcquire attempts to reserve a slot without blocking.
|
|
//
|
|
// Returns a release function and true if successful, otherwise false.
|
|
// The release function must be called exactly once.
|
|
//
|
|
// If the key was not defined in the limiter, no limit is applied.
|
|
Acquire(key string) (ReleaseFunc, bool)
|
|
}
|
|
|
|
type semaphore interface {
|
|
acquire() bool
|
|
release()
|
|
}
|
|
|
|
type semaphoreLimiter[T semaphore] struct {
|
|
m map[string]T
|
|
}
|
|
|
|
// KeyLimit defines a concurrency limit for a set of keys.
|
|
//
|
|
// All keys of one set share the same limit.
|
|
// Keys of different sets have separate limits.
|
|
//
|
|
// Sets must not overlap.
|
|
type KeyLimit struct {
|
|
Keys []string
|
|
Limit int64
|
|
}
|
|
|
|
var NewAtomicLimiter = func(limits []KeyLimit) (Limiter, error) {
|
|
return newSemaphoreLimiter(limits, newAtomicSemaphore)
|
|
}
|
|
|
|
var NewBurstAtomicLimiter = func(limits []KeyLimit) (Limiter, error) {
|
|
return newSemaphoreLimiter(limits, newBurstAtomicSemaphore)
|
|
}
|
|
|
|
var NewChannelLimiter = func(limits []KeyLimit) (Limiter, error) {
|
|
return newSemaphoreLimiter(limits, newChannelSemaphore)
|
|
}
|
|
|
|
func newSemaphoreLimiter[T semaphore](limits []KeyLimit, newSemaphore func(size int64) T) (*semaphoreLimiter[T], error) {
|
|
lr := semaphoreLimiter[T]{make(map[string]T)}
|
|
for _, limit := range limits {
|
|
if err := lr.addLimit(&limit, newSemaphore); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
return &lr, nil
|
|
}
|
|
|
|
func (lr *semaphoreLimiter[T]) addLimit(limit *KeyLimit, newSemaphore func(size int64) T) error {
|
|
if limit.Limit < 0 {
|
|
return fmt.Errorf("invalid limit %d", limit.Limit)
|
|
}
|
|
|
|
sem := newSemaphore(limit.Limit)
|
|
for _, key := range limit.Keys {
|
|
if _, exists := lr.m[key]; exists {
|
|
return fmt.Errorf("duplicate key %q", key)
|
|
}
|
|
lr.m[key] = sem
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (lr *semaphoreLimiter[T]) Acquire(key string) (ReleaseFunc, bool) {
|
|
sem, ok := lr.m[key]
|
|
if !ok {
|
|
return func() {}, true
|
|
}
|
|
|
|
if ok := sem.acquire(); ok {
|
|
return func() { sem.release() }, true
|
|
}
|
|
return nil, false
|
|
}
|