@peerbit/document 13.1.5 → 13.1.7

package/src/search.ts

@@ -24,7 +24,7 @@ import {
 import { CachedIndex, type QueryCacheOptions } from "@peerbit/indexer-cache";
 import * as indexerTypes from "@peerbit/indexer-interface";
 import { HashmapIndex } from "@peerbit/indexer-simple";
-import { BORSH_ENCODING, type Encoding, Entry } from "@peerbit/log";
+import { BORSH_ENCODING, type Encoding, type Entry } from "@peerbit/log";
 import { logger as loggerFn } from "@peerbit/logger";
 import { ClosedError, Program } from "@peerbit/program";
 import {
@@ -47,7 +47,7 @@ import {
 } from "@peerbit/stream-interface";
 import { AbortError, TimeoutError, waitFor } from "@peerbit/time";
 import pDefer, { type DeferredPromise } from "p-defer";
-import { concat, fromString } from "uint8arrays";
+import { concat, equals, fromString } from "uint8arrays";
 import { copySerialization } from "./borsh.js";
 import { MAX_BATCH_SIZE } from "./constants.js";
 import type { DocumentEvents, DocumentsChange } from "./events.js";
@@ -511,6 +511,7 @@ type QueryDetailedOptions<
 		from: PublicSignKey,
 	) => void | Promise<void>;
 	onMissingResponses?: (error: MissingResponsesError) => void | Promise<void>;
+	onRemoteTargets?: (targets: string[]) => void;
 	remote?: {
 		from?: string[]; // if specified, only query these peers
 	};
@@ -742,6 +743,7 @@ function isSubclassOf(
 const DEFAULT_TIMEOUT = 1e4;
 const DEFAULT_KEEP_REMOTE_ITERATOR_TIMEOUT = 3e5;
 const DISCOVER_TIMEOUT_FALLBACK = 500;
+const CLOSE_ITERATOR_REQUEST_TIMEOUT = 5e3;
 
 const DEFAULT_INDEX_BY = "id";
 
@@ -1179,26 +1181,56 @@ export class DocumentIndex<
 		}
 	}
 
-	/** Head entry safe to borsh-serialize into ResultIndexedValue.entries over the
-	 * document RPC. Under the native block store, this._log.log.get can return a
-	 * hollow entry whose payload bytes were never materialized on the JS side;
-	 * serializing it throws and aborts the whole RPC response. The complete entry is
-	 * in the block store keyed by hash, so recover it there. No-op for pure JS. */
-	private async getSerializableHead(
-		hash: string,
-	): Promise<Entry<any> | undefined> {
-		const head = await this._log.log.get(hash);
-		if (!head?.hash) return head ?? undefined;
-		try {
-			head.getStorageBytes(); // = serialize(head); throws on a hollow native entry
-			return head;
-		} catch {
-			try {
-				return await Entry.fromMultihash(this._log.log.blocks, head.hash);
-			} catch {
-				return head; // block absent (e.g. pruned) — preserve today's behavior
+	/**
+	 * Resolve documents without bypassing the resolver/program caches. Cache and
+	 * identity-index hits are collected first; only the remaining heads cross the
+	 * log read boundary, and those reads stay aligned in one batch.
+	 */
+	private async resolveDocumentsWithBatchedHeads(
+		values: Array<{
+			id?: indexerTypes.IdPrimitive;
+			indexed: I;
+			head: string;
+		}>,
+	): Promise<{
+		resolved: Array<{ value: T } | undefined>;
+		heads: Array<Entry<Operation> | undefined>;
+	}> {
+		const resolved: Array<{ value: T } | undefined> = new Array(values.length);
+		const heads: Array<Entry<Operation> | undefined> = new Array(values.length);
+		const unresolvedPositions: number[] = [];
+
+		for (let i = 0; i < values.length; i++) {
+			const value = values[i]!;
+			const cached = await this.resolveDocument({
+				...value,
+				headEntry: null,
+			});
+			if (cached) {
+				resolved[i] = cached;
+			} else {
+				unresolvedPositions.push(i);
 			}
 		}
+
+		if (unresolvedPositions.length === 0) {
+			return { resolved, heads };
+		}
+
+		const unresolvedHeads = await this._log.log.getMany(
+			unresolvedPositions.map((position) => values[position]!.head),
+		);
+		for (let i = 0; i < unresolvedPositions.length; i++) {
+			const position = unresolvedPositions[i]!;
+			const head = unresolvedHeads[i];
+			heads[position] = head;
+			resolved[position] = await this.resolveDocument({
+				...values[position]!,
+				headEntry: head ?? null,
+			});
+		}
+
+		return { resolved, heads };
 	}
 
 	private async wrapPushResults(
@@ -1206,8 +1238,41 @@ export class DocumentIndex<
 		resolve: boolean,
 	): Promise<types.Result[]> {
 		if (!matches.length) return [];
+		const headsByMatch: Array<Entry<Operation> | undefined> = [];
+		const resolvedByMatch: Array<{ value: T } | undefined> = [];
+		if (!resolve) {
+			headsByMatch.push(
+				...(await this._log.log.getMany(
+					matches.map((match) => match.__context.head),
+				)),
+			);
+		} else {
+			const indexedPositions: number[] = [];
+			for (let i = 0; i < matches.length; i++) {
+				if (!(matches[i] instanceof this.documentType)) {
+					indexedPositions.push(i);
+				}
+			}
+			const indexedBatch = indexedPositions.length
+				? await this.resolveDocumentsWithBatchedHeads(
+						indexedPositions.map((position) => {
+							const indexed = matches[position] as WithContext<I>;
+							return {
+								indexed,
+								head: indexed.__context.head,
+							};
+						}),
+					)
+				: { resolved: [], heads: [] };
+			for (let i = 0; i < indexedPositions.length; i++) {
+				const position = indexedPositions[i]!;
+				resolvedByMatch[position] = indexedBatch.resolved[i];
+				headsByMatch[position] = indexedBatch.heads[i];
+			}
+		}
 		const results: types.Result[] = [];
-		for (const match of matches) {
+		for (let i = 0; i < matches.length; i++) {
+			const match = matches[i]!;
 			if (resolve) {
 				if (match instanceof this.documentType) {
 					const doc = match as WithContext<T>;
@@ -1225,10 +1290,7 @@ export class DocumentIndex<
 				}
 
 				const indexed = match as WithContext<I>;
-				const resolved = await this.resolveDocument({
-					indexed,
-					head: indexed.__context.head,
-				});
+				const resolved = resolvedByMatch[i];
 
 				if (resolved) {
 					results.push(
@@ -1242,7 +1304,7 @@ export class DocumentIndex<
 					continue;
 				}
 
-				const head = await this.getSerializableHead(indexed.__context.head);
+				const head = headsByMatch[i];
 				results.push(
 					new types.ResultIndexedValue({
 						context: indexed.__context,
@@ -1253,7 +1315,7 @@ export class DocumentIndex<
 				);
 			} else {
 				const indexed = match as WithContext<I>;
-				const head = await this.getSerializableHead(indexed.__context.head);
+				const head = headsByMatch[i];
 				results.push(
 					new types.ResultIndexedValue({
 						context: indexed.__context,
@@ -1275,17 +1337,28 @@ export class DocumentIndex<
 			return [];
 		}
 		const drained = queueEntries.splice(0);
+		const resolvedBatch = resolve
+			? await this.resolveDocumentsWithBatchedHeads(
+					drained.map((entry) => ({
+						indexed: entry.value,
+						head: entry.value.__context.head,
+					})),
+				)
+			: undefined;
+		const heads = resolvedBatch
+			? resolvedBatch.heads
+			: await this._log.log.getMany(
+					drained.map((entry) => entry.value.__context.head),
+				);
 		const results: types.Result[] = [];
-		for (const entry of drained) {
+		for (let i = 0; i < drained.length; i++) {
+			const entry = drained[i]!;
 			const indexedUnwrapped = Object.assign(
 				Object.create(this.indexedType.prototype),
 				entry.value,
 			);
 			if (resolve) {
-				const value = await this.resolveDocument({
-					indexed: entry.value,
-					head: entry.value.__context.head,
-				});
+				const value = resolvedBatch!.resolved[i];
 				if (value) {
 					results.push(
 						new types.ResultValue({
@@ -1298,7 +1371,7 @@ export class DocumentIndex<
 					continue;
 				}
 
-				const head = await this.getSerializableHead(entry.value.__context.head);
+				const head = heads[i];
 				results.push(
 					new types.ResultIndexedValue({
 						context: entry.value.__context,
@@ -1308,7 +1381,7 @@ export class DocumentIndex<
 					}),
 				);
 			} else {
-				const head = await this.getSerializableHead(entry.value.__context.head);
+				const head = heads[i];
 				results.push(
 					new types.ResultIndexedValue({
 						context: entry.value.__context,
@@ -1358,6 +1431,7 @@ export class DocumentIndex<
 				let pendingAdded = added;
 				do {
 					const batches: types.Result[] = [];
+					let claimedIds: indexerTypes.IdKey[] = [];
 					const queued = await this.drainQueuedResults(
 						queue.queue,
 						resolveFlag,
@@ -1379,6 +1453,21 @@ export class DocumentIndex<
 						const wrapped = await this.wrapPushResults(matches, resolveFlag);
 						if (wrapped.length) {
 							batches.push(...wrapped);
+							claimedIds = wrapped
+								.map((result) =>
+									result instanceof types.ResultValue && result.indexed
+										? indexerTypes.toId(
+												this.indexByResolver(result.indexed),
+											)
+										: result instanceof types.ResultIndexedValue
+											? indexerTypes.toId(
+													this.indexByResolver(result.value),
+												)
+											: undefined,
+								)
+								.filter(
+									(id): id is indexerTypes.IdKey => id !== undefined,
+								);
 						}
 					}
 					if (batches.length) {
@@ -1396,6 +1485,12 @@ export class DocumentIndex<
 								redundancy: 1,
 							}),
 						});
+						if (claimedIds.length) {
+							await this._resumableIterators.markYielded(
+								_iteratorId,
+								claimedIds,
+							);
+						}
 					}
 					pendingAdded = queue.pendingAdded ?? [];
 					queue.pendingAdded = undefined;
@@ -3474,6 +3569,7 @@ export class DocumentIndex<
 		id?: indexerTypes.IdPrimitive;
 		indexed: I;
 		head: string;
+		headEntry?: Entry<Operation> | null;
 	}): Promise<{ value: T } | undefined> {
 		const id =
 			value.id ??
@@ -3495,7 +3591,10 @@ export class DocumentIndex<
 			return { value: obj as T };
 		}
 
-		const head = await this._log.log.get(value.head);
+		const head =
+			value.headEntry === undefined
+				? await this._log.log.get(value.head)
+				: (value.headEntry ?? undefined);
 		if (!head) {
 			return undefined; // we could end up here if we recently pruned the document and other peers never persisted the entry
 			// TODO update changes in index before removing entries from log entry storage
@@ -3633,7 +3732,10 @@ export class DocumentIndex<
 			this._resultQueue.set(query.idString, prevQueued);
 		}
 
-		const filteredResults: types.Result[] = [];
+		const candidates: Array<{
+			result: indexerTypes.IndexedResult<WithContext<I>>;
+			indexed: I;
+		}> = [];
 		const resolveDocumentsFlag = resolvesDocuments(fromQuery);
 		const replicateIndexFlag = replicatesIndex(fromQuery);
 		for (const result of toIterate) {
@@ -3655,11 +3757,29 @@ export class DocumentIndex<
 			) {
 				continue;
 			}
+			candidates.push({ result, indexed: indexedUnwrapped });
+		}
+
+		const resolvedBatch = resolveDocumentsFlag
+			? await this.resolveDocumentsWithBatchedHeads(
+					candidates.map(({ result }) => ({
+						indexed: result.value,
+						head: result.value.__context.head,
+					})),
+				)
+			: undefined;
+		const heads = resolvedBatch
+			? resolvedBatch.heads
+			: candidates.length > 0
+				? await this._log.log.getMany(
+						candidates.map(({ result }) => result.value.__context.head),
+					)
+				: [];
+		const filteredResults: types.Result[] = [];
+		for (let i = 0; i < candidates.length; i++) {
+			const { result, indexed: indexedUnwrapped } = candidates[i]!;
 			if (resolveDocumentsFlag) {
-				const value = await this.resolveDocument({
-					indexed: result.value,
-					head: result.value.__context.head,
-				});
+				const value = resolvedBatch!.resolved[i];
 
 				if (!value) {
 					continue;
@@ -3675,7 +3795,7 @@ export class DocumentIndex<
 				);
 			} else {
 				const context = result.value.__context;
-				const head = await this.getSerializableHead(context.head);
+				const head = heads[i];
 				if (replicateIndexFlag) {
 					if (!head) {
 						continue;
@@ -4165,6 +4285,7 @@ export class DocumentIndex<
 				let extraPromises: Promise<void>[] | undefined = undefined;
 
 				const seenRemoteHashes = new Set<string>();
+				const selectedRemoteHashes: string[] = [];
 				const groupHashes: string[][] = replicatorGroups
 					.filter((hash) => {
 						if (hash === this.node.identity.publicKey.hashcode()) {
@@ -4181,6 +4302,7 @@ export class DocumentIndex<
 							return false;
 						}
 						fetchFirstForRemote?.add(hash);
+						selectedRemoteHashes.push(hash);
 
 						const resultAlready = this._prefetch?.accumulator.consume(
 							queryRequest,
@@ -4206,6 +4328,7 @@ export class DocumentIndex<
 					})
 					.map((x) => [x]);
 
+				options?.onRemoteTargets?.(selectedRemoteHashes);
 				extraPromises && (await Promise.all(extraPromises));
 				let tearDown: (() => void) | undefined = undefined;
 				const search = this;
@@ -4364,23 +4487,25 @@ export class DocumentIndex<
 
 		const allResults: ValueTypeFromRequest<Resolve, T, I>[] = [];
 
-		while (
-			iterator.done() !== true &&
-			coercedRequest.fetch > allResults.length
-		) {
-			// We might need to pull .next multiple time due to data message size limitations
+		try {
+			while (
+				iterator.done() !== true &&
+				coercedRequest.fetch > allResults.length
+			) {
+				// We might need to pull .next multiple time due to data message size limitations
 
-			for (const result of await iterator.next(
-				coercedRequest.fetch - allResults.length,
-			)) {
-				allResults.push(result as ValueTypeFromRequest<Resolve, T, I>);
+				for (const result of await iterator.next(
+					coercedRequest.fetch - allResults.length,
+				)) {
+					allResults.push(result as ValueTypeFromRequest<Resolve, T, I>);
+				}
 			}
-		}
-
-		await iterator.close();
 
-		// Deduplicate and return values directly
-		return dedup(allResults, this.indexByResolver);
+			// Deduplicate and return values directly
+			return dedup(allResults, this.indexByResolver);
+		} finally {
+			await iterator.close();
+		}
 	}
 
 	private resolveIndexed<R>(
@@ -4672,6 +4797,50 @@ export class DocumentIndex<
 				buffer: BufferedResult<types.ResultTypeFromRequest<R, T, I> | I, I>[];
 			}
 		> = new Map();
+		const remoteIteratorPeersToClose = new Set<string>();
+		const staleRemoteIteratorIdsToClose = new Map<string, Uint8Array[]>();
+		const retiredRemoteIteratorPeers = new Set<string>();
+		const pendingMissingResponseRetryPeers = new Set<string>();
+		const missingResponseRetryAttempts = new Map<string, number>();
+		const maxMissingResponseRetryAttempts = 2;
+		const retireRemoteIteratorPeer = (peer: string) => {
+			remoteIteratorPeersToClose.add(peer);
+			retiredRemoteIteratorPeers.add(peer);
+			const peerBuffer = peerBufferMap.get(peer);
+			if (!peerBuffer || peerBuffer.buffer.length === 0) {
+				peerBufferMap.delete(peer);
+			} else {
+				// Keep already-received values available to the caller, but do not
+				// issue another CollectNextRequest until a fresh iteration succeeds.
+				peerBuffer.kept = 0;
+			}
+		};
+		const recordMissingResponseGroups = (missingGroups: string[][]) => {
+			const selfHash = this.node.identity.publicKey.hashcode();
+			for (const group of missingGroups) {
+				for (const hash of group) {
+					if (hash && hash !== selfHash) {
+						retireRemoteIteratorPeer(hash);
+					}
+				}
+
+				if (!retryMissingResponseGroups) {
+					continue;
+				}
+
+				const target = group.find((hash) => {
+					if (!hash || hash === selfHash) return false;
+					const attempts = missingResponseRetryAttempts.get(hash) ?? 0;
+					return attempts < maxMissingResponseRetryAttempts;
+				});
+				if (!target) continue;
+				pendingMissingResponseRetryPeers.add(target);
+				missingResponseRetryAttempts.set(
+					target,
+					(missingResponseRetryAttempts.get(target) ?? 0) + 1,
+				);
+			}
+		};
 		const visited = new Set<indexerTypes.IdPrimitive>();
 		let indexedPlaceholders:
 			| Map<
@@ -4903,6 +5072,17 @@ export class DocumentIndex<
 			n: number,
 			fetchOptions?: { from?: string[]; fetchedFirstForRemote?: Set<string> },
 		): Promise<boolean> => {
+			const remoteRequestOptions =
+				typeof options?.remote === "object" ? options.remote : undefined;
+			const fetchSignals = [
+				options?.signal,
+				remoteRequestOptions?.signal,
+				ensureController().signal,
+			].filter((signal): signal is AbortSignal => signal != null);
+			const fetchSignal =
+				fetchSignals.length === 1
+					? fetchSignals[0]
+					: AbortSignal.any(fetchSignals);
 			await warmupPromise;
 			let hasMore = false;
 			let missingResponses = false;
@@ -4918,33 +5098,43 @@ export class DocumentIndex<
 				typeof options?.remote === "object" &&
 				options.remote.reach?.discover &&
 				discoveredTargetHashes?.length === 0;
+			const queryRemote =
+				options?.remote !== false && !skipRemoteDueToDiscovery;
+			const remoteFrom =
+				fetchOptions?.from ??
+				initialRemoteTargets ??
+				remoteRequestOptions?.from;
+			if (queryRemote) {
+				const selfHash = this.node.identity.publicKey.hashcode();
+				for (const peer of remoteFrom ?? []) {
+					if (peer !== selfHash) {
+						remoteIteratorPeersToClose.add(peer);
+					}
+				}
+			}
 
 			queryRequestCoerced.fetch = n;
 			await this.queryCommence(
 				queryRequestCoerced,
 				{
 					local: fetchOptions?.from != null ? false : options?.local,
-					remote:
-						options?.remote !== false && !skipRemoteDueToDiscovery
-							? {
-									...(typeof options?.remote === "object"
-										? options.remote
-										: {}),
-									from:
-										fetchOptions?.from ??
-										initialRemoteTargets ??
-										(typeof options?.remote === "object"
-											? options.remote.from
-											: undefined),
-								}
-							: false,
+					remote: queryRemote
+						? {
+								...remoteRequestOptions,
+								from: remoteFrom,
+								signal: fetchSignal,
+							}
+						: false,
 					resolve,
-					signal: options?.signal,
+					signal: fetchSignal,
 					onResponse: async (response, from) => {
 						if (!from) {
 							logger.error("Missing response from");
 							return;
 						}
+						const fromHash = from.hashcode();
+						remoteIteratorPeersToClose.add(fromHash);
+						retiredRemoteIteratorPeers.delete(fromHash);
 						if (response instanceof types.NoAccess) {
 							logger.error("Dont have access");
 							return;
@@ -4953,7 +5143,7 @@ export class DocumentIndex<
 								types.ResultTypeFromRequest<R, T, I>
 							>;
 
-							const existingBuffer = peerBufferMap.get(from.hashcode());
+							const existingBuffer = peerBufferMap.get(fromHash);
 							const buffer: BufferedResult<
 								types.ResultTypeFromRequest<R, T, I> | I,
 								I
@@ -4961,10 +5151,12 @@ export class DocumentIndex<
 
 							if (results.kept === 0n && results.results.length === 0) {
 								if (keepRemoteAlive) {
-									peerBufferMap.set(from.hashcode(), {
+									peerBufferMap.set(fromHash, {
 										buffer,
 										kept: 0,
 									});
+								} else {
+									remoteIteratorPeersToClose.delete(fromHash);
 								}
 								return;
 							}
@@ -4979,6 +5171,8 @@ export class DocumentIndex<
 
 							if (effectiveKept > 0) {
 								hasMore = true;
+							} else if (!keepRemoteAlive) {
+								remoteIteratorPeersToClose.delete(fromHash);
 							}
 
 							for (const result of results.results) {
@@ -5042,7 +5236,7 @@ export class DocumentIndex<
 								}
 							}
 
-							peerBufferMap.set(from.hashcode(), {
+							peerBufferMap.set(fromHash, {
 								buffer,
 								kept: effectiveKept,
 							});
@@ -5054,9 +5248,6 @@ export class DocumentIndex<
 					},
 					onMissingResponses: (error) => {
 						missingResponses = true;
-						if (!retryMissingResponseGroups) {
-							return;
-						}
 						const missingGroups = (
 							error as MissingResponsesError & {
 								missingGroups?: string[][];
@@ -5065,27 +5256,22 @@ export class DocumentIndex<
 						if (!missingGroups?.length) {
 							return;
 						}
-
-						const selfHash = this.node.identity.publicKey.hashcode();
-						for (const group of missingGroups) {
-							const target = group.find((hash) => {
-								if (!hash || hash === selfHash) return false;
-								const attempts = missingResponseRetryAttempts.get(hash) ?? 0;
-								return attempts < maxMissingResponseRetryAttempts;
-							});
-							if (!target) continue;
-							pendingMissingResponseRetryPeers.add(target);
-							missingResponseRetryAttempts.set(
-								target,
-								(missingResponseRetryAttempts.get(target) ?? 0) + 1,
-							);
+						recordMissingResponseGroups(missingGroups);
+					},
+					onRemoteTargets: (targets) => {
+						for (const peer of targets) {
+							remoteIteratorPeersToClose.add(peer);
 						}
 					},
 				},
 				fetchOptions?.fetchedFirstForRemote,
 			);
 
-			if (missingResponses && retryMissingResponseGroups) {
+			if (
+				missingResponses &&
+				retryMissingResponseGroups &&
+				pendingMissingResponseRetryPeers.size > 0
+			) {
 				hasMore = true;
 				unsetDone();
 			}
@@ -5117,6 +5303,19 @@ export class DocumentIndex<
 			if (pendingMissingResponseRetryPeers.size > 0) {
 				const retryTargets = [...pendingMissingResponseRetryPeers];
 				pendingMissingResponseRetryPeers.clear();
+				const idTranslation =
+					this._prefetch?.accumulator.getTranslationMap(queryRequestCoerced);
+				for (const peer of retryTargets) {
+					const staleRemoteIteratorId = idTranslation?.get(peer);
+					if (staleRemoteIteratorId) {
+						const staleIds = staleRemoteIteratorIdsToClose.get(peer) ?? [];
+						if (!staleIds.some((id) => equals(id, staleRemoteIteratorId))) {
+							staleIds.push(staleRemoteIteratorId);
+							staleRemoteIteratorIdsToClose.set(peer, staleIds);
+						}
+						idTranslation!.delete(peer);
+					}
+				}
 				return setFetchPromise(
 					fetchFirst(n, {
 						from: retryTargets,
@@ -5130,6 +5329,12 @@ export class DocumentIndex<
 			let resultsLeft = 0;
 
 			for (const [peer, buffer] of peerBufferMap) {
+				if (retiredRemoteIteratorPeers.has(peer)) {
+					if (buffer.buffer.length === 0) {
+						peerBufferMap.delete(peer);
+					}
+					continue;
+				}
 				if (buffer.buffer.length < n) {
 					const hasExistingRemoteResults = buffer.kept > 0;
 					if (!hasExistingRemoteResults && !keepRemoteAlive) {
@@ -5272,10 +5477,11 @@ export class DocumentIndex<
 										"Failed to collect sorted results from self. " + e?.message,
 									);
 									peerBufferMap.delete(peer);
-								}),
+										}),
 						);
 					} else {
 						// Fetch remotely
+						remoteIteratorPeersToClose.add(peer);
 						const idTranslation =
 							this._prefetch?.accumulator.getTranslationMap(
 								queryRequestCoerced,
@@ -5287,22 +5493,45 @@ export class DocumentIndex<
 								amount: collectRequest.amount,
 							});
 						}
+						const remoteRequestOptions =
+							typeof options?.remote === "object" ? options.remote : undefined;
+						const collectSignals = [
+							options?.signal,
+							remoteRequestOptions?.signal,
+							ensureController().signal,
+						].filter((signal): signal is AbortSignal => signal != null);
+						const collectSignal =
+							collectSignals.length === 1
+								? collectSignals[0]
+								: AbortSignal.any(collectSignals);
 
 						promises.push(
 							this._query
 								.request(remoteCollectRequest, {
 									...options,
-									signal: options?.signal
-										? AbortSignal.any([
-												options.signal,
-												ensureController().signal,
-											])
-										: ensureController().signal,
+									...remoteRequestOptions,
+									signal: collectSignal,
 									priority: getRemoteQueryPriority(options?.remote),
 									mode: new SilentDelivery({ to: [peer], redundancy: 1 }),
 								})
-								.then((response) =>
-									introduceEntries(
+								.then((response) => {
+									if (
+										!response.some((result) => result.from?.hashcode() === peer)
+									) {
+										const missingGroups = [[peer]];
+										if (remoteRequestOptions?.throwOnMissing) {
+											retireRemoteIteratorPeer(peer);
+											throw new MissingResponsesError(
+												"Did not receive responses from all shards: " +
+													JSON.stringify(missingGroups),
+												missingGroups,
+											);
+										}
+										recordMissingResponseGroups(missingGroups);
+										return;
+									}
+
+									return introduceEntries(
 										queryRequestCoerced,
 										response,
 										this.documentType,
@@ -5319,6 +5548,12 @@ export class DocumentIndex<
 														logger.error("Missing from for sorted query");
 														return;
 													}
+													if (
+														!keepRemoteAlive &&
+														response.response.kept === 0n
+													) {
+														remoteIteratorPeersToClose.delete(peer);
+													}
 
 													if (response.response.results.length === 0) {
 														if (
@@ -5439,8 +5674,8 @@ export class DocumentIndex<
 													e?.message,
 											);
 											peerBufferMap.delete(peer);
-										}),
-								),
+										});
+								}),
 						);
 					}
 				} else {
@@ -5466,7 +5701,9 @@ export class DocumentIndex<
 							}
 						}
 					}
-					return resultsLeft === 0; // 0 results left to fetch and 0 pending results
+					return (
+						resultsLeft === 0 && pendingMissingResponseRetryPeers.size === 0
+					); // 0 results left to fetch and 0 pending results
 				}),
 			);
 		};
@@ -5587,32 +5824,89 @@ export class DocumentIndex<
 			done = true;
 		};
 
-		let close = async () => {
+		const outerSignal = options?.signal;
+		let outerAbortListener: (() => void) | undefined;
+		let closeStarted = false;
+		let closePromise: Promise<void> | undefined;
+		const performClose = async () => {
+			const idTranslation =
+				this._prefetch?.accumulator.getTranslationMap(queryRequestCoerced);
+			const remoteIteratorIds = idTranslation
+				? new Map(idTranslation)
+				: undefined;
 			cleanupAndDone();
 
 			// Keep-open iterators can still have active remote state even when
 			// their pending count has already drained to zero.
-			const closeRequest = new types.CloseIteratorRequest({
-				id: queryRequestCoerced.id,
-			});
 			const selfHash = this.node.identity.publicKey.hashcode();
-			const remotePeers = keepRemoteAlive
-				? [...peerBufferMap.keys()].filter((peer) => peer !== selfHash)
-				: [...peerBufferMap.entries()]
-						.filter(([peer, buffer]) => peer !== selfHash && buffer.kept > 0)
-						.map(([peer]) => peer);
+			const activeRemotePeers = new Set(
+				keepRemoteAlive
+					? [...peerBufferMap.keys()].filter((peer) => peer !== selfHash)
+					: [...peerBufferMap.entries()]
+							.filter(
+								([peer, buffer]) => peer !== selfHash && buffer.kept > 0,
+							)
+							.map(([peer]) => peer),
+			);
+			for (const peer of remoteIteratorPeersToClose) {
+				if (peer !== selfHash) {
+					activeRemotePeers.add(peer);
+				}
+			}
+			const staleRemoteIteratorIds = new Map(
+				[...staleRemoteIteratorIdsToClose].map(([peer, ids]) => [
+					peer,
+					[...ids],
+				]),
+			);
+			const remotePeers = new Set([
+				...activeRemotePeers,
+				...staleRemoteIteratorIds.keys(),
+			]);
 			peerBufferMap.clear();
+			retiredRemoteIteratorPeers.clear();
+			remoteIteratorPeersToClose.clear();
+			staleRemoteIteratorIdsToClose.clear();
+			if (remotePeers.size === 0) {
+				return;
+			}
+			const remoteCloseOptions =
+				typeof options?.remote === "object" ? options.remote : undefined;
+			const closeSignal = AbortSignal.timeout(CLOSE_ITERATOR_REQUEST_TIMEOUT);
 			await Promise.allSettled(
-				remotePeers.map((peer) =>
-					this._query.send(closeRequest, {
-						...options,
-						priority: getRemoteQueryPriority(options?.remote),
-						mode: new SilentDelivery({ to: [peer], redundancy: 1 }),
-					}),
-				),
+				[...remotePeers].flatMap((peer) => {
+					const ids: Uint8Array[] = [];
+					if (activeRemotePeers.has(peer)) {
+						ids.push(remoteIteratorIds?.get(peer) ?? queryRequestCoerced.id);
+					}
+					for (const staleRemoteIteratorId of
+						staleRemoteIteratorIds.get(peer) ?? []) {
+						if (!ids.some((id) => equals(id, staleRemoteIteratorId))) {
+							ids.push(staleRemoteIteratorId);
+						}
+					}
+					return ids.map((id) => {
+						const closeRequest = new types.CloseIteratorRequest({ id });
+						return this._query.send(closeRequest, {
+							...remoteCloseOptions,
+							signal: closeSignal,
+							priority: getRemoteQueryPriority(options?.remote),
+							mode: new SilentDelivery({ to: [peer], redundancy: 1 }),
+						});
+					});
+				}),
 			);
 		};
-		options?.signal && options.signal.addEventListener("abort", close);
+		const close = () => {
+			if (closeStarted) return closePromise!;
+			closeStarted = true;
+			if (outerAbortListener) {
+				outerSignal?.removeEventListener("abort", outerAbortListener);
+				outerAbortListener = undefined;
+			}
+			closePromise = performClose();
+			return closePromise;
+		};
 
 		let doneFn = () => {
 			return done;
@@ -5621,9 +5915,6 @@ export class DocumentIndex<
 		let joinListener: (() => void) | undefined;
 
 		let fetchedFirstForRemote: Set<string> | undefined = undefined;
-		const pendingMissingResponseRetryPeers = new Set<string>();
-		const missingResponseRetryAttempts = new Map<string, number>();
-		const maxMissingResponseRetryAttempts = 2;
 		let joinFetchesInFlight = 0;
 
 		let updateDeferred: ReturnType<typeof pDefer> | undefined;
@@ -6144,9 +6435,10 @@ export class DocumentIndex<
 									continue;
 								}
 							}
-							const id = indexerTypes.toId(
+							const indexId = indexerTypes.toId(
 								this.indexByResolver(indexedCandidate),
-							).primitive;
+							);
+							const id = indexId.primitive;
 							const existingIndexed = indexedPlaceholders?.get(id);
 							if (existingIndexed) {
 								if (resolve) {
@@ -6187,6 +6479,10 @@ export class DocumentIndex<
 							if (!resolve) {
 								ensureIndexedPlaceholders().set(id, placeholder);
 							}
+							await this._resumableIterators.markYielded(
+								queryRequestCoerced.idString,
+								[indexId],
+							);
 							hasRelevantChange = true;
 						}
 					}
@@ -6458,6 +6754,16 @@ export class DocumentIndex<
 			}
 		};
 
+		if (outerSignal) {
+			outerAbortListener = () => {
+				void close();
+			};
+			outerSignal.addEventListener("abort", outerAbortListener, { once: true });
+			if (outerSignal.aborted) {
+				void close();
+			}
+		}
+
 		return {
 			close,
 			next,
@@ -6531,53 +6837,62 @@ export class DocumentIndex<
 				const drainBatchSize = replicate ? 1000 : 100;
 				let result: ValueTypeFromRequest<Resolve, T, I>[] = [];
 				let c = 0;
-				while (doneFn() !== true) {
-					let batch = await next(drainBatchSize);
-					c += batch.length;
-					if (c > WARNING_WHEN_ITERATING_FOR_MORE_THAN) {
-						warn(
-							"Iterating for more than " +
-								WARNING_WHEN_ITERATING_FOR_MORE_THAN +
-								" results",
-						);
-					}
-					if (batch.length > 0) {
-						result.push(...batch);
-						continue;
+				try {
+					while (doneFn() !== true) {
+						let batch = await next(drainBatchSize);
+						c += batch.length;
+						if (c > WARNING_WHEN_ITERATING_FOR_MORE_THAN) {
+							warn(
+								"Iterating for more than " +
+									WARNING_WHEN_ITERATING_FOR_MORE_THAN +
+									" results",
+							);
+						}
+						if (batch.length > 0) {
+							result.push(...batch);
+							continue;
+						}
+						await waitForUpdateAndResetDeferred();
 					}
-					await waitForUpdateAndResetDeferred();
+					return result;
+				} finally {
+					await close();
 				}
-				cleanupAndDone();
-				return result;
 			},
 			first: async () => {
-				if (doneFn()) {
-					return undefined;
+				try {
+					if (doneFn()) {
+						return undefined;
+					}
+					let batch = await next(1);
+					return batch[0];
+				} finally {
+					await close();
 				}
-				let batch = await next(1);
-				cleanupAndDone();
-				return batch[0];
 			},
 			[Symbol.asyncIterator]: async function* () {
 				drain = true;
 				const drainBatchSize = replicate ? 1000 : 100;
 				let c = 0;
-				while (doneFn() !== true) {
-					const batch = await next(drainBatchSize);
-					c += batch.length;
-					if (c > WARNING_WHEN_ITERATING_FOR_MORE_THAN) {
-						warn(
-							"Iterating for more than " +
-								WARNING_WHEN_ITERATING_FOR_MORE_THAN +
-								" results",
-						);
-					}
-					for (const entry of batch) {
-						yield entry;
+				try {
+					while (doneFn() !== true) {
+						const batch = await next(drainBatchSize);
+						c += batch.length;
+						if (c > WARNING_WHEN_ITERATING_FOR_MORE_THAN) {
+							warn(
+								"Iterating for more than " +
+									WARNING_WHEN_ITERATING_FOR_MORE_THAN +
+									" results",
+							);
+						}
+						for (const entry of batch) {
+							yield entry;
+						}
+						await waitForUpdateAndResetDeferred();
 					}
-					await waitForUpdateAndResetDeferred();
+				} finally {
+					await close();
 				}
-				cleanupAndDone();
 			},
 		};
 	}