[Bugfix] Miscalculated load with cache aware routing

[simondanielsson]

Purpose

Caused by double-decrements on request counters, making the load look lower than it is.

Details

The code previously handled increments and decrements manually, which was very error prone because decrements need to happen in many places and conditions.

Solution: Use a RAII guard pattern to automatically decrement on every exit path of send_typed_request (including early returns).

There were at least two instance of double decrementing:

1. This is clearly a double count since line 906 always runs

router/src/routers/http/router.rs

Lines 882 to 909 in 5002469

	let response = match res.bytes().await {
	Ok(body) => {
	let mut response = Response::new(axum::body::Body::from(body));
	*response.status_mut() = status;
	*response.headers_mut() = response_headers;
	response
	}
	Err(e) => {
	// IMPORTANT: Decrement load on error before returning
	if load_incremented {
	if let Some(worker) = self.worker_registry.get_by_url(worker_url) {
	worker.decrement_load();
	RouterMetrics::set_running_requests(worker_url, worker.load());
	}
	}
	let error_msg = format!("Failed to get response body: {}", e);
	(StatusCode::INTERNAL_SERVER_ERROR, error_msg).into_response()
	}
	};
	// Decrement load counter for non-streaming requests if it was incremented
	if load_incremented {
	if let Some(worker) = self.worker_registry.get_by_url(worker_url) {
	worker.decrement_load();
	RouterMetrics::set_running_requests(worker_url, worker.load());
	}
	}

2. Not needed as line 906 always decrements anyways, even with retries

router/src/routers/http/router.rs

Lines 605 to 615 in 5002469

	// 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 {
	cleanup_worker.decrement_load();
	RouterMetrics::set_running_requests(
	cleanup_worker.url(),
	cleanup_worker.load(),
	);
	}
	}

Test Plan

Added new test cases to ensure load is correct in many different cases:

1. streamed/not streamed
2. retried/not retried
3. response outcome (success, retryable, not retryable, send failed)

cargo test --lib load

Test Result

Without fixes: 4 fail

running 21 tests
test core::worker::tests::test_load_counter_operations ... ok
test core::worker::tests::test_concurrent_load_decrements ... ok
test core::worker::tests::test_load_guard_multiple_workers ... ok
test core::worker::tests::test_concurrent_load_increments ... ok
test core::worker::tests::test_load_guard_single_worker ... ok
test core::worker::tests::test_total_load_calculation ... ok
test policies::cache_aware::tests::test_cache_aware_with_imbalanced_load ... ok
test policies::power_of_two::tests::test_power_of_two_with_cached_loads ... ok
test policies::cache_aware::tests::test_cache_aware_with_balanced_load ... ok
test core::worker::tests::test_load_guard_panic_safety ... ok
test core::worker::tests::test_load_counter_performance ... ok
test routers::http::router::tests::test_load_returns_to_baseline_non_streaming_success ... ok
test routers::http::router::tests::test_load_returns_to_baseline_streaming_success ... ok
test routers::http::pd_router::tests::test_load_monitor_updates ... ok
test routers::http::router::tests::test_load_returns_to_baseline_on_send_failure ... FAILED
test routers::http::router::tests::test_load_returns_to_baseline_non_streaming_non_retryable ... ok
test routers::http::router::tests::test_load_counter_not_double_decremented_on_retryable_status ... FAILED
test routers::http::router::tests::test_load_returns_to_baseline_non_streaming_retries_exhausted ... FAILED
test routers::http::router::tests::test_streaming_retryable_does_not_leak_load ... ok
test routers::http::router::tests::test_load_returns_to_baseline_streaming_with_retries ... FAILED
test tokenizer::factory::tests::test_download_tokenizer_from_hf ... ok

failures:

---- routers::http::router::tests::test_load_returns_to_baseline_on_send_failure stdout ----

thread 'routers::http::router::tests::test_load_returns_to_baseline_on_send_failure' (98149) panicked at src/routers/http/router.rs:2497:9:
assertion `left == right` failed
  left: 4
 right: 5

---- routers::http::router::tests::test_load_counter_not_double_decremented_on_retryable_status stdout ----

thread 'routers::http::router::tests::test_load_counter_not_double_decremented_on_retryable_status' (98145) panicked at src/routers/http/router.rs:2281:9:
assertion `left == right` failed: worker load must return exactly to baseline (no double decrement)
  left: 4
 right: 5

---- routers::http::router::tests::test_load_returns_to_baseline_non_streaming_retries_exhausted stdout ----

thread 'routers::http::router::tests::test_load_returns_to_baseline_non_streaming_retries_exhausted' (98147) panicked at src/routers/http/router.rs:2369:9:
assertion `left == right` failed
  left: 2
 right: 5

---- routers::http::router::tests::test_load_returns_to_baseline_streaming_with_retries stdout ----

thread 'routers::http::router::tests::test_load_returns_to_baseline_streaming_with_retries' (98151) panicked at src/routers/http/router.rs:2479:9:
assertion `left == right` failed
  left: 2
 right: 5


failures:
    routers::http::router::tests::test_load_counter_not_double_decremented_on_retryable_status
    routers::http::router::tests::test_load_returns_to_baseline_non_streaming_retries_exhausted
    routers::http::router::tests::test_load_returns_to_baseline_on_send_failure
    routers::http::router::tests::test_load_returns_to_baseline_streaming_with_retries

test result: FAILED. 17 passed; 4 failed; 0 ignored; 0 measured; 471 filtered out; finished in 1.47s

With fix:

running 22 tests
test core::worker::tests::test_load_counter_operations ... ok
test core::worker::tests::test_load_guard_multiple_workers ... ok
test core::worker::tests::test_load_guard_panic_safety ... ok
test core::worker::tests::test_concurrent_load_decrements ... ok
test core::worker::tests::test_concurrent_load_increments ... ok
test core::worker::tests::test_load_guard_single_worker ... ok
test core::worker::tests::test_total_load_calculation ... ok
test policies::cache_aware::tests::test_cache_aware_with_imbalanced_load ... ok
test policies::cache_aware::tests::test_cache_aware_with_balanced_load ... ok
test policies::power_of_two::tests::test_power_of_two_with_cached_loads ... ok
test core::worker::tests::test_load_counter_performance ... ok
test routers::http::pd_router::tests::test_load_monitor_updates ... ok
test routers::http::router::tests::test_load_returns_to_baseline_on_dp_parse_error ... ok
test routers::http::router::tests::test_load_returns_to_baseline_on_send_failure ... ok
test routers::http::router::tests::test_load_returns_to_baseline_streaming_success ... ok
test routers::http::router::tests::test_load_returns_to_baseline_non_streaming_non_retryable ... ok
test routers::http::router::tests::test_load_returns_to_baseline_non_streaming_success ... ok
test routers::http::router::tests::test_load_returns_to_baseline_non_streaming_retries_exhausted ... ok
test routers::http::router::tests::test_load_counter_not_double_decremented_on_retryable_status ... ok
test routers::http::router::tests::test_load_returns_to_baseline_streaming_with_retries ... ok
test routers::http::router::tests::test_streaming_retryable_does_not_leak_load ... ok
test tokenizer::factory::tests::test_download_tokenizer_from_hf ... ok

test result: ok. 22 passed; 0 failed; 0 ignored; 0 measured; 471 filtered out; finished in 1.20s

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results

[simondanielsson]

@BugenZhao Happy to get a review when you have time 🙏

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3e68b6c46f

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[simondanielsson]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 794d18c571

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[simondanielsson]

@codex review

[simondanielsson]

Failing Trivy CI fixed in : #212

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Hooray!

Reviewed commit: 7d1227ce58

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".