easyai-ai-gateway/apps/api/internal/workerload/controller_test.go

package workerload

import (
	"context"
	"testing"
	"time"
)

func TestControllerStartsConservativelyAndGrowsUnderSustainedDemand(t *testing.T) {
	controller := New(Config{Mode: ModeAdaptive, HardLimit: 16, InitialActive: 4, InitialHeavy: 1, HealthySamples: 2})
	leases := make([]*Lease, 0, 4)
	for range 4 {
		lease, ok := controller.TryStart()
		if !ok {
			t.Fatal("initial task was not admitted")
		}
		_ = lease.EnterWaiting()
		leases = append(leases, lease)
	}
	for range 2 {
		controller.Observe(ResourceSample{MemoryCurrentBytes: 40, MemoryLimitBytes: 100, CPUUtilization: .4, DBConnections: 4, DBMaxConnections: 20})
	}
	snapshot := controller.Snapshot()
	if snapshot.ActiveLimit != 5 || snapshot.HeavyLimit != 2 {
		t.Fatalf("grown snapshot=%+v, want active=5 heavy=2", snapshot)
	}
	for _, lease := range leases {
		lease.Release()
	}
}

func TestControllerRemembersDemandAcrossShortPollingRuns(t *testing.T) {
	controller := New(Config{Mode: ModeAdaptive, HardLimit: 8, InitialActive: 4, InitialHeavy: 4, HealthySamples: 2})
	for range 2 {
		leases := make([]*Lease, 0, 4)
		for range 4 {
			lease, ok := controller.TryStart()
			if !ok {
				t.Fatal("short polling task was not admitted")
			}
			leases = append(leases, lease)
		}
		for _, lease := range leases {
			lease.Release()
		}
		controller.Observe(ResourceSample{
			MemoryCurrentBytes: 30,
			MemoryLimitBytes:   100,
			CPUUtilization:     .3,
			DBConnections:      3,
			DBMaxConnections:   20,
		})
	}
	if snapshot := controller.Snapshot(); snapshot.ActiveLimit != 5 || snapshot.SafeCapacity != 5 {
		t.Fatalf("grown snapshot=%+v, want active=5 safe=5", snapshot)
	}
}

func TestControllerDoesNotGrowWhileIdleWithoutDemand(t *testing.T) {
	controller := New(Config{Mode: ModeAdaptive, HardLimit: 8, InitialActive: 4, InitialHeavy: 1, HealthySamples: 2})
	for range 4 {
		controller.Observe(ResourceSample{MemoryCurrentBytes: 20, MemoryLimitBytes: 100, CPUUtilization: .2})
	}
	if snapshot := controller.Snapshot(); snapshot.ActiveLimit != 4 {
		t.Fatalf("idle snapshot=%+v, want active=4", snapshot)
	}
}

func TestControllerBusyAndCriticalPressureReduceNewClaims(t *testing.T) {
	controller := New(Config{Mode: ModeAdaptive, HardLimit: 16, InitialActive: 8, InitialHeavy: 2})
	busy := controller.Observe(ResourceSample{MemoryCurrentBytes: 80, MemoryLimitBytes: 100})
	if busy.PressureState != PressureBusy || busy.SafeCapacity >= 8 {
		t.Fatalf("busy snapshot=%+v", busy)
	}
	critical := controller.Observe(ResourceSample{MemoryCurrentBytes: 95, MemoryLimitBytes: 100})
	if critical.PressureState != PressureCritical || critical.SafeCapacity != 0 {
		t.Fatalf("critical snapshot=%+v", critical)
	}
	controller.SetClaimLimit(0)
	if _, ok := controller.TryStart(); ok {
		t.Fatal("critical controller admitted a new task")
	}
}

func TestWaitingReleasesHeavyPermitAndFinalizingReacquiresIt(t *testing.T) {
	controller := New(Config{Mode: ModeAdaptive, HardLimit: 4, InitialActive: 4, InitialHeavy: 1})
	first, ok := controller.TryStart()
	if !ok {
		t.Fatal("first lease unavailable")
	}
	if _, ok := controller.TryStart(); ok {
		t.Fatal("second preparing task bypassed heavy limit")
	}
	if err := first.EnterWaiting(); err != nil {
		t.Fatal(err)
	}
	second, ok := controller.TryStart()
	if !ok {
		t.Fatal("waiting task did not release heavy permit")
	}
	if err := second.EnterWaiting(); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := first.EnterFinalizing(ctx); err != nil {
		t.Fatal(err)
	}
	blocked := make(chan error, 1)
	go func() { blocked <- second.EnterFinalizing(ctx) }()
	select {
	case err := <-blocked:
		t.Fatalf("second finalizer did not wait: %v", err)
	case <-time.After(20 * time.Millisecond):
	}
	first.Release()
	if err := <-blocked; err != nil {
		t.Fatal(err)
	}
	second.Release()
	if snapshot := controller.Snapshot(); snapshot.ActiveTasks != 0 {
		t.Fatalf("active tasks=%d, want 0", snapshot.ActiveTasks)
	}
}

func TestPreparingLeaseCanTransitionDirectlyToFinalizingAtHeavyLimitOne(t *testing.T) {
	controller := New(Config{Mode: ModeAdaptive, HardLimit: 1, InitialActive: 1, InitialHeavy: 1})
	lease, ok := controller.TryStart()
	if !ok {
		t.Fatal("preparing lease unavailable")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	if err := lease.EnterFinalizing(ctx); err != nil {
		t.Fatalf("preparing lease deadlocked while entering finalizing: %v", err)
	}
	if phase := lease.Phase(); phase != PhaseFinalizing {
		t.Fatalf("lease phase=%s, want %s", phase, PhaseFinalizing)
	}
	lease.Release()
}

func TestLegacyModeUsesHardLimit(t *testing.T) {
	controller := New(Config{Mode: ModeLegacy, HardLimit: 7})
	snapshot := controller.Observe(ResourceSample{MemoryCurrentBytes: 99, MemoryLimitBytes: 100, CPUUtilization: 1, CPUThrottled: true})
	if snapshot.ActiveLimit != 7 || snapshot.HeavyLimit != 7 || snapshot.SafeCapacity != 7 || snapshot.PressureReason != "legacy" {
		t.Fatalf("legacy snapshot=%+v", snapshot)
	}
}