110 lines
3.6 KiB
Go
110 lines
3.6 KiB
Go
package proxy_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.wntrmute.dev/kyle/crossbar/internal/store"
|
|
)
|
|
|
|
// A client that goes away mid-stream is still a request that happened: it held a slot, it cost
|
|
// prefill, and it belongs in the accounting. The row records status 499 and a non-empty err.
|
|
func TestClientCancelMidStreamIsRecorded(t *testing.T) {
|
|
slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/health":
|
|
fmt.Fprint(w, `{"status":"ok"}`)
|
|
case "/v1/models":
|
|
fmt.Fprint(w, `{"object":"list","data":[{"id":"shared"}]}`)
|
|
default:
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(200)
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"first\"}}]}\n\n")
|
|
w.(http.Flusher).Flush()
|
|
select {
|
|
case <-r.Context().Done():
|
|
case <-time.After(3 * time.Second):
|
|
}
|
|
}
|
|
}))
|
|
t.Cleanup(slow.Close)
|
|
beta := newUpstream(t, "beta")
|
|
r := newRig(t, twoHosts, &upstream{name: "alpha", srv: slow}, beta)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
body := `{"model":"alpha-only","stream":true,"messages":[{"role":"user","content":"cancel me"}]}`
|
|
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, r.front.URL+"/r/v1/chat/completions", strings.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
buf := make([]byte, 64)
|
|
if _, err := resp.Body.Read(buf); err != nil {
|
|
t.Fatalf("first chunk: %v", err)
|
|
}
|
|
cancel()
|
|
resp.Body.Close()
|
|
|
|
deadline := time.Now().Add(3 * time.Second)
|
|
var counts []store.StatusCount
|
|
for time.Now().Before(deadline) {
|
|
counts, _ = r.store.StatusCounts(time.Time{})
|
|
if len(counts) > 0 {
|
|
break
|
|
}
|
|
time.Sleep(25 * time.Millisecond)
|
|
}
|
|
if len(counts) != 1 || counts[0].Status != 499 || counts[0].Route != "r" || counts[0].Count != 1 {
|
|
t.Fatalf("status counts after a cancelled stream = %+v, want one row: route r, status 499", counts)
|
|
}
|
|
rows, _ := r.store.Usage(time.Time{}, store.ByRoute)
|
|
if len(rows) != 1 || rows[0].Requests != 1 || rows[0].Errors != 1 {
|
|
t.Errorf("usage = %+v, want 1 request counted as an error", rows)
|
|
}
|
|
}
|
|
|
|
// The same when the client gives up while waiting in the queue: a 499 row, no slot leaked.
|
|
func TestClientCancelWhileQueuedIsRecorded(t *testing.T) {
|
|
alpha := newUpstream(t, "alpha")
|
|
alpha.delay = 800 * time.Millisecond
|
|
r := newRig(t, `
|
|
listen = "127.0.0.1:1"
|
|
queue_max = 2
|
|
[hosts.alpha]
|
|
base_url = %q
|
|
models = { "shared" = { parallel = 1 } }
|
|
[routes.r]
|
|
hosts = ["alpha"]
|
|
default_model = "shared"
|
|
`, alpha)
|
|
go func() { drain(r.post("/r/v1/chat/completions", conversation(1, 1))) }() // holds the one slot
|
|
waitUntil(t, func() bool { return r.lim.InFlight("alpha", "shared") == 1 })
|
|
ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
|
|
defer cancel()
|
|
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, r.front.URL+"/r/v1/chat/completions", strings.NewReader(conversation(2, 1)))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if _, err := http.DefaultClient.Do(req); err == nil {
|
|
t.Fatal("the queued request should have been cancelled by its context")
|
|
}
|
|
deadline := time.Now().Add(3 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
counts, _ := r.store.StatusCounts(time.Time{})
|
|
for _, c := range counts {
|
|
if c.Status == 499 {
|
|
if r.lim.Queued("alpha", "shared") != 0 {
|
|
t.Errorf("queued = %d after the waiter cancelled", r.lim.Queued("alpha", "shared"))
|
|
}
|
|
return
|
|
}
|
|
}
|
|
time.Sleep(25 * time.Millisecond)
|
|
}
|
|
t.Fatal("no 499 row recorded for the request cancelled while queued")
|
|
}
|