lib/llm/src/mocker/scheduler.rs (1)

294-346: 🛠️ Refactor suggestion

⚠️ Potential issue

Locks are held across an .await; this blocks every other task

state_guard and kv_manager_guard stay alive while the task sleeps:

let adjusted_time = Duration::from_secs_f64(total_time.as_secs_f64() / speedup_ratio);
if adjusted_time.as_millis() > 0 {
    tokio::time::sleep(adjusted_time).await;   // <-- await with locks held
}

While sleeping, all other async tasks that need the same Mutexes (e.g. receive, metrics queries, other scheduler ticks) will starve, drastically reducing concurrency and possibly causing dead-locks when used with additional locks elsewhere.

Fix by dropping the guards before the await:

-// Sleep once for the adjusted duration
-let adjusted_time = Duration::from_secs_f64(total_time.as_secs_f64() / speedup_ratio);
-if adjusted_time.as_millis() > 0 {
-    tokio::time::sleep(adjusted_time).await;
-}
+// Sleep once for the adjusted duration
+let adjusted_time = Duration::from_secs_f64(total_time.as_secs_f64() / speedup_ratio);
+
+// Explicitly drop locks to avoid holding them across await points
+drop(state_guard);
+drop(kv_manager_guard);
+
+if adjusted_time.as_millis() > 0 {
+    tokio::time::sleep(adjusted_time).await;
+}

Please scope the guard usages or call drop() as above to restore non-blocking behaviour.

lib/llm/src/mocker/protocols.rs (1)

58-63: Consider documenting the new new_blocks field

PrefillCost now exposes new_blocks, but the struct-level doc comment still talks about “the cost of prefilling content” without explaining that both token and block deltas are returned. A short description will make downstream usage clearer and avoid confusion with the pre–refactor signature.

 /// Represents the cost of prefilling content in the cache
-#[derive(Debug, Clone, Serialize, Deserialize)]
+/// `new_blocks` – number of blocks not currently resident in **active** cache  
+/// `new_tokens` – number of tokens that need to be copied for those blocks  
+/// `prefill_compute` – empirical latency model (µs)
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PrefillCost {

lib/llm/src/mocker/sequence.rs (3)

69-76: Parameter name still indicates an option, but the struct has a fixed block_size

The constructor now takes block_size: Option<usize>, which is fine, yet the earlier discussion in the PR implied the removal of a configurable size.
If external callers are no longer expected to pass None (i.e. we always default to 64) consider:

1. Dropping the Option wrapper and accepting usize directly, or
2. Renaming to block_size_override (or similar) to clarify semantics.

This will prevent accidental None propagation or double-wrapped options elsewhere.

---

107-110: new_with_signal leaks the same default-handling logic twice

new_with_signal forwards to Self::new which already expands the default for block_size.
If you accept the previous suggestion and drop the Option, you can simplify the call and the tests that pass Some(16).

---

229-234: Tests rely on magic numbers – make them constants

Some(16) is repeated in every test case. Replacing the literal with a const BLOCK: usize = 16; at the module level improves readability and prevents drift when the default changes.

lib/llm/src/mocker/kv_manager.rs (1)

203-206: Floating-point division on integers – cast once

A micro-nit: you convert both operands every call. Caching max_capacity as f64 in the struct saves a few casts in hot loops.

lib/llm/src/mocker/evictor.rs (1)

108-113: Risk of i64 counter under-/overflow & eventual wrap-around

negative_counter monotonically decreases and positive_counter increases forever.
With long-running services (billions of insertions) the counters may reach the limits of i64, wrapping to the opposite sign and breaking ordering guarantees.

Mitigations:

• Use i128, or reset both counters when they approach the limits (e.g. every 2^60 ops) while re-normalising all stored counters.
• Alternatively, store (Instant, u64) or (u64, GenerationId) and rely on wrapping add with safe tiebreakers.

Although unlikely in tests, overflow bugs manifest in persistent services.

lib/llm/src/mocker/scheduler.rs (1)

108-113: Decode LRU is not refreshed on every access

run() merely checks membership; it does not update recency, so eviction order is insertion order, not true LRU.
If the intent is genuine LRU pre-emption, consider calling decode.insert(uuid) inside run() (or providing a specialised touch method) to bump the counter each time the sequence is decoded.

This is a functional not a correctness bug – but it biases eviction unfairly.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro