frostfs-qos/limiting/limiter_test.go

package limiting_test

import (
	"math/rand"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"git.frostfs.info/TrueCloudLab/frostfs-qos/limiting"
	"github.com/stretchr/testify/require"
)

const operationDuration = time.Second
const operationCount = 64

type testKeyLimit struct {
	keys         []string
	limit        int
	withoutLimit bool
	failCount    atomic.Int32
}

func getLimits(testLimits []*testKeyLimit) []limiting.KeyLimit {
	var limits []limiting.KeyLimit
	for _, tl := range testLimits {
		if tl.withoutLimit {
			continue
		}
		limits = append(limits, limiting.KeyLimit{
			Keys:  tl.keys,
			Limit: uint64(tl.limit),
		})
	}
	return limits
}

func createTestTasks(testLimits []*testKeyLimit, limiter *limiting.Limiter) []func() {
	var tasks []func()
	for _, tl := range testLimits {
		for _, key := range tl.keys {
			tasks = append(tasks, func() {
				runParallelN(operationCount, func() {
					release, ok := limiter.TryAcquire(key)
					if !ok {
						tl.failCount.Add(1)
						return
					}
					defer release()
					time.Sleep(operationDuration)
				})
			})
		}
	}
	rand.Shuffle(len(tasks), func(i, j int) {
		tasks[i], tasks[j] = tasks[j], tasks[i]
	})
	return tasks
}

func runParallel(tasks ...func()) {
	var g sync.WaitGroup

	g.Add(len(tasks))
	for _, task := range tasks {
		go func() {
			defer g.Done()
			task()
		}()
	}
	g.Wait()
}

func runParallelN(N int, tsk func()) {
	tasks := make([]func(), N)
	for i := range N {
		tasks[i] = tsk
	}
	runParallel(tasks...)
}

func verifyResults(t *testing.T, testLimits []*testKeyLimit) {
	for _, tl := range testLimits {
		if tl.withoutLimit {
			require.Equal(t, 0, int(tl.failCount.Load()))
		} else {
			require.Equal(t, max(operationCount*len(tl.keys)-tl.limit, 0), int(tl.failCount.Load()))
		}
	}
}

func resetFailCounts(testLimits []*testKeyLimit) {
	for _, tl := range testLimits {
		tl.failCount.Store(0)
	}
}

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},
	}

	lr := limiting.New(getLimits(testLimits))
	tasks := createTestTasks(testLimits, lr)

	t.Run("run once", func(t *testing.T) {
		runParallel(tasks...)
		verifyResults(t, testLimits)
	})

	t.Run("run again", func(t *testing.T) {
		resetFailCounts(testLimits)
		runParallel(tasks...)
		verifyResults(t, testLimits)
	})
}