fix(quota-share): give each in-flight request its own lease

changelog.d/fixes/11547-quota-share-inflight-lease.md

@@ -0,0 +1 @@
+- **fix(quota-share):** in-flight leases are per request, so an aborted request's slot ages out instead of being kept alive by later traffic on the same connection ([#11547](https://github.com/diegosouzapw/OmniRoute/pull/11547)) — thanks @abhisheksharma2411

open-sse/services/combo/quotaShareInflight.ts

@@ -6,13 +6,21 @@
  *
  * Decrement-on-abort safety (TTL/lease):
  *   The generic combo dispatch path is intentionally NOT instrumented (so this
- *   feature cannot regress existing strategies). Instead, each in-flight slot
- *   carries an expiry: incrementInflight() stamps `nowMs + leaseMs`. The normal
- *   path calls decrementInflight() (returned to the caller as a callback) once
- *   the request settles, which clears the slot immediately. If a request is
- *   aborted or crashes before that callback runs, the slot still auto-expires
- *   after DEFAULT_LEASE_MS — so the counter can never leak forever, even without
- *   touching the generic dispatch.
+ *   feature cannot regress existing strategies). Instead, EACH IN-FLIGHT REQUEST
+ *   carries its own expiry: incrementInflight() appends `nowMs + leaseMs`. The
+ *   normal path calls decrementInflight() (returned to the caller as a callback)
+ *   once the request settles, which retires one lease immediately. If a request
+ *   is aborted or crashes before that callback runs, only that request's lease
+ *   remains, and it expires after DEFAULT_LEASE_MS — so the counter cannot leak
+ *   forever, even without touching the generic dispatch.
+ *
+ *   The lease is per request rather than per connection on purpose. A single
+ *   shared `expiresAtMs` was refreshed by every subsequent increment on the same
+ *   connection, so under sustained traffic an orphaned count rode along
+ *   indefinitely and never expired — the exact leak this mechanism exists to
+ *   bound. It also meant a request settling after that shared lease lapsed
+ *   deleted the whole entry, zeroing the count for its still-active neighbours
+ *   and presenting a busy connection to P2C as idle.
  *
  * Fail-open: getInflight() returns 0 for an unknown / empty connectionId.
  * All time input is injectable (the `nowMs` param) so unit tests drive the
@@ -35,11 +43,24 @@ export const DEFAULT_LEASE_MS = 120_000;
 // Types
 // ---------------------------------------------------------------------------
 
+/**
+ * Expiry timestamps for the requests currently in flight on one connection,
+ * kept in ascending order. The count IS `leases.length` — there is no separate
+ * counter to drift out of step with the leases.
+ */
 interface InflightSlot {
-  count: number;
-  expiresAtMs: number;
+  leases: number[];
 }
 
+/**
+ * Upper bound on tracked leases per connection. A connection with more than this
+ * many simultaneous in-flight requests is already far past any sane concurrency
+ * cap; beyond the bound the oldest lease is retired so memory stays bounded.
+ * Under-counting a saturated connection is the fail-open direction this module
+ * already takes elsewhere (getInflight returns 0 for anything unknown).
+ */
+const MAX_LEASES_PER_CONNECTION = 4096;
+
 // ---------------------------------------------------------------------------
 // In-process store. Key: connectionId.
 // ---------------------------------------------------------------------------
@@ -65,11 +86,14 @@ export function incrementInflight(
 ): number {
   if (!connectionId) return 0;
   pruneExpired(nowMs);
-  const slot = _inflightMap.get(connectionId);
-  const base = slot && slot.expiresAtMs > nowMs ? slot.count : 0;
-  const newCount = base + 1;
-  _inflightMap.set(connectionId, { count: newCount, expiresAtMs: nowMs + leaseMs });
-  return newCount;
+  const slot = _inflightMap.get(connectionId) ?? { leases: [] };
+  // Appending keeps `leases` ascending whenever leaseMs is constant, which is the
+  // only case in production; a shorter bespoke lease can land out of order, and
+  // retiring the earliest expiry below stays correct either way.
+  slot.leases.push(nowMs + leaseMs);
+  if (slot.leases.length > MAX_LEASES_PER_CONNECTION) slot.leases.shift();
+  _inflightMap.set(connectionId, slot);
+  return slot.leases.length;
 }
 
 /**
@@ -82,16 +106,43 @@ export function incrementInflight(
 export function decrementInflight(connectionId: string, nowMs: number = Date.now()): void {
   if (!connectionId) return;
   const slot = _inflightMap.get(connectionId);
-  if (!slot || slot.expiresAtMs <= nowMs) {
-    _inflightMap.delete(connectionId);
-    return;
+  if (!slot) return;
+
+  // A release says "one request settled" without saying which, so the lease to
+  // retire has to be inferred. Two rules, in order:
+  //
+  //  1. If any lease has already expired, retire the newest of those. A request
+  //     that outlived its own lease is by definition the longest-running one, so
+  //     an expired lease is the best match for the caller — and consuming it
+  //     leaves the still-live neighbours alone. Retiring a live lease here
+  //     instead would decrement twice for one settled request: once when the
+  //     expiry lapsed, once again now.
+  //  2. Otherwise retire the newest live lease. Never the oldest: a normal
+  //     request's release would then retire an *orphaned* lease and leave its
+  //     own newer one behind, so the orphan never ages out on its own schedule
+  //     and the count never converges — which is the leak the lease exists to
+  //     bound.
+  const expiredIndex = lastIndexWhere(slot.leases, (expiresAtMs) => expiresAtMs <= nowMs);
+  if (expiredIndex >= 0) {
+    slot.leases.splice(expiredIndex, 1);
+  } else {
+    slot.leases.pop();
   }
-  const newCount = Math.max(0, slot.count - 1);
-  if (newCount === 0) {
+  retireExpired(slot, nowMs);
+
+  if (slot.leases.length === 0) {
     _inflightMap.delete(connectionId);
   } else {
-    _inflightMap.set(connectionId, { count: newCount, expiresAtMs: slot.expiresAtMs });
+    _inflightMap.set(connectionId, slot);
+  }
+}
+
+/** Index of the last element satisfying `predicate`, or -1. */
+function lastIndexWhere(values: number[], predicate: (value: number) => boolean): number {
+  for (let i = values.length - 1; i >= 0; i--) {
+    if (predicate(values[i]!)) return i;
   }
+  return -1;
 }
 
 /**
@@ -103,18 +154,26 @@ export function decrementInflight(connectionId: string, nowMs: number = Date.now
 export function getInflight(connectionId: string, nowMs: number = Date.now()): number {
   if (!connectionId) return 0;
   const slot = _inflightMap.get(connectionId);
-  if (!slot || slot.expiresAtMs <= nowMs) return 0;
-  return slot.count;
+  if (!slot) return 0;
+  retireExpired(slot, nowMs);
+  return slot.leases.length;
 }
 
 // ---------------------------------------------------------------------------
 // Internal helpers
 // ---------------------------------------------------------------------------
 
-/** Drop all expired slots so the map cannot grow unbounded with stale leases. */
+/** Drop one slot's expired leases in place. */
+function retireExpired(slot: InflightSlot, nowMs: number): void {
+  if (slot.leases.length === 0) return;
+  slot.leases = slot.leases.filter((expiresAtMs) => expiresAtMs > nowMs);
+}
+
+/** Drop all expired leases so the map cannot grow unbounded with stale entries. */
 function pruneExpired(nowMs: number): void {
   for (const [key, slot] of _inflightMap) {
-    if (slot.expiresAtMs <= nowMs) _inflightMap.delete(key);
+    retireExpired(slot, nowMs);
+    if (slot.leases.length === 0) _inflightMap.delete(key);
   }
 }
 

tests/unit/quota-share-strategy.test.ts

@@ -102,6 +102,52 @@ describe("quotaShareInflight", () => {
     assert.equal(getInflight("conn-b", NOW + leaseMs + 1), 0);
   });
 
+  test("an orphaned lease expires even while the connection keeps taking traffic", () => {
+    // The lease exists to bound a leak from a request that aborts before its
+    // release callback runs. A single per-connection expiry was refreshed by
+    // every later increment, so under sustained traffic the orphaned count rode
+    // along forever — the leak was unbounded exactly when the connection was
+    // busy, which is when P2C's view of load matters most.
+    const minute = 60_000;
+    incrementInflight("conn-leak", DEFAULT_LEASE_MS, NOW); // request that never releases
+
+    for (let i = 1; i <= 10; i++) {
+      const at = NOW + i * minute;
+      incrementInflight("conn-leak", DEFAULT_LEASE_MS, at);
+      decrementInflight("conn-leak", at + 1);
+    }
+
+    // 10 minutes on, five times the 120s lease: the orphan must be gone.
+    assert.equal(getInflight("conn-leak", NOW + 10 * minute + 2), 0);
+  });
+
+  test("a late release retires only its own lease, not its neighbours'", () => {
+    // Two concurrent requests on one connection. The first settles after the
+    // lease it was stamped with has lapsed. It must not take the second
+    // request's live slot with it — presenting a busy connection as idle sends
+    // P2C straight at the connection that is already loaded.
+    // First at NOW (lease lapses at NOW+120s), second at NOW+100s (lapses at
+    // NOW+220s). Releasing the first at NOW+130s is therefore a release that
+    // arrives after its own lease expired but while its neighbour is still live.
+    incrementInflight("conn-pair", DEFAULT_LEASE_MS, NOW);
+    incrementInflight("conn-pair", DEFAULT_LEASE_MS, NOW + 100_000);
+    assert.equal(getInflight("conn-pair", NOW + 100_000), 2);
+
+    decrementInflight("conn-pair", NOW + 130_000);
+
+    assert.equal(getInflight("conn-pair", NOW + 130_000), 1);
+  });
+
+  test("expired leases are retired without an explicit decrement, per request", () => {
+    const leaseMs = 1_000;
+    incrementInflight("conn-mix", leaseMs, NOW);
+    incrementInflight("conn-mix", leaseMs * 10, NOW);
+    assert.equal(getInflight("conn-mix", NOW), 2);
+    // The short lease lapses; the long one does not.
+    assert.equal(getInflight("conn-mix", NOW + leaseMs + 1), 1);
+    assert.equal(getInflight("conn-mix", NOW + leaseMs * 10 + 1), 0);
+  });
+
   test("empty connectionId returns 0 (fail-open)", () => {
     assert.equal(getInflight("", NOW), 0);
     incrementInflight("", DEFAULT_LEASE_MS, NOW); // must not throw / not store