Let coalesced live searches hand back their admission slot

packages/realm-server/handlers/handle-search.ts

@@ -22,6 +22,7 @@ import {
 } from '@cardstack/runtime-common';
 import {
   fetchRequestFromContext,
+  releaseSearchAdmission,
   sendResponseForBadRequest,
   setContextResponse,
 } from '../middleware/index.ts';
@@ -381,16 +382,25 @@ async function respondWithJobScopedSearchCache(
       query,
       opts: { ...(args.opts as Record<string, unknown>), generations },
       populate: runSearch,
+      // A joiner or a hit holds no result document of its own, so it stops
+      // counting toward the search admission ceiling here rather than when
+      // its response ends; the ceiling is then a bound on concurrent
+      // computations, which is what holds the heap.
+      onOutcome: (decided) => {
+        if (decided !== 'miss') {
+          releaseSearchAdmission(ctxt);
+        }
+      },
     });
-    await setContextResponse(
-      ctxt,
-      new Response(body, {
-        headers: {
-          'content-type': SupportedMimeType.CardJson,
-          [LIVE_SEARCH_CACHE_HEADER]: outcome,
-        },
-      }),
-    );
+    // The body is a string the cache may be handing to many requests at once,
+    // so it goes to Koa as-is. Wrapping it in a `Response` would encode it into
+    // a stream that `setContextResponse` decodes back into a per-request copy;
+    // this way a joiner's or a hit's cost on the way out is Koa's own write of
+    // the shared string and nothing more.
+    ctxt.status = 200;
+    ctxt.set('content-type', SupportedMimeType.CardJson);
+    ctxt.set(LIVE_SEARCH_CACHE_HEADER, outcome);
+    ctxt.body = body;
     emitTimeline();
     return;
   }

packages/realm-server/live-search-cache.ts

@@ -175,11 +175,18 @@ export class LiveSearchCache {
     this.#emitTelemetry = opts?.emitTelemetry ?? defaultEmitTelemetry;
   }
 
+  // `onOutcome` fires the moment the cache decides how it will satisfy the
+  // request — before a joiner starts waiting on the in-flight compute, before
+  // a hit returns, and as a miss starts its populate — so a caller can act on
+  // the decision while the request is still in progress. The realm-server's
+  // search admission uses it to hand back the slot of a request that holds no
+  // result document of its own.
   async getOrPopulate(args: {
     realms: string[];
     query: Query;
     opts: unknown | undefined;
     populate: () => Promise<string>;
+    onOutcome?: (outcome: LiveSearchOutcome) => void;
   }): Promise<{ body: string; outcome: LiveSearchOutcome }> {
     try {
       return await this.#getOrPopulate(args);
@@ -193,9 +200,11 @@ export class LiveSearchCache {
     query: Query;
     opts: unknown | undefined;
     populate: () => Promise<string>;
+    onOutcome?: (outcome: LiveSearchOutcome) => void;
   }): Promise<{ body: string; outcome: LiveSearchOutcome }> {
     this.#reapExpiredHead();
     let key = searchRequestKeyHash(args.realms, args.query, args.opts);
+    let onOutcome = args.onOutcome ?? (() => {});
 
     let entry = this.#entries.get(key);
     if (entry) {
@@ -206,6 +215,7 @@ export class LiveSearchCache {
         this.#entries.set(key, entry);
         this.#counters.hits += 1;
         this.#counters.hitBytes += entry.body.length;
+        onOutcome('hit');
         return { body: entry.body, outcome: 'hit' };
       }
       this.#delete(key, entry);
@@ -215,11 +225,13 @@ export class LiveSearchCache {
     let inFlight = this.#inFlight.get(key);
     if (inFlight) {
       this.#counters.joins += 1;
+      onOutcome('join');
       let body = await inFlight;
       this.#counters.joinBytes += body.length;
       return { body, outcome: 'join' };
     }
 
+    onOutcome('miss');
     let promise = args.populate();
     this.#inFlight.set(key, promise);
     try {

packages/realm-server/middleware/index.ts

@@ -222,12 +222,34 @@ export function httpLogging(ctxt: Koa.Context, next: Koa.Next) {
   return next();
 }
 
-// Puts a search through the admission gate (`search-inflight.ts`) and holds
-// its slot for the request's full lifecycle (parse → SQL → serialize → send),
-// which is the window in which it holds heap and in which a saturated event
-// loop would leave it unserviced. Mounted after CORS so that a shed response
-// carries the headers a cross-origin client needs to read its status and
-// Retry-After; before the body is parsed so that a shed costs nothing.
+// Where `searchAdmission` leaves the release for the slot it granted, so a
+// handler can hand the slot back before the response ends.
+const SEARCH_ADMISSION_RELEASE = 'searchAdmissionRelease';
+
+// Hand back the admission slot a search request holds, if it holds one. A
+// request that the live-search cache satisfies from another request's
+// computation — a `join` or a `hit` — builds no result document of its own, so
+// it calls this the moment the cache says so and stops counting toward the
+// ceiling; only the request doing the computing keeps its slot until its
+// response ends. That holds for an indexing-lane admission too: an in-render
+// search served from another request's computation holds no document either,
+// and the count is of computations, whichever lane admitted them. Idempotent,
+// and a no-op for requests the gate never saw.
+export function releaseSearchAdmission(ctxt: Koa.Context): void {
+  let release = ctxt.state[SEARCH_ADMISSION_RELEASE];
+  if (typeof release === 'function') {
+    release();
+  }
+}
+
+// Puts a search through the admission gate (`search-inflight.ts`). A search
+// that computes its own result holds its slot for the request's full lifecycle
+// (parse → SQL → serialize → send), which is the window in which it holds heap
+// and in which a saturated event loop would leave it unserviced; one that the
+// live-search cache serves from another's computation hands the slot back
+// early via `releaseSearchAdmission`. Mounted after CORS so that a shed
+// response carries the headers a cross-origin client needs to read its status
+// and Retry-After; before the body is parsed so that a shed costs nothing.
 export async function searchAdmission(ctxt: Koa.Context, next: Koa.Next) {
   if (
     !SEARCH_PATH_PATTERN.test(ctxt.path) ||
@@ -249,6 +271,7 @@ export async function searchAdmission(ctxt: Koa.Context, next: Koa.Next) {
     release?.();
     release = undefined;
   };
+  ctxt.state[SEARCH_ADMISSION_RELEASE] = releaseSlot;
   // `finish` fires on a fully-sent response; `close` covers a connection
   // torn down before that, so a slot can't leak on an abort.
   ctxt.res.on('finish', releaseSlot);

packages/realm-server/search-inflight.ts

@@ -4,14 +4,18 @@ import {
 } from '@cardstack/runtime-common';
 
 // Admission gate for the realm-server's search endpoints (`/_search`,
-// `/_federated-search`). Every in-flight search holds tens of MB of heap while
-// its result set is assembled, and the process is a single event loop, so the
-// number of searches running at once is what decides whether the heap survives
-// a burst. The gate bounds it: up to `limit` searches run concurrently,
-// arrivals above that wait up to a bounded time for a slot in FIFO order, and a
-// request still waiting when its time is up is shed — the middleware answers
-// 429 + Retry-After without having parsed a body or touched the index, so a
-// shed costs the process almost nothing.
+// `/_federated-search`). A search that assembles its own result document holds
+// tens of MB of heap while it does, and the process is a single event loop, so
+// the number of such computations running at once is what decides whether the
+// heap survives a burst. The gate bounds it: up to `limit` searches hold a
+// slot at once, arrivals above that wait up to a bounded time for a slot in
+// FIFO order, and a request still waiting when its time is up is shed — the
+// middleware answers 429 + Retry-After without having parsed a body or
+// touched the index, so a shed costs the process almost nothing. A request
+// that the live-search cache serves from another request's computation hands
+// its slot back as soon as the cache decides so, so in steady state the slots
+// are held by computations plus the requests briefly between admission and
+// the cache lookup.
 //
 // Indexing traffic (a request stamped with a prerender job id or the
 // during-prerender header) is admitted unconditionally: shedding an in-render

packages/realm-server/tests/live-search-cache-test.ts

@@ -67,6 +67,42 @@ module(basename(import.meta.filename), function () {
     assert.strictEqual(a.body, b.body, 'both callers share the body');
   });
 
+  test('onOutcome fires when the cache decides, not when the body resolves', async function (assert) {
+    let cache = new LiveSearchCache({ ttlMs: 60_000 });
+    let deferred = deferredPopulate('{"data":[1]}');
+    let realms = ['http://a/'];
+    let decided: string[] = [];
+    let load = () =>
+      cache.getOrPopulate({
+        realms,
+        query: personQuery(),
+        opts: undefined,
+        populate: deferred.populate,
+        onOutcome: (outcome) => decided.push(outcome),
+      });
+
+    let first = load();
+    assert.deepEqual(decided, ['miss'], 'the miss is announced as it starts');
+    let second = load();
+    assert.deepEqual(
+      decided,
+      ['miss', 'join'],
+      'the join is announced before the joiner has anything to wait on',
+    );
+    assert.strictEqual(deferred.calls, 1, 'one populate');
+
+    deferred.resolve();
+    await Promise.all([first, second]);
+
+    let third = load();
+    assert.deepEqual(
+      decided,
+      ['miss', 'join', 'hit'],
+      'a hit is announced synchronously on the call',
+    );
+    assert.strictEqual((await third).outcome, 'hit');
+  });
+
   test('different queries do not coalesce', async function (assert) {
     let cache = new LiveSearchCache({ ttlMs: 60_000 });
     let realms = ['http://a/'];

packages/realm-server/tests/search-admission-test.ts

@@ -12,7 +12,11 @@ import {
   resetSearchAdmissionForTests,
   setSearchAdmissionForTests,
 } from '../search-inflight.ts';
-import { httpLogging, searchAdmission } from '../middleware/index.ts';
+import {
+  httpLogging,
+  releaseSearchAdmission,
+  searchAdmission,
+} from '../middleware/index.ts';
 
 // The admission gate is what stands between a burst of searches and a heap
 // exhausted by their concurrent result sets. These tests pin the contract the
@@ -177,6 +181,9 @@ module(basename(import.meta.filename), function () {
       let app = new Koa();
       let router = new Router();
       let search = async (ctxt: Koa.Context) => {
+        if (ctxt.query.releaseEarly) {
+          releaseSearchAdmission(ctxt);
+        }
         if (ctxt.query.hold) {
           await new Promise<void>((resolve) => {
             holds.push(resolve);
@@ -225,7 +232,8 @@ module(basename(import.meta.filename), function () {
       headers: Record<string, string> = {},
     ) {
       let held = new Promise<void>((resolve) => (onHeld = resolve));
-      let response = send(app, `${path}?hold=1`, headers);
+      let separator = path.includes('?') ? '&' : '?';
+      let response = send(app, `${path}${separator}hold=1`, headers);
       return { response, held };
     }
 
@@ -252,12 +260,11 @@ module(basename(import.meta.filename), function () {
       return [first.response, second.response];
     }
 
-    function assert_inFlight(expected: number) {
-      QUnit.assert.strictEqual(
-        getSearchInFlight(),
-        expected,
-        `inFlight=${expected}`,
-      );
+    function assert_inFlight(
+      expected: number,
+      message = `inFlight=${expected}`,
+    ) {
+      QUnit.assert.strictEqual(getSearchInFlight(), expected, message);
     }
 
     test('a search arriving above the ceiling is shed with 429 and Retry-After', async function (assert) {
@@ -353,6 +360,41 @@ module(basename(import.meta.filename), function () {
       await Promise.all(held);
     });
 
+    test('a handler can hand its slot back before its response ends', async function (assert) {
+      setSearchAdmissionForTests({ limit: 2, waitMs: 2000 });
+      let app = buildApp();
+      let releasing = holdSearch(app, '/_federated-search?releaseEarly=1');
+      await releasing.held;
+      assert_inFlight(0, 'released while the response is still open');
+
+      // The freed slot is real: with the gate full again, a waiter is blocked
+      // by the two computing searches, not by the early-released one.
+      let held = await fillGate(app);
+      let waiting = send(app, '/_federated-search');
+      await wait(50);
+      let blocked = await settledWithin(waiting, 20);
+      assert.false(blocked.settled, 'the gate is full');
+
+      // Ending the early-released response hands back nothing more.
+      holds.shift()!();
+      await releasing.response;
+      let stillBlocked = await settledWithin(waiting, 20);
+      assert.false(stillBlocked.settled, 'no second release on finish');
+      assert_inFlight(2);
+
+      holds.shift()!();
+      assert.strictEqual(
+        (await waiting).status,
+        200,
+        'a computing search ending admits the waiter',
+      );
+      for (let release of holds) {
+        release();
+      }
+      await Promise.all(held);
+      assert_inFlight(0);
+    });
+
     test('a non-search request is not counted', async function (assert) {
       let app = buildApp();
       let { response, held } = holdSearch(app, '/some-realm/cards');