chore: initial commit
This commit is contained in:
@@ -0,0 +1,203 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user