145 lines
3.1 KiB
Go
145 lines
3.1 KiB
Go
package limiting_test
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.frostfs.info/TrueCloudLab/frostfs-qos/limiting"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
const (
|
|
operationDuration = 100 * time.Millisecond
|
|
operationCount = 64
|
|
)
|
|
|
|
type testKeyLimit struct {
|
|
keys []string
|
|
limit int
|
|
withoutLimit bool
|
|
failCount atomic.Int32
|
|
}
|
|
|
|
func TestLimiter(t *testing.T) {
|
|
testLimits := []*testKeyLimit{
|
|
{keys: []string{"A"}, limit: operationCount / 4},
|
|
{keys: []string{"B"}, limit: operationCount / 2},
|
|
{keys: []string{"C", "D"}, limit: operationCount / 4},
|
|
{keys: []string{"E"}, limit: 2 * operationCount},
|
|
{keys: []string{"F"}, withoutLimit: true},
|
|
}
|
|
|
|
t.Run("non-blocking mode", func(t *testing.T) {
|
|
testLimiter(t, testLimits, false)
|
|
})
|
|
|
|
resetFailCounts(testLimits)
|
|
|
|
t.Run("blocking mode", func(t *testing.T) {
|
|
testLimiter(t, testLimits, true)
|
|
})
|
|
}
|
|
|
|
func testLimiter(t *testing.T, testCases []*testKeyLimit, blocking bool) {
|
|
lr := limiting.New(getLimits(testCases))
|
|
tasks := createTestTasks(testCases, lr, blocking)
|
|
|
|
t.Run("first run", func(t *testing.T) {
|
|
executeTasks(tasks...)
|
|
verifyResults(t, testCases, blocking)
|
|
})
|
|
|
|
t.Run("repeated run", func(t *testing.T) {
|
|
resetFailCounts(testCases)
|
|
executeTasks(tasks...)
|
|
verifyResults(t, testCases, blocking)
|
|
})
|
|
}
|
|
|
|
func getLimits(testCases []*testKeyLimit) []limiting.KeyLimit {
|
|
var limits []limiting.KeyLimit
|
|
for _, tc := range testCases {
|
|
if tc.withoutLimit {
|
|
continue
|
|
}
|
|
limits = append(limits, limiting.KeyLimit{
|
|
Keys: tc.keys,
|
|
Limit: uint64(tc.limit),
|
|
})
|
|
}
|
|
return limits
|
|
}
|
|
|
|
func createTestTasks(testCases []*testKeyLimit, lr *limiting.Limiter, blocking bool) []func() {
|
|
var tasks []func()
|
|
for _, tc := range testCases {
|
|
for _, key := range tc.keys {
|
|
tasks = append(tasks, func() {
|
|
executeTaskN(operationCount, func() { acquireAndExecute(tc, lr, key, blocking) })
|
|
})
|
|
}
|
|
}
|
|
return tasks
|
|
}
|
|
|
|
func acquireAndExecute(tc *testKeyLimit, lr *limiting.Limiter, key string, blocking bool) {
|
|
if blocking {
|
|
release, err := lr.Acquire(context.Background(), key)
|
|
if err != nil {
|
|
tc.failCount.Add(1)
|
|
return
|
|
}
|
|
defer release()
|
|
} else {
|
|
release, ok := lr.TryAcquire(key)
|
|
if !ok {
|
|
tc.failCount.Add(1)
|
|
return
|
|
}
|
|
defer release()
|
|
}
|
|
time.Sleep(operationDuration)
|
|
}
|
|
|
|
func executeTasks(tasks ...func()) {
|
|
var g sync.WaitGroup
|
|
|
|
g.Add(len(tasks))
|
|
for _, task := range tasks {
|
|
go func() {
|
|
defer g.Done()
|
|
task()
|
|
}()
|
|
}
|
|
g.Wait()
|
|
}
|
|
|
|
func executeTaskN(N int, task func()) {
|
|
tasks := make([]func(), N)
|
|
for i := range N {
|
|
tasks[i] = task
|
|
}
|
|
executeTasks(tasks...)
|
|
}
|
|
|
|
func verifyResults(t *testing.T, testCases []*testKeyLimit, blocking bool) {
|
|
for _, tc := range testCases {
|
|
var expectedFailCount int
|
|
if blocking || tc.withoutLimit {
|
|
expectedFailCount = 0
|
|
} else {
|
|
expectedFailCount = max(operationCount*len(tc.keys)-tc.limit, 0)
|
|
}
|
|
actualFailCount := int(tc.failCount.Load())
|
|
require.Equal(t, expectedFailCount, actualFailCount)
|
|
}
|
|
}
|
|
|
|
func resetFailCounts(testLimits []*testKeyLimit) {
|
|
for _, tc := range testLimits {
|
|
tc.failCount.Store(0)
|
|
}
|
|
}
|