Super tiny extract route_typed_request_once (#14951)
This commit is contained in:
@@ -175,71 +175,8 @@ impl Router {
|
||||
&self.retry_config,
|
||||
// operation per attempt
|
||||
|_: u32| async {
|
||||
let worker = match self.select_worker_for_model(model_id, Some(&text)) {
|
||||
Some(w) => w,
|
||||
None => {
|
||||
RouterMetrics::record_request_error(route, "no_available_workers");
|
||||
return (
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"No available workers (all circuits open or unhealthy)",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
// Optional load tracking for cache-aware policy
|
||||
// Get the policy for this model to check if it's cache-aware
|
||||
let policy = match model_id {
|
||||
Some(model) => self.policy_registry.get_policy_or_default(model),
|
||||
None => self.policy_registry.get_default_policy(),
|
||||
};
|
||||
|
||||
let load_incremented = if policy.name() == "cache_aware" {
|
||||
increment_load(&worker);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
// Keep a clone for potential cleanup on retry
|
||||
let worker_for_cleanup = if load_incremented {
|
||||
Some(worker.clone())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
events::RequestSentEvent {
|
||||
url: worker.url().to_string(),
|
||||
}
|
||||
.emit();
|
||||
let mut headers_with_trace = headers.cloned().unwrap_or_default();
|
||||
inject_trace_context_http(&mut headers_with_trace);
|
||||
let headers = Some(&headers_with_trace);
|
||||
|
||||
let response = self
|
||||
.send_typed_request(
|
||||
headers,
|
||||
typed_req,
|
||||
route,
|
||||
worker.url(),
|
||||
is_stream,
|
||||
load_incremented,
|
||||
)
|
||||
.await;
|
||||
|
||||
events::RequestReceivedEvent {}.emit();
|
||||
|
||||
worker.record_outcome(response.status().is_success());
|
||||
|
||||
// For retryable failures, we need to decrement load since send_typed_request
|
||||
// won't have done it (it only decrements on success or non-retryable failures)
|
||||
if is_retryable_status(response.status()) && load_incremented {
|
||||
if let Some(cleanup_worker) = worker_for_cleanup {
|
||||
decrement_load(&cleanup_worker);
|
||||
}
|
||||
}
|
||||
|
||||
response
|
||||
self.route_typed_request_once(headers, typed_req, route, model_id, is_stream, &text)
|
||||
.await
|
||||
},
|
||||
// should_retry predicate
|
||||
|res, _attempt| is_retryable_status(res.status()),
|
||||
@@ -264,6 +201,82 @@ impl Router {
|
||||
response
|
||||
}
|
||||
|
||||
async fn route_typed_request_once<T: GenerationRequest + serde::Serialize + Clone>(
|
||||
&self,
|
||||
headers: Option<&HeaderMap>,
|
||||
typed_req: &T,
|
||||
route: &str,
|
||||
model_id: Option<&str>,
|
||||
is_stream: bool,
|
||||
text: &str,
|
||||
) -> Response {
|
||||
let worker = match self.select_worker_for_model(model_id, Some(text)) {
|
||||
Some(w) => w,
|
||||
None => {
|
||||
RouterMetrics::record_request_error(route, "no_available_workers");
|
||||
return (
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"No available workers (all circuits open or unhealthy)",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
// Optional load tracking for cache-aware policy
|
||||
// Get the policy for this model to check if it's cache-aware
|
||||
let policy = match model_id {
|
||||
Some(model) => self.policy_registry.get_policy_or_default(model),
|
||||
None => self.policy_registry.get_default_policy(),
|
||||
};
|
||||
|
||||
let load_incremented = if policy.name() == "cache_aware" {
|
||||
increment_load(&worker);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
// Keep a clone for potential cleanup on retry
|
||||
let worker_for_cleanup = if load_incremented {
|
||||
Some(worker.clone())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
events::RequestSentEvent {
|
||||
url: worker.url().to_string(),
|
||||
}
|
||||
.emit();
|
||||
let mut headers_with_trace = headers.cloned().unwrap_or_default();
|
||||
inject_trace_context_http(&mut headers_with_trace);
|
||||
let headers = Some(&headers_with_trace);
|
||||
|
||||
let response = self
|
||||
.send_typed_request(
|
||||
headers,
|
||||
typed_req,
|
||||
route,
|
||||
worker.url(),
|
||||
is_stream,
|
||||
load_incremented,
|
||||
)
|
||||
.await;
|
||||
|
||||
events::RequestReceivedEvent {}.emit();
|
||||
|
||||
worker.record_outcome(response.status().is_success());
|
||||
|
||||
// For retryable failures, we need to decrement load since send_typed_request
|
||||
// won't have done it (it only decrements on success or non-retryable failures)
|
||||
if is_retryable_status(response.status()) && load_incremented {
|
||||
if let Some(cleanup_worker) = worker_for_cleanup {
|
||||
decrement_load(&cleanup_worker);
|
||||
}
|
||||
}
|
||||
|
||||
response
|
||||
}
|
||||
|
||||
// Helper: return base worker URL (strips DP suffix when enabled)
|
||||
fn worker_base_url(&self, worker_url: &str) -> String {
|
||||
if self.dp_aware {
|
||||
|
||||
Reference in New Issue
Block a user