Support tracking worker routing key loads in gateway (#16765)

This commit is contained in:
fzyzcjy
2026-01-13 08:07:17 +08:00
committed by GitHub
parent d3c08fb07c
commit 1f0e3d7fd8
8 changed files with 287 additions and 24 deletions
@@ -37,7 +37,7 @@ async fn test_guard_dropped_when_response_body_consumed() {
let response = Response::new(body);
// Attach guard
let guard = WorkerLoadGuard::new(worker.clone());
let guard = WorkerLoadGuard::new(worker.clone(), None);
assert_eq!(worker.load(), 1);
let guarded_response = AttachedBody::wrap_response(response, guard);
@@ -62,7 +62,7 @@ async fn test_guard_dropped_when_response_dropped_without_consumption() {
let body = Body::from("Hello, World!");
let response = Response::new(body);
let guard = WorkerLoadGuard::new(worker.clone());
let guard = WorkerLoadGuard::new(worker.clone(), None);
assert_eq!(worker.load(), 1);
let _guarded_response = AttachedBody::wrap_response(response, guard);
@@ -86,7 +86,7 @@ async fn test_streaming_guard_dropped_when_stream_ends() {
let (tx, rx) = mpsc::unbounded_channel::<Bytes>();
let response = create_sse_response(rx);
let guard = WorkerLoadGuard::new(worker.clone());
let guard = WorkerLoadGuard::new(worker.clone(), None);
assert_eq!(worker.load(), 1);
let guarded_response = AttachedBody::wrap_response(response, guard);
@@ -133,7 +133,7 @@ async fn test_streaming_guard_dropped_on_client_disconnect() {
let (tx, rx) = mpsc::unbounded_channel::<Bytes>();
let response = create_sse_response(rx);
let guard = WorkerLoadGuard::new(worker.clone());
let guard = WorkerLoadGuard::new(worker.clone(), None);
assert_eq!(worker.load(), 1);
let guarded_response = AttachedBody::wrap_response(response, guard);
@@ -171,8 +171,8 @@ async fn test_multiple_guards_all_dropped() {
let response = Response::new(body);
// Create guards for both workers (simulates dual prefill/decode)
let guard1 = WorkerLoadGuard::new(worker1.clone());
let guard2 = WorkerLoadGuard::new(worker2.clone());
let guard1 = WorkerLoadGuard::new(worker1.clone(), None);
let guard2 = WorkerLoadGuard::new(worker2.clone(), None);
assert_eq!(worker1.load(), 1);
assert_eq!(worker2.load(), 1);
@@ -197,7 +197,7 @@ async fn test_guard_with_empty_body() {
let body = Body::empty();
let response = Response::new(body);
let guard = WorkerLoadGuard::new(worker.clone());
let guard = WorkerLoadGuard::new(worker.clone(), None);
assert_eq!(worker.load(), 1);
let guarded_response = AttachedBody::wrap_response(response, guard);