[#4] limiting: Add Limiter

Signed-off-by: Aleksey Savchuk <a.savchuk@yadro.com>

limiting/limiter_test.go

@@ -0,0 +1,145 @@
+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)
+	}
+}