AIRouter/test/performance/benchmark_test.go

package performance

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/edgeai/gateway/internal/config"
	"github.com/edgeai/gateway/internal/observability"
	"github.com/edgeai/gateway/internal/server"
)

var perfServer *httptest.Server

func TestMain(m *testing.M) {
	cfg := &config.Config{
		Server: config.ServerConfig{Host: "127.0.0.1", Port: 0, MaxRequestBodyMB: 20},
		Auth:   config.AuthConfig{Enabled: true, Methods: []string{"api_key"}},
		Scheduler: config.SchedulerConfig{
			MaxRunningTasks: 16, MaxQueuedTasks: 1000,
			Fairness: "weighted_fair_queue", PriorityAgingSeconds: 30,
		},
		Timeouts: config.TimeoutConfig{
			DefaultConnectMs: 5000, DefaultQueueMs: 5000, DefaultFirstTokenMs: 10000,
			DefaultInferenceMs: 60000, DefaultIdleMs: 15000, DefaultTotalMs: 90000,
		},
		Context: config.ContextConfig{SafetyMarginRatio: 0.08, DefaultPolicy: "recent_only"},
		Models: map[string]config.ModelConfig{
			"test-chat": {
				Provider: "ollama", ActualModel: "qwen2.5:0.5b",
				Endpoint: "http://127.0.0.1:11434", ContextWindow: 4096,
				MaxOutputTokens: 256, MaxConcurrency: 4, CancelSupported: true,
			},
		},
		Observability: config.ObservabilityConfig{MetricsPath: "/metrics", LogLevel: "warn"},
		Storage: config.StorageConfig{
			SessionDB: "sqlite:///tmp/edgeai-perf/sessions.db",
			TaskState: "sqlite:///tmp/edgeai-perf/tasks.db",
		},
	}

	os.MkdirAll("/tmp/edgeai-perf", 0755)
	defer os.RemoveAll("/tmp/edgeai-perf")

	logger := observability.NewLogger(observability.LevelWarn, os.Stdout, "metadata_only")
	srv, err := server.New(cfg, logger)
	if err != nil {
		panic("failed to create perf server: " + err.Error())
	}

	perfServer = httptest.NewServer(srv.HTTPSrv.Handler)
	defer perfServer.Close()

	m.Run()
}

// PERF-001: Health check latency under 5ms
func TestPerf001_HealthLatency(t *testing.T) {
	var total time.Duration
	iterations := 100

	for i := 0; i < iterations; i++ {
		start := time.Now()
		resp, err := http.Get(perfServer.URL + "/health")
		if err != nil {
			t.Fatal(err)
		}
		resp.Body.Close()
		total += time.Since(start)
	}

	avgMs := total.Milliseconds() / int64(iterations)
	if avgMs > 5 {
		t.Errorf("average health check latency %dms exceeds 5ms target", avgMs)
	}
	t.Logf("average health check latency: %dms", avgMs)
}

// PERF-002: Concurrent health checks
func TestPerf002_ConcurrentHealth(t *testing.T) {
	concurrency := 50
	var wg sync.WaitGroup
	wg.Add(concurrency)

	start := time.Now()
	for i := 0; i < concurrency; i++ {
		go func() {
			defer wg.Done()
			resp, err := http.Get(perfServer.URL + "/health")
			if err != nil {
				t.Error(err)
				return
			}
			resp.Body.Close()
		}()
	}
	wg.Wait()
	elapsed := time.Since(start)
	t.Logf("%d concurrent health checks completed in %v", concurrency, elapsed)
}

// PERF-003: Metrics endpoint latency under 10ms
func TestPerf003_MetricsLatency(t *testing.T) {
	var total time.Duration
	iterations := 50

	for i := 0; i < iterations; i++ {
		start := time.Now()
		resp, err := http.Get(perfServer.URL + "/metrics")
		if err != nil {
			t.Fatal(err)
		}
		resp.Body.Close()
		total += time.Since(start)
	}

	avgMs := total.Milliseconds() / int64(iterations)
	if avgMs > 10 {
		t.Errorf("average metrics latency %dms exceeds 10ms target", avgMs)
	}
	t.Logf("average metrics latency: %dms", avgMs)
}

// PERF-004: Auth check latency under 2ms
func TestPerf004_AuthLatency(t *testing.T) {
	var total time.Duration
	iterations := 100

	for i := 0; i < iterations; i++ {
		start := time.Now()
		req, _ := http.NewRequest("GET", perfServer.URL+"/v1/models", nil)
		req.Header.Set("Authorization", "Bearer test-key")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		resp.Body.Close()
		total += time.Since(start)
	}

	avgUs := total.Microseconds() / int64(iterations)
	t.Logf("average auth check latency: %dus", avgUs)
}

// PERF-005: Scheduler throughput
func TestPerf005_SchedulerThroughput(t *testing.T) {
	// Submit and complete many tasks rapidly
	ctx := context.Background()
	_ = ctx
	iterations := 1000
	start := time.Now()

	for i := 0; i < iterations; i++ {
		req, _ := http.NewRequest("GET", perfServer.URL+"/health", nil)
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		resp.Body.Close()
	}

	elapsed := time.Since(start)
	rps := float64(iterations) / elapsed.Seconds()
	t.Logf("Throughput: %.0f requests/sec (%d requests in %v)", rps, iterations, elapsed)
}

// PERF-006: Memory usage stable under load
func TestPerf006_MemoryStability(t *testing.T) {
	// Run requests for 2 seconds and check no panic
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	count := 0
	for {
		select {
		case <-ctx.Done():
			t.Logf("completed %d requests in 2s without crash", count)
			return
		default:
			resp, err := http.Get(perfServer.URL + "/health")
			if err != nil {
				// Port exhaustion is acceptable under extreme load
				continue
			}
			resp.Body.Close()
			count++
			if count%50 == 0 {
				time.Sleep(10 * time.Millisecond)
			}
		}
	}
}

// PERF-007: Cancellation timing
func TestPerf007_CancellationTiming(t *testing.T) {
	// Cancel a request and verify it returns quickly
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	start := time.Now()
	req, _ := http.NewRequestWithContext(ctx, "GET", perfServer.URL+"/health", nil)
	resp, err := http.DefaultClient.Do(req)
	elapsed := time.Since(start)

	if err != nil && elapsed > 200*time.Millisecond {
		t.Errorf("cancellation took %v, expected under 200ms", elapsed)
	}
	if resp != nil {
		resp.Body.Close()
	}
	t.Logf("cancellation response time: %v", elapsed)
}

// PERF-008: SSE throughput benchmark
func TestPerf008_SSEThroughput(t *testing.T) {
	// Benchmark SSE channel throughput (without real inference)
	ch := make(chan string, 1000)
	go func() {
		for i := 0; i < 1000; i++ {
			ch <- fmt.Sprintf("chunk-%d", i)
		}
		close(ch)
	}()

	count := 0
	start := time.Now()
	for range ch {
		count++
	}
	elapsed := time.Since(start)
	t.Logf("SSE channel: %d chunks in %v (%.0f chunks/sec)", count, elapsed,
		float64(count)/elapsed.Seconds())
}