perf(zql): pull-function protocol for the fetch path

packages/shared/src/btree-set.ts

@@ -96,7 +96,7 @@ export class BTreeSet<K> {
     return valuesFrom(this.#root, this.comparator, undefined, true);
   }
 
-  valuesFrom(lowestKey?: K, inclusive: boolean = true): IterableIterator<K> {
+  valuesFrom(lowestKey?: K, inclusive: boolean = true): ValueIterator<K> {
     return valuesFrom(this.#root, this.comparator, lowestKey, inclusive);
   }
 
@@ -113,7 +113,7 @@ export class BTreeSet<K> {
   valuesFromReversed(
     highestKey?: K,
     inclusive: boolean = true,
-  ): IterableIterator<K> {
+  ): ValueIterator<K> {
     return valuesFromReversed(
       this.#maxKey(),
       this.#root,
@@ -187,7 +187,16 @@ export class BTreeSet<K> {
   }
 }
 
-class BTreeForwardIterator<K> implements IterableIterator<K> {
+/**
+ * An iterator that can also hand back values without the iterator protocol's
+ * per-value result object. `nextValue()` returns `undefined` once exhausted,
+ * so it is only meaningful for sets whose keys are never `undefined`.
+ */
+export interface ValueIterator<K> extends IterableIterator<K> {
+  nextValue(): K | undefined;
+}
+
+class BTreeForwardIterator<K> implements ValueIterator<K> {
   readonly #nodeQueue: BNode<K>[][];
   readonly #nodeIndex: number[];
   #leaf: BNode<K>;
@@ -205,16 +214,17 @@ class BTreeForwardIterator<K> implements IterableIterator<K> {
     this.#i = startI;
   }
 
-  next(): IteratorResult<K> {
+  /** Moves to the next key; false once exhausted. */
+  #advance(): boolean {
     for (;;) {
       if (++this.#i < this.#leaf.keys.length) {
-        return {done: false, value: this.#leaf.keys[this.#i]};
+        return true;
       }
 
       let level = -1;
       for (;;) {
         if (++level >= this.#nodeQueue.length) {
-          return {done: true, value: undefined as unknown as K};
+          return false;
         }
         if (++this.#nodeIndex[level] < this.#nodeQueue[level].length) {
           break;
@@ -231,12 +241,22 @@ class BTreeForwardIterator<K> implements IterableIterator<K> {
     }
   }
 
+  next(): IteratorResult<K> {
+    return this.#advance()
+      ? {done: false, value: this.#leaf.keys[this.#i]}
+      : {done: true, value: undefined as unknown as K};
+  }
+
+  nextValue(): K | undefined {
+    return this.#advance() ? this.#leaf.keys[this.#i] : undefined;
+  }
+
   [Symbol.iterator]() {
     return this;
   }
 }
 
-class BTreeReverseIterator<K> implements IterableIterator<K> {
+class BTreeReverseIterator<K> implements ValueIterator<K> {
   readonly #nodeQueue: BNode<K>[][];
   readonly #nodeIndex: number[];
   #leaf: BNode<K>;
@@ -254,17 +274,18 @@ class BTreeReverseIterator<K> implements IterableIterator<K> {
     this.#i = startI;
   }
 
-  next(): IteratorResult<K> {
+  /** Moves to the previous key; false once exhausted. */
+  #advance(): boolean {
     for (;;) {
       if (--this.#i >= 0) {
-        return {done: false, value: this.#leaf.keys[this.#i]};
+        return true;
       }
 
       let level;
       // Advance to the next leaf node
       for (level = -1; ;) {
         if (++level >= this.#nodeQueue.length) {
-          return {done: true, value: undefined as unknown as K};
+          return false;
         }
         if (--this.#nodeIndex[level] >= 0) {
           break;
@@ -281,6 +302,16 @@ class BTreeReverseIterator<K> implements IterableIterator<K> {
     }
   }
 
+  next(): IteratorResult<K> {
+    return this.#advance()
+      ? {done: false, value: this.#leaf.keys[this.#i]}
+      : {done: true, value: undefined as unknown as K};
+  }
+
+  nextValue(): K | undefined {
+    return this.#advance() ? this.#leaf.keys[this.#i] : undefined;
+  }
+
   [Symbol.iterator]() {
     return this;
   }
@@ -291,10 +322,10 @@ function valuesFrom<K>(
   comparator: Comparator<K>,
   lowestKey: K | undefined,
   inclusive: boolean,
-): IterableIterator<K> {
+): ValueIterator<K> {
   const info = findPath(lowestKey, root, comparator);
   if (info === undefined) {
-    return iterator<K>(() => ({done: true, value: undefined}));
+    return emptyValueIterator<K>();
   }
 
   let [nodeQueue, nodeIndex, leaf] = info;
@@ -322,11 +353,11 @@ function valuesFromReversed<K>(
   comparator: Comparator<K>,
   highestKey: K | undefined,
   inclusive: boolean,
-): IterableIterator<K> {
+): ValueIterator<K> {
   if (highestKey === undefined) {
     highestKey = maxKey;
     if (highestKey === undefined) {
-      return iterator<K>(() => ({done: true, value: undefined}));
+      return emptyValueIterator<K>();
     } // collection is empty
   }
   let [nodeQueue, nodeIndex, leaf] =
@@ -371,9 +402,10 @@ function findPath<K>(
   return [nodeQueue, nodeIndex, nextNode];
 }
 
-function iterator<T>(next: () => IteratorResult<T>): IterableIterator<T> {
+function emptyValueIterator<K>(): ValueIterator<K> {
   return {
-    next,
+    next: () => ({done: true, value: undefined as unknown as K}),
+    nextValue: () => undefined,
     [Symbol.iterator]() {
       return this;
     },

packages/zero-cache/src/auth/write-authorizer.ts

@@ -569,10 +569,13 @@ export class WriteAuthorizerImpl implements WriteAuthorizer {
     const input = buildPipeline(rowQueryAst, this.#builderDelegate, 'query-id');
     try {
       const res = input.fetch({});
-      for (const _ of res) {
-        // if any row is returned at all, the
-        // rule passes.
-        return true;
+      try {
+        // if any row is returned at all, the rule passes.
+        if (res.next() !== undefined) {
+          return true;
+        }
+      } finally {
+        res.close();
       }
     } finally {
       input.destroy();

packages/zero-cache/src/services/run-ast.ts

@@ -102,8 +102,13 @@ export async function runAst(
     // triggering early return on Take's #initialFetch assertion.
     // The subquery AST already has limit: 1, so at most one row is produced.
     let node: Node | undefined;
-    for (const n of skipYields(input.fetch({}))) {
-      node ??= n;
+    const rows = skipYields(input.fetch({}));
+    try {
+      for (let n = rows.next(); n !== undefined; n = rows.next()) {
+        node ??= n;
+      }
+    } finally {
+      rows.close();
     }
     input.destroy();
     return node ? ((node.row[childField] as LiteralValue) ?? null) : undefined;

packages/zero-cache/src/services/view-syncer/flipped-exists-fetch-filter.bench.ts

@@ -21,7 +21,8 @@ type Fetch = {
 
 function fetchRowCount(input: Input, filter: NoSubqueryCondition | undefined) {
   let count = 0;
-  for (const node of input.fetch({filter})) {
+  const stream = input.fetch({filter});
+  for (let node = stream.next(); node !== undefined; node = stream.next()) {
     if (node !== 'yield') {
       count++;
     }

packages/zero-cache/src/services/view-syncer/pipeline-driver.ts

@@ -15,7 +15,7 @@ import {
 import {ChangeIndex} from '../../../../zql/src/ivm/change-index.ts';
 import {ChangeType} from '../../../../zql/src/ivm/change-type.ts';
 import type {Change} from '../../../../zql/src/ivm/change.ts';
-import type {Node} from '../../../../zql/src/ivm/data.ts';
+import type {Node, RelationshipStream} from '../../../../zql/src/ivm/data.ts';
 import {
   skipYields,
   throwOutput,
@@ -33,6 +33,7 @@ import {
   makeSourceChangeEdit,
   makeSourceChangeRemove,
 } from '../../../../zql/src/ivm/source.ts';
+import {pullOf, type PullStream} from '../../../../zql/src/ivm/stream.ts';
 import type {ConnectionCostModel} from '../../../../zql/src/planner/planner-connection.ts';
 import {MeasurePushOperator} from '../../../../zql/src/query/measure-push-operator.ts';
 import type {ClientGroupStorage} from '../../../../zqlite/src/database-storage.ts';
@@ -560,8 +561,13 @@ export class PipelineDriver {
       // triggering early return on Take's #initialFetch assertion.
       // The subquery AST already has limit: 1, so at most one row is produced.
       let node: Node | undefined;
-      for (const n of skipYields(input.fetch({}))) {
-        node ??= n;
+      const rows = skipYields(input.fetch({}));
+      try {
+        for (let n = rows.next(); n !== undefined; n = rows.next()) {
+          node ??= n;
+        }
+      } finally {
+        rows.close();
       }
       if (!node) {
         return undefined;
@@ -1369,9 +1375,9 @@ class Streamer {
       switch (type) {
         case ChangeType.REMOVE:
         case ChangeType.ADD: {
-          yield* this.#streamNodes(queryID, schema, type, () => [
-            change[ChangeIndex.NODE],
-          ]);
+          yield* this.#streamNodes(queryID, schema, type, () =>
+            pullOf([change[ChangeIndex.NODE]]),
+          );
           break;
         }
 
@@ -1385,9 +1391,9 @@ class Streamer {
           break;
         }
         case ChangeType.EDIT:
-          yield* this.#streamNodes(queryID, schema, type, () => [
-            {row: change[ChangeIndex.NODE].row, relationships: {}},
-          ]);
+          yield* this.#streamNodes(queryID, schema, type, () =>
+            pullOf([{row: change[ChangeIndex.NODE].row, relationships: {}}]),
+          );
           break;
         default:
           unreachable(change[ChangeIndex.TYPE]);
@@ -1399,7 +1405,7 @@ class Streamer {
     queryID: string,
     schema: SourceSchema,
     op: ChangeType.ADD | ChangeType.REMOVE | ChangeType.EDIT,
-    nodes: () => Iterable<Node | 'yield'>,
+    nodes: () => RelationshipStream,
   ): Iterable<RowChange | 'yield'> {
     const {tableName: table, system} = schema;
 
@@ -1412,36 +1418,46 @@ class Streamer {
       return;
     }
 
-    for (const node of nodes()) {
-      if (node === 'yield') {
-        yield node;
-        continue;
-      }
-      const {relationships} = node;
-      let {row} = node;
-      const rowKey = getRowKey(primaryKey, row);
-      if (op !== ChangeType.REMOVE) {
-        const rowVersion = row[ZERO_VERSION_COLUMN_NAME];
-        if (
-          typeof rowVersion === 'string' &&
-          rowVersion < (spec.minRowVersion ?? '00')
-        ) {
-          row = {...row, [ZERO_VERSION_COLUMN_NAME]: spec.minRowVersion};
+    const stream = nodes();
+    try {
+      for (let node = stream.next(); node !== undefined; node = stream.next()) {
+        if (node === 'yield') {
+          yield node;
+          continue;
+        }
+        const {relationships} = node;
+        let {row} = node;
+        const rowKey = getRowKey(primaryKey, row);
+        if (op !== ChangeType.REMOVE) {
+          const rowVersion = row[ZERO_VERSION_COLUMN_NAME];
+          if (
+            typeof rowVersion === 'string' &&
+            rowVersion < (spec.minRowVersion ?? '00')
+          ) {
+            row = {...row, [ZERO_VERSION_COLUMN_NAME]: spec.minRowVersion};
+          }
         }
-      }
 
-      yield {
-        type: op,
-        queryID,
-        table,
-        rowKey,
-        row: op === ChangeType.REMOVE ? undefined : row,
-      } as RowChange;
-
-      for (const [relationship, children] of Object.entries(relationships)) {
-        const childSchema = must(schema.relationships[relationship]);
-        yield* this.#streamNodes(queryID, childSchema, op, children);
+        yield {
+          type: op,
+          queryID,
+          table,
+          rowKey,
+          row: op === ChangeType.REMOVE ? undefined : row,
+        } as RowChange;
+
+        for (const [relationship, children] of Object.entries(relationships)) {
+          const childSchema = must(schema.relationships[relationship]);
+          yield* this.#streamNodes(
+            queryID,
+            childSchema,
+            op,
+            children as () => RelationshipStream,
+          );
+        }
       }
+    } finally {
+      stream.close();
     }
   }
 }
@@ -1481,7 +1497,7 @@ class QueryFailureLoggingOperator implements Input, Output {
     this.#input.destroy();
   }
 
-  fetch(req: FetchRequest): Iterable<Node | 'yield'> {
+  fetch(req: FetchRequest): PullStream<Node | 'yield'> {
     return this.#input.fetch(req);
   }
 
@@ -1525,13 +1541,24 @@ function logQueryFailure(
   queryLC.error?.(message, error);
 }
 
-function* toAdds(nodes: Iterable<Node | 'yield'>): Iterable<Change | 'yield'> {
-  for (const node of nodes) {
-    if (node === 'yield') {
-      yield node;
-      continue;
+function* toAdds(
+  nodes: PullStream<Node | 'yield'>,
+): Iterable<Change | 'yield'> {
+  // `finally` is load-bearing: abandoning this generator -- an aborted or
+  // evicted hydration -- must still close the stream. `for...of` did that
+  // implicitly via `.return()`; the pull protocol makes it explicit, and a
+  // leaked SQLite cursor leaves later writes failing with "database
+  // connection is busy executing a query".
+  try {
+    for (let node = nodes.next(); node !== undefined; node = nodes.next()) {
+      if (node === 'yield') {
+        yield node;
+        continue;
+      }
+      yield [ChangeType.ADD, node, null];
     }
-    yield [ChangeType.ADD, node, null];
+  } finally {
+    nodes.close();
   }
 }