fix: self-heal in-flight counts by pruning stale request ids
On a long-lived SSE connection a missed remove leaves a stale request id in the tracker forever; llama-swap only sends a fresh snapshot on connect, so counts drifted up and stayed stuck. The tracker now timestamps each request id (upserts refresh it) and prunes ids not updated within 120s, on every inflight event and on a 30s runtime interval. Removes also accept numeric ids defensively.
This commit is contained in:
@@ -17,7 +17,7 @@ export function decodeEvent(msg: SseMessage): FeedEvent | null {
|
||||
if (outer.type === "inflight") {
|
||||
const operation = inner.operation as string;
|
||||
if (operation === "remove") {
|
||||
const id = typeof inner.id === "string" ? inner.id : undefined;
|
||||
const id = typeof inner.id === "string" ? inner.id : typeof inner.id === "number" ? String(inner.id) : undefined;
|
||||
if (id) return { type: "inflight", operation: "remove", id };
|
||||
return null;
|
||||
}
|
||||
|
||||
@@ -21,27 +21,35 @@ export type FeedEvent =
|
||||
export type ModelRuntimeState = "stopped" | "loading" | "ready";
|
||||
|
||||
export class InflightTracker {
|
||||
private requests = new Map<string, string>();
|
||||
private requests = new Map<string, { model: string; seen: number }>();
|
||||
private states = new Map<string, ModelRuntimeState>();
|
||||
|
||||
constructor(private now: () => number = () => Date.now()) {}
|
||||
|
||||
apply(event: FeedEvent): void {
|
||||
if (event.type === "inflight") {
|
||||
if (event.operation === "snapshot") {
|
||||
this.requests = new Map<string, string>();
|
||||
for (const r of event.requests) if (r.id) this.requests.set(r.id, r.model);
|
||||
this.requests = new Map<string, { model: string; seen: number }>();
|
||||
for (const r of event.requests) if (r.id) this.requests.set(r.id, { model: r.model, seen: this.now() });
|
||||
} else if (event.operation === "add") {
|
||||
for (const r of event.requests) if (r.id) this.requests.set(r.id, r.model);
|
||||
for (const r of event.requests) if (r.id) this.requests.set(r.id, { model: r.model, seen: this.now() });
|
||||
} else if (event.operation === "remove") {
|
||||
this.requests.delete(event.id);
|
||||
}
|
||||
this.prune();
|
||||
} else if (event.type === "modelStatus") {
|
||||
for (const m of event.models) this.states.set(m.id, normalizeState(m.state));
|
||||
}
|
||||
}
|
||||
|
||||
prune(maxAgeMs = 120_000): void {
|
||||
const cutoff = this.now() - maxAgeMs;
|
||||
for (const [id, entry] of this.requests) if (entry.seen < cutoff) this.requests.delete(id);
|
||||
}
|
||||
|
||||
count(modelId: string): number {
|
||||
let n = 0;
|
||||
for (const model of this.requests.values()) if (model === modelId) n++;
|
||||
for (const entry of this.requests.values()) if (entry.model === modelId) n++;
|
||||
return n;
|
||||
}
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ class Runtime {
|
||||
private feed?: EventFeed;
|
||||
private poller?: MetricsPoller;
|
||||
private statsCache?: StatsCache;
|
||||
private pruneTimer?: ReturnType<typeof setInterval>;
|
||||
private listeners = new Set<() => void>();
|
||||
|
||||
ensureConnections(cfg: LlamaSwapConfig): void {
|
||||
@@ -45,6 +46,9 @@ class Runtime {
|
||||
this.statsCache = new StatsCache(cfg);
|
||||
this.statsCache.onChange(() => this.emit());
|
||||
}
|
||||
if (!this.pruneTimer) {
|
||||
this.pruneTimer = setInterval(() => this.tracker.prune(), PRUNE_INTERVAL_MS);
|
||||
}
|
||||
}
|
||||
|
||||
get pollerInstance(): MetricsPoller | undefined {
|
||||
@@ -73,3 +77,5 @@ class Runtime {
|
||||
}
|
||||
|
||||
export const runtime = new Runtime();
|
||||
|
||||
const PRUNE_INTERVAL_MS = 30_000;
|
||||
|
||||
@@ -67,6 +67,15 @@ test("decodeEvent parses an inflight remove with a bare id", () => {
|
||||
assert.deepEqual(ev, { type: "inflight", operation: "remove", id: "817" });
|
||||
});
|
||||
|
||||
test("decodeEvent accepts a numeric id on remove", () => {
|
||||
const msg: SseMessage = {
|
||||
event: "message",
|
||||
data: JSON.stringify({ type: "inflight", data: JSON.stringify({ operation: "remove", id: 820 }) }),
|
||||
};
|
||||
const ev = decodeEvent(msg);
|
||||
assert.deepEqual(ev, { type: "inflight", operation: "remove", id: "820" });
|
||||
});
|
||||
|
||||
test("decodeEvent parses modelStatus", () => {
|
||||
const msg: SseMessage = {
|
||||
event: "message",
|
||||
|
||||
@@ -73,3 +73,34 @@ test("total sums in-flight requests across all models", () => {
|
||||
tracker.apply({ type: "inflight", operation: "remove", id: "1" });
|
||||
assert.equal(tracker.total(), 2);
|
||||
});
|
||||
|
||||
test("prune drops request ids that stopped receiving updates", () => {
|
||||
let now = 0;
|
||||
const tracker = new InflightTracker(() => now);
|
||||
tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "1" }] });
|
||||
now = 200_000;
|
||||
tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "2" }] });
|
||||
tracker.prune(60_000);
|
||||
assert.equal(tracker.total(), 1);
|
||||
assert.equal(tracker.count("A"), 1);
|
||||
});
|
||||
|
||||
test("active requests refreshed by upserts survive pruning", () => {
|
||||
let now = 0;
|
||||
const tracker = new InflightTracker(() => now);
|
||||
tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "1" }] });
|
||||
now = 200_000;
|
||||
tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "1" }] });
|
||||
tracker.prune(60_000);
|
||||
assert.equal(tracker.total(), 1);
|
||||
});
|
||||
|
||||
test("prune is applied on every inflight event", () => {
|
||||
let now = 0;
|
||||
const tracker = new InflightTracker(() => now);
|
||||
tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "stale" }] });
|
||||
now = 200_000;
|
||||
tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "B", id: "live" }] });
|
||||
assert.equal(tracker.total(), 1);
|
||||
assert.equal(tracker.count("B"), 1);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user