204 lines
6.4 KiB
Go
204 lines
6.4 KiB
Go
package server
|
|
|
|
import (
|
|
"ai-router/internal/protocol"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// Explicitly opt in: ROUTER_PERF=1 go test ./internal/server -run TestPerformanceHTTP -v -count=1.
|
|
func TestPerformanceHTTP(t *testing.T) {
|
|
if os.Getenv("ROUTER_PERF") != "1" {
|
|
t.Skip("opt-in load test")
|
|
}
|
|
for _, stream := range []bool{false, true} {
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
io.Copy(io.Discard, r.Body)
|
|
r.Body.Close()
|
|
if stream {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
for _, f := range strings.Split(integrationSSE("openai-chat"), "\n\n") {
|
|
if f != "" {
|
|
fmt.Fprint(w, f+"\n\n")
|
|
w.(http.Flusher).Flush()
|
|
}
|
|
}
|
|
} else {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
fmt.Fprint(w, integrationResponse("openai-chat"))
|
|
}
|
|
}))
|
|
for _, mode := range []string{"direct", "memory", "disk"} {
|
|
for _, down := range []string{"openai-chat", "anthropic"} {
|
|
if mode == "direct" && down != "openai-chat" {
|
|
continue
|
|
}
|
|
for _, concurrency := range []int{1, 8, 32} {
|
|
t.Run(fmt.Sprintf("stream=%t/%s/%s/c%d", stream, mode, down, concurrency), func(t *testing.T) {
|
|
target := up.URL + integrationPath("openai-chat")
|
|
raw := "test-key"
|
|
var s *Server
|
|
var router *httptest.Server
|
|
if mode != "direct" {
|
|
path := ":memory:"
|
|
if mode == "disk" {
|
|
path = filepath.Join(t.TempDir(), "router.db")
|
|
}
|
|
var err error
|
|
s, err = New(path, "test-password", nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer s.Close()
|
|
if err := s.SetLimits(1000000, 64, 128); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
cookie := loginCookie(t, s)
|
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
|
key := integrationCreateKey(t, s, cookie, "perf", ch)
|
|
raw = key.raw
|
|
router = httptest.NewServer(s)
|
|
defer router.Close()
|
|
target = router.URL + integrationPath(down)
|
|
}
|
|
tr := http.DefaultTransport.(*http.Transport).Clone()
|
|
tr.MaxIdleConnsPerHost = 64
|
|
defer tr.CloseIdleConnections()
|
|
client := &http.Client{Transport: tr, Timeout: 10 * time.Second}
|
|
body := integrationRequest(down, "abc/gpt-5", stream)
|
|
request := func() (time.Duration, time.Duration, bool) {
|
|
start := time.Now()
|
|
r, _ := http.NewRequest("POST", target, strings.NewReader(body))
|
|
r.Header.Set("Authorization", "Bearer "+raw)
|
|
r.Header.Set("Content-Type", "application/json")
|
|
resp, err := client.Do(r)
|
|
if err != nil {
|
|
return time.Since(start), 0, false
|
|
}
|
|
first := make([]byte, 1)
|
|
_, err = resp.Body.Read(first)
|
|
ttfb := time.Since(start)
|
|
_, e := io.Copy(io.Discard, resp.Body)
|
|
resp.Body.Close()
|
|
return time.Since(start), ttfb, resp.StatusCode == 200 && (err == nil || err == io.EOF) && e == nil
|
|
}
|
|
for i := 0; i < 20; i++ {
|
|
if _, _, ok := request(); !ok {
|
|
t.Fatal("warmup failed")
|
|
}
|
|
}
|
|
const n = 400
|
|
dur := make([]float64, n)
|
|
first := make([]float64, n)
|
|
var index atomic.Int64
|
|
var failures atomic.Int64
|
|
var wg sync.WaitGroup
|
|
runtime.GC()
|
|
var before, after runtime.MemStats
|
|
runtime.ReadMemStats(&before)
|
|
start := time.Now()
|
|
for c := 0; c < concurrency; c++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
for {
|
|
i := int(index.Add(1)) - 1
|
|
if i >= n {
|
|
return
|
|
}
|
|
d, f, ok := request()
|
|
dur[i] = float64(d) / 1e6
|
|
first[i] = float64(f) / 1e6
|
|
if !ok {
|
|
failures.Add(1)
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
elapsed := time.Since(start)
|
|
runtime.ReadMemStats(&after)
|
|
sort.Float64s(dur)
|
|
sort.Float64s(first)
|
|
var total float64
|
|
for _, d := range dur {
|
|
total += d
|
|
}
|
|
drained := elapsed
|
|
if s != nil {
|
|
router.Close()
|
|
if err := s.flushLogs(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
drained = time.Since(start)
|
|
var count int
|
|
if err := s.db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&count); err != nil || count != n+20 {
|
|
t.Fatalf("log count %d err %v", count, err)
|
|
}
|
|
}
|
|
t.Logf("RESULT rps=%.1f drained_rps=%.1f mean_ms=%.3f p50_ms=%.3f p95_ms=%.3f p99_ms=%.3f first_p50_ms=%.3f first_p95_ms=%.3f alloc_B_req=%d errors=%d", float64(n)/elapsed.Seconds(), float64(n)/drained.Seconds(), total/n, dur[n/2], dur[n*95/100], dur[n*99/100], first[n/2], first[n*95/100], (after.TotalAlloc-before.TotalAlloc)/n, failures.Load())
|
|
if failures.Load() != 0 {
|
|
t.Fail()
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
up.Close()
|
|
}
|
|
}
|
|
|
|
func BenchmarkPerformanceConversion(b *testing.B) {
|
|
for _, size := range []int{1024, 65536} {
|
|
for _, to := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
|
|
b.Run(fmt.Sprintf("%d/%s", size, to), func(b *testing.B) {
|
|
body := []byte(`{"model":"abc/gpt-5","max_tokens":1024,"messages":[{"role":"user","content":"` + strings.Repeat("x", size) + `"}]}`)
|
|
b.ReportAllocs()
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
if _, err := protocol.ConvertRequest(protocol.Chat, to, body, "gpt-4"); err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func BenchmarkPerformanceStats(b *testing.B) {
|
|
for _, n := range []int{1000, 10000, 100000} {
|
|
b.Run(fmt.Sprint(n), func(b *testing.B) {
|
|
s, err := New(filepath.Join(b.TempDir(), "stats.db"), "test-password", nil)
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
defer s.Close()
|
|
_, err = s.db.Exec(`WITH RECURSIVE seq(x) AS (SELECT 1 UNION ALL SELECT x+1 FROM seq WHERE x<?) INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) SELECT '2026-10-20T10:00:00Z','key','perf','abc','abc/gpt-5','gpt-4','openai-chat','openai-chat',200,10,0,1000,100,100,0,1,'','{}' FROM seq`, n)
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
r := httptest.NewRequest("GET", "/admin/api/stats", nil)
|
|
b.ReportAllocs()
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
w := httptest.NewRecorder()
|
|
s.stats(w, r)
|
|
if w.Code != 200 {
|
|
b.Fatal(w.Code)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|