All checks were successful
DCO action / DCO (pull_request) Successful in 29s
Vulncheck / Vulncheck (pull_request) Successful in 44s
Tests and linters / Run gofumpt (pull_request) Successful in 48s
Tests and linters / Staticcheck (pull_request) Successful in 50s
Pre-commit hooks / Pre-commit (pull_request) Successful in 1m23s
Tests and linters / Lint (pull_request) Successful in 1m24s
Tests and linters / gopls check (pull_request) Successful in 1m20s
Tests and linters / Tests (pull_request) Successful in 1m29s
Tests and linters / Tests with -race (pull_request) Successful in 1m46s
Vulncheck / Vulncheck (push) Successful in 35s
Tests and linters / Tests with -race (push) Successful in 56s
Tests and linters / Tests (push) Successful in 1m12s
Pre-commit hooks / Pre-commit (push) Successful in 1m15s
Tests and linters / Lint (push) Successful in 1m23s
Tests and linters / Run gofumpt (push) Successful in 1m17s
Tests and linters / Staticcheck (push) Successful in 1m27s
Tests and linters / gopls check (push) Successful in 1m42s
Previously, `Acquire` on exising key did 1 allocation because `func() { sem.Release() }` was a closure capturing different variables. Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>
75 lines
1.7 KiB
Go
75 lines
1.7 KiB
Go
package limiting
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"git.frostfs.info/TrueCloudLab/frostfs-qos/limiting/semaphore"
|
|
)
|
|
|
|
type ReleaseFunc func()
|
|
|
|
type Limiter interface {
|
|
// Acquire attempts to reserve a slot without blocking.
|
|
//
|
|
// Returns a release function and true if successful. The function must be
|
|
// called to release the limiter. The function must be called exactly once.
|
|
// Calling the function more that once will cause incorrect behavior of the
|
|
// limiter.
|
|
//
|
|
// Returns nil and false if fails.
|
|
//
|
|
// If the key was not defined in the limiter, no limit is applied.
|
|
Acquire(key string) (ReleaseFunc, bool)
|
|
}
|
|
|
|
type SemaphoreLimiter struct {
|
|
m map[string]*semaphore.Semaphore
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
func NewSemaphoreLimiter(limits []KeyLimit) (*SemaphoreLimiter, error) {
|
|
lr := SemaphoreLimiter{make(map[string]*semaphore.Semaphore)}
|
|
for _, limit := range limits {
|
|
if limit.Limit < 0 {
|
|
return nil, fmt.Errorf("invalid limit %d", limit.Limit)
|
|
}
|
|
sem := semaphore.NewSemaphore(limit.Limit)
|
|
|
|
if err := lr.addLimit(&limit, sem); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
return &lr, nil
|
|
}
|
|
|
|
func (lr *SemaphoreLimiter) addLimit(limit *KeyLimit, sem *semaphore.Semaphore) error {
|
|
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) Acquire(key string) (ReleaseFunc, bool) {
|
|
sem, ok := lr.m[key]
|
|
if !ok {
|
|
return func() {}, true
|
|
}
|
|
|
|
if ok := sem.Acquire(); ok {
|
|
return sem.Release, true
|
|
}
|
|
return nil, false
|
|
}
|