Avoid weakset for executed generators (#64830)

Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: dragomirtitian

packages/typescript/src/api/sync/generatorSupport.ts

@@ -23,6 +23,26 @@ export function* apiRequest(method: PropertyKey, params: unknown): Generator<{ m
     return yield { method, params };
 }
 const deferredGeneratorMarker: unique symbol = Symbol();
+// Marks executed generators on the generator itself: a WeakSet holding every generator of a long-running batch is
+// costly to grow and to process during GC
+const executedMarker: unique symbol = Symbol();
+interface ExecutedGenerator {
+    [executedMarker]?: true;
+}
+const executedNonExtensibleGenerators = new WeakSet<AnyAPIRequestGenerator>();
+
+/** Marks the generator as executed, returning false if it already was. */
+function markExecuted(generator: AnyAPIRequestGenerator): boolean {
+    if (!Object.isExtensible(generator)) {
+        if (executedNonExtensibleGenerators.has(generator)) return false;
+        executedNonExtensibleGenerators.add(generator);
+        return true;
+    }
+    const target = generator as ExecutedGenerator;
+    if (target[executedMarker]) return false;
+    target[executedMarker] = true;
+    return true;
+}
 const allGeneratorMarker: unique symbol = Symbol();
 interface DeferredAPIRequest {
     readonly method: "__defer";
@@ -125,7 +145,6 @@ export function executeRequestGenerators<T extends readonly AnyAPIRequestGenerat
     requestGenerators: T,
     executeRequests: (requests: APIRequest[]) => readonly GeneratorResponse[],
 ): ExecutedGeneratorsResults<T> {
-    const registeredGenerators = new WeakSet<AnyAPIRequestGenerator>();
     let firstGenerator: GeneratorState | undefined;
     let lastGenerator: GeneratorState | undefined;
     const root = startGroup(requestGenerators);
@@ -220,13 +239,12 @@ export function executeRequestGenerators<T extends readonly AnyAPIRequestGenerat
     }
 
     function addGenerator(generator: AnyAPIRequestGenerator, parent?: GeneratorGroup): void {
-        if (registeredGenerators.has(generator)) {
+        if (!markExecuted(generator)) {
             const error = new Error("Cannot execute the same generator instance more than once");
             if (!parent) throw error;
             failGroup(parent, error);
             return;
         }
-        registeredGenerators.add(generator);
         const next = parent?.parent;
         const previous = next ? next.previous : lastGenerator;
         const state: GeneratorState = {
@@ -265,12 +283,11 @@ export function executeRequestGenerators<T extends readonly AnyAPIRequestGenerat
             const wrapper = generator as AllRequestGenerator<unknown>;
             const pending = wrapper.generators;
             if (!pending) break;
-            if (registeredGenerators.has(wrapper)) {
+            if (!markExecuted(wrapper)) {
                 failGroup(group, new Error("Cannot execute the same generator instance more than once"));
                 group.initializing = false;
                 return group;
             }
-            registeredGenerators.add(wrapper);
             wrapper.next();
             (group.wrappers ??= []).push(wrapper);
             generators = pending;

packages/typescript/test/api-comparison.bench.ts

@@ -56,6 +56,7 @@ export async function runBenchmarks(options?: { filter?: string; singleIteration
     addGeneratorBenchmarks();
     addExecutorBenchmarks();
     addDeferBenchmarks();
+    addLongBatchBenchmarks();
 
     if (filter) {
         const pattern = filter.toLowerCase();
@@ -291,6 +292,62 @@ export async function runBenchmarks(options?: { filter?: string; singleIteration
         }
     }
 
+    function addLongBatchBenchmarks() {
+        // Enough short-lived generators in a single executor run that per-generator bookkeeping dominates
+        const generatorCount = 100_000;
+        const rounds = 100;
+        const generatorsPerRound = generatorCount / rounds;
+        for (const mode of ["direct", "yield* defer"] as const) {
+            bench.add(`executor long batch - ${generatorCount} generators (${mode})`, () => run(false), {
+                async: false,
+                beforeAll: () => {
+                    run(true);
+                },
+            });
+
+            function run(verify: boolean) {
+                let completed = 0;
+                let executedRequests = 0;
+                let executedRounds = 0;
+                const results = executeRequestGenerators(
+                    mode === "direct" ? Array.from({ length: generatorCount }, (_, index) => runRequest(index)) : [schedule()],
+                    requests => {
+                        executedRequests += requests.length;
+                        executedRounds++;
+                        return requests.map(request => ({ result: request.params }));
+                    },
+                );
+                if (verify) {
+                    assert.equal(completed, generatorCount);
+                    if (mode === "direct") {
+                        assert.equal(executedRequests, generatorCount);
+                        assert.equal(executedRounds, 1);
+                        assert.equal(results.length, generatorCount);
+                    }
+                    else {
+                        assert.equal(executedRequests, generatorCount + rounds);
+                        assert.equal(executedRounds, rounds);
+                        assert.deepEqual(results, [undefined]);
+                    }
+                }
+
+                function* runRequest(index: number): APIRequestGenerator<number> {
+                    const result = yield { method: "benchmark", params: index } as unknown as APIRequest;
+                    completed++;
+                    return result;
+                }
+
+                // A long-running scheduler that keeps spawning new generators across many request rounds
+                function* schedule(): APIRequestGenerator {
+                    for (let round = 0; round < rounds; round++) {
+                        for (let lane = 0; lane < generatorsPerRound; lane++) yield* defer(runRequest(round * generatorsPerRound + lane));
+                        yield { method: "benchmark", params: round } as unknown as APIRequest;
+                    }
+                }
+            }
+        }
+    }
+
     function* runDependentRequests(project: SyncProject) {
         let index = 0;
         for (let request = 0; request < requestCount; request++) {

packages/typescript/test/sync/api-generators.test.ts

@@ -579,7 +579,7 @@ describe("API - generator batching", { concurrency: areTestsFiltered() }, () =>
         assert.equal(started, false);
     });
 
-    test("retains deferred branding after completion", () => {
+    test("rejects deferred generator reuse after completion", () => {
         let completed = false;
         function* child(): Generator<APIRequest, string, string> {
             const result = yield { method: "request", params: null } as unknown as APIRequest;
@@ -591,8 +591,7 @@ describe("API - generator batching", { concurrency: areTestsFiltered() }, () =>
         assert.deepEqual(results, []);
         assert.equal(completed, true);
         const noRequests = () => assert.fail("Unexpected request round");
-        assert.deepEqual(executeRequestGenerators([deferred], noRequests), []);
-        assert.throws(() => executeRequestGenerators([deferred, deferred], noRequests), /same generator instance more than once/);
+        assert.throws(() => executeRequestGenerators([deferred], noRequests), /same generator instance more than once/);
     });
 
     test("joins nested, empty, and synchronously completed groups in result order", () => {
@@ -897,23 +896,15 @@ describe("API - generator batching", { concurrency: areTestsFiltered() }, () =>
         }
     });
 
-    test("tracks all identity independently across executors", () => {
+    test("rejects generator reuse across executors", () => {
         const joined = all();
         const noRequests = () => assert.fail("Unexpected request round");
         function* parent() {
             assert.deepEqual(yield* all(joined), [[]]);
-            assert.deepEqual(executeRequestGenerators([joined], noRequests), [undefined]);
-            try {
-                yield* all(joined);
-                assert.fail("Expected duplicate generator rejection");
-            }
-            catch (error) {
-                assert.match((error as Error).message, /same generator instance more than once/);
-            }
+            assert.throws(() => executeRequestGenerators([joined], noRequests), /same generator instance more than once/);
         }
         executeRequestGenerators([parent()], noRequests);
-        assert.deepEqual(executeRequestGenerators([joined], noRequests), [undefined]);
-        assert.throws(() => executeRequestGenerators([joined, joined], noRequests), /same generator instance more than once/);
+        assert.throws(() => executeRequestGenerators([joined], noRequests), /same generator instance more than once/);
     });
 
     test("deduplicates initialization across nested joins", () => {