[#4] limiting: Add check for duplicated keys
Some checks failed
DCO action / DCO (pull_request) Successful in 26s
Vulncheck / Vulncheck (pull_request) Failing after 45s
Tests and linters / Run gofumpt (pull_request) Successful in 48s
Tests and linters / Staticcheck (pull_request) Successful in 1m8s
Pre-commit hooks / Pre-commit (pull_request) Successful in 1m27s
Tests and linters / Lint (pull_request) Successful in 1m30s
Tests and linters / Tests with -race (pull_request) Successful in 1m41s
Tests and linters / gopls check (pull_request) Successful in 1m52s
Tests and linters / Tests (pull_request) Successful in 2m0s

Signed-off-by: Aleksey Savchuk <a.savchuk@yadro.com>
This commit is contained in:
Aleksey Savchuk 2025-02-10 13:00:29 +03:00
parent 240501c1b7
commit ae6938c61c
Signed by: a-savchuk
GPG key ID: 70C0A7FF6F9C4639
2 changed files with 26 additions and 6 deletions

View file

@ -2,6 +2,7 @@ package limiting
import (
"context"
"fmt"
)
type semaphore struct {
@ -51,15 +52,24 @@ type KeyLimit struct {
type ReleaseFunc func()
func New(limits []KeyLimit) *Limiter {
func New(limits []KeyLimit) (*Limiter, error) {
lr := Limiter{m: make(map[string]*semaphore)}
for _, l := range limits {
sem := newSemaphore(l.Limit)
for _, key := range l.Keys {
lr.m[key] = sem
if err := addLimit(&lr, l.Keys, newSemaphore(l.Limit)); err != nil {
return nil, err
}
}
return &lr
return &lr, nil
}
func addLimit(lr *Limiter, keys []string, sem *semaphore) error {
for _, key := range keys {
if _, exists := lr.m[key]; exists {
return fmt.Errorf("duplicate key %q", key)
}
lr.m[key] = sem
}
return nil
}
// Acquire reserves a slot for the given key, blocking if necessary.

View file

@ -24,6 +24,14 @@ type testKeyLimit struct {
}
func TestLimiter(t *testing.T) {
t.Run("duplicate key", func(t *testing.T) {
_, err := limiting.New([]limiting.KeyLimit{
{[]string{"A", "B"}, 10},
{[]string{"B", "C"}, 10},
})
require.Error(t, err)
})
testLimits := []*testKeyLimit{
{keys: []string{"A"}, limit: operationCount / 4},
{keys: []string{"B"}, limit: operationCount / 2},
@ -44,7 +52,9 @@ func TestLimiter(t *testing.T) {
}
func testLimiter(t *testing.T, testCases []*testKeyLimit, blocking bool) {
lr := limiting.New(getLimits(testCases))
lr, err := limiting.New(getLimits(testCases))
require.NoError(t, err)
tasks := createTestTasks(testCases, lr, blocking)
t.Run("first run", func(t *testing.T) {