224 lines
6.8 KiB
Go
224 lines
6.8 KiB
Go
//go:build integration
|
|
|
|
package redisadmission
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
|
|
"proxy-pool/internal/platform/admission"
|
|
)
|
|
|
|
var integrationNamespaceSequence atomic.Uint64
|
|
|
|
func TestAdaptersShareGlobalAndPerKeyLimits(t *testing.T) {
|
|
fixture := newRedisFixture(t)
|
|
first := fixture.adapter(t, fixture.namespace, time.Minute, 3, 2)
|
|
second := fixture.adapter(t, fixture.namespace, time.Minute, 3, 2)
|
|
|
|
if err := first.Admit(t.Context(), "client-a"); err != nil {
|
|
t.Fatalf("first Admit(client-a): %v", err)
|
|
}
|
|
if err := second.Admit(t.Context(), "client-a"); err != nil {
|
|
t.Fatalf("second Admit(client-a): %v", err)
|
|
}
|
|
if err := first.Admit(t.Context(), "client-a"); !errors.Is(err, admission.ErrPerKeyLimit) {
|
|
t.Fatalf("shared per-key limit error = %v, want ErrPerKeyLimit", err)
|
|
}
|
|
if err := second.Admit(t.Context(), "client-b"); err != nil {
|
|
t.Fatalf("second Admit(client-b): %v", err)
|
|
}
|
|
if err := first.Admit(t.Context(), "client-c"); !errors.Is(err, admission.ErrGlobalLimit) {
|
|
t.Fatalf("shared global limit error = %v, want ErrGlobalLimit", err)
|
|
}
|
|
}
|
|
|
|
func TestConcurrentAdmissionsAreExactAcrossAdapters(t *testing.T) {
|
|
fixture := newRedisFixture(t)
|
|
tests := []struct {
|
|
name string
|
|
global int64
|
|
perKey int64
|
|
identity func(int) string
|
|
want int64
|
|
wantReject error
|
|
}{
|
|
{
|
|
name: "global", global: 37, perKey: 1000, want: 37, wantReject: admission.ErrGlobalLimit,
|
|
identity: func(index int) string { return fmt.Sprintf("client-%d", index) },
|
|
},
|
|
{
|
|
name: "per-key", global: 1000, perKey: 23, want: 23, wantReject: admission.ErrPerKeyLimit,
|
|
identity: func(int) string { return "shared-client" },
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
namespace := fixture.namespace + "-" + tt.name
|
|
first := fixture.adapter(t, namespace, time.Minute, tt.global, tt.perKey)
|
|
second := fixture.adapter(t, namespace, time.Minute, tt.global, tt.perKey)
|
|
var accepted atomic.Int64
|
|
var wrong atomic.Int64
|
|
var wait sync.WaitGroup
|
|
for index := range 200 {
|
|
wait.Add(1)
|
|
go func() {
|
|
defer wait.Done()
|
|
adapter := first
|
|
if index%2 == 1 {
|
|
adapter = second
|
|
}
|
|
err := adapter.Admit(t.Context(), tt.identity(index))
|
|
switch {
|
|
case err == nil:
|
|
accepted.Add(1)
|
|
case !errors.Is(err, tt.wantReject):
|
|
wrong.Add(1)
|
|
}
|
|
}()
|
|
}
|
|
wait.Wait()
|
|
if got := accepted.Load(); got != tt.want {
|
|
t.Fatalf("accepted = %d, want %d", got, tt.want)
|
|
}
|
|
if got := wrong.Load(); got != 0 {
|
|
t.Fatalf("unexpected rejection count = %d", got)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestWindowRolloverRemovesOldClientFields(t *testing.T) {
|
|
fixture := newRedisFixture(t)
|
|
const window = 100 * time.Millisecond
|
|
adapter := fixture.adapter(t, fixture.namespace, window, 100, 10)
|
|
if err := adapter.Admit(t.Context(), "old-client"); err != nil {
|
|
t.Fatalf("Admit(old-client): %v", err)
|
|
}
|
|
if err := fixture.client.Persist(t.Context(), adapter.key).Err(); err != nil {
|
|
t.Fatalf("PERSIST admission key: %v", err)
|
|
}
|
|
storedWindow, err := fixture.client.HGet(t.Context(), adapter.key, "window").Int64()
|
|
if err != nil {
|
|
t.Fatalf("read stored window: %v", err)
|
|
}
|
|
waitForRedisWindow(t, fixture.client, window, storedWindow)
|
|
|
|
if err := adapter.Admit(t.Context(), "new-client"); err != nil {
|
|
t.Fatalf("Admit(new-client): %v", err)
|
|
}
|
|
oldField := hashedIdentityField("old-client")
|
|
if exists, err := fixture.client.HExists(t.Context(), adapter.key, oldField).Result(); err != nil || exists {
|
|
t.Fatalf("old client field exists = %v, error = %v", exists, err)
|
|
}
|
|
newField := hashedIdentityField("new-client")
|
|
if exists, err := fixture.client.HExists(t.Context(), adapter.key, newField).Result(); err != nil || !exists {
|
|
t.Fatalf("new client field exists = %v, error = %v", exists, err)
|
|
}
|
|
if length, err := fixture.client.HLen(t.Context(), adapter.key).Result(); err != nil || length != 3 {
|
|
t.Fatalf("current window hash length = %d, error = %v, want 3", length, err)
|
|
}
|
|
}
|
|
|
|
func TestNamespacesAreIsolated(t *testing.T) {
|
|
fixture := newRedisFixture(t)
|
|
first := fixture.adapter(t, fixture.namespace+"-a", time.Minute, 1, 1)
|
|
second := fixture.adapter(t, fixture.namespace+"-b", time.Minute, 1, 1)
|
|
|
|
if err := first.Admit(t.Context(), "same-client"); err != nil {
|
|
t.Fatalf("first namespace Admit(): %v", err)
|
|
}
|
|
if err := second.Admit(t.Context(), "same-client"); err != nil {
|
|
t.Fatalf("second namespace Admit(): %v", err)
|
|
}
|
|
if err := first.Admit(t.Context(), "same-client"); !errors.Is(err, admission.ErrGlobalLimit) {
|
|
t.Fatalf("first namespace second Admit() error = %v", err)
|
|
}
|
|
}
|
|
|
|
type redisFixture struct {
|
|
client *redis.Client
|
|
namespace string
|
|
}
|
|
|
|
func newRedisFixture(t *testing.T) redisFixture {
|
|
t.Helper()
|
|
redisURL := os.Getenv("PROXY_POOL_TEST_REDIS_URL")
|
|
if redisURL == "" {
|
|
t.Skip("PROXY_POOL_TEST_REDIS_URL is not set")
|
|
}
|
|
options, err := redis.ParseURL(redisURL)
|
|
if err != nil {
|
|
t.Fatalf("parse PROXY_POOL_TEST_REDIS_URL: %v", err)
|
|
}
|
|
client := redis.NewClient(options)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
if err := client.Ping(ctx).Err(); err != nil {
|
|
_ = client.Close()
|
|
t.Fatalf("ping Redis: %v", err)
|
|
}
|
|
namespace := fmt.Sprintf("admission-it-%d-%d-%d", os.Getpid(), time.Now().UnixNano(), integrationNamespaceSequence.Add(1))
|
|
t.Cleanup(func() {
|
|
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cleanupCancel()
|
|
pattern := redisAdmissionKey + namespace + "*"
|
|
var cursor uint64
|
|
for {
|
|
keys, next, scanErr := client.Scan(cleanupCtx, cursor, pattern, 128).Result()
|
|
if scanErr != nil {
|
|
t.Errorf("scan Redis admission keys: %v", scanErr)
|
|
break
|
|
}
|
|
if len(keys) > 0 {
|
|
if unlinkErr := client.Unlink(cleanupCtx, keys...).Err(); unlinkErr != nil {
|
|
t.Errorf("remove Redis admission keys: %v", unlinkErr)
|
|
break
|
|
}
|
|
}
|
|
cursor = next
|
|
if cursor == 0 {
|
|
break
|
|
}
|
|
}
|
|
_ = client.Close()
|
|
})
|
|
return redisFixture{client: client, namespace: namespace}
|
|
}
|
|
|
|
func (fixture redisFixture) adapter(t *testing.T, namespace string, window time.Duration, global, perKey int64) *Adapter {
|
|
t.Helper()
|
|
return mustAdapter(t, fixture.client, Options{
|
|
Namespace: namespace,
|
|
Window: window,
|
|
Global: global,
|
|
PerKey: perKey,
|
|
})
|
|
}
|
|
|
|
func waitForRedisWindow(t *testing.T, client *redis.Client, window time.Duration, previous int64) {
|
|
t.Helper()
|
|
deadline := time.Now().Add(2 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
redisTime, err := client.Time(t.Context()).Result()
|
|
if err != nil {
|
|
t.Fatalf("Redis TIME: %v", err)
|
|
}
|
|
windowID := redisTime.UnixMilli() / window.Milliseconds()
|
|
if windowID != previous {
|
|
return
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
}
|
|
t.Fatal("timed out waiting for Redis window rollover")
|
|
}
|