Skip to content

Commit 9d73f45

Browse files
nbbeekennbbeeken
authored
refactor(NODE-6057): implement CursorResponse for lazy document parsing (#4085)
1 parent 6d8ad33 commit 9d73f45

File tree

20 files changed

+584
-138
lines changed

20 files changed

+584
-138
lines changed

‎src/cmap/connection.ts

+14-5
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,11 @@ import type { ClientMetadata } from './handshake/client_metadata';
6262
import { StreamDescription, type StreamDescriptionOptions } from './stream_description';
6363
import { type CompressorName, decompressResponse } from './wire_protocol/compression';
6464
import { onData } from './wire_protocol/on_data';
65-
import { MongoDBResponse, type MongoDBResponseConstructor } from './wire_protocol/responses';
65+
import {
66+
isErrorResponse,
67+
MongoDBResponse,
68+
type MongoDBResponseConstructor
69+
} from './wire_protocol/responses';
6670
import { getReadPreference, isSharded } from './wire_protocol/shared';
6771

6872
/** @internal */
@@ -443,7 +447,12 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
443447
this.socket.setTimeout(0);
444448
const bson = response.parse();
445449

446-
const document = new (responseType ?? MongoDBResponse)(bson, 0, false);
450+
const document =
451+
responseType == null
452+
? new MongoDBResponse(bson)
453+
: isErrorResponse(bson)
454+
? new MongoDBResponse(bson)
455+
: new responseType(bson);
447456

448457
yield document;
449458
this.throwIfAborted();
@@ -739,7 +748,7 @@ export class CryptoConnection extends Connection {
739748
ns: MongoDBNamespace,
740749
cmd: Document,
741750
options?: CommandOptions,
742-
responseType?: T | undefined
751+
_responseType?: T | undefined
743752
): Promise<Document> {
744753
const { autoEncrypter } = this;
745754
if (!autoEncrypter) {
@@ -753,7 +762,7 @@ export class CryptoConnection extends Connection {
753762
const serverWireVersion = maxWireVersion(this);
754763
if (serverWireVersion === 0) {
755764
// This means the initial handshake hasn't happened yet
756-
return await super.command<T>(ns, cmd, options, responseType);
765+
return await super.command<T>(ns, cmd, options, undefined);
757766
}
758767

759768
if (serverWireVersion < 8) {
@@ -787,7 +796,7 @@ export class CryptoConnection extends Connection {
787796
}
788797
}
789798

790-
const response = await super.command<T>(ns, encrypted, options, responseType);
799+
const response = await super.command<T>(ns, encrypted, options, undefined);
791800

792801
return await autoEncrypter.decrypt(response, options);
793802
}

‎src/cmap/wire_protocol/on_demand/document.ts

+36-21
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ export class OnDemandDocument {
5858
private readonly indexFound: Record<number, boolean> = Object.create(null);
5959

6060
/** All bson elements in this document */
61-
private readonly elements: BSONElement[];
61+
private readonly elements: ReadonlyArray<BSONElement>;
6262

6363
constructor(
6464
/** BSON bytes, this document begins at offset */
@@ -97,14 +97,30 @@ export class OnDemandDocument {
9797
* @param name - a basic latin string name of a BSON element
9898
* @returns
9999
*/
100-
private getElement(name: string): CachedBSONElement | null {
100+
private getElement(name: string | number): CachedBSONElement | null {
101101
const cachedElement = this.cache[name];
102102
if (cachedElement === false) return null;
103103

104104
if (cachedElement != null) {
105105
return cachedElement;
106106
}
107107

108+
if (typeof name === 'number') {
109+
if (this.isArray) {
110+
if (name < this.elements.length) {
111+
const element = this.elements[name];
112+
const cachedElement = { element, value: undefined };
113+
this.cache[name] = cachedElement;
114+
this.indexFound[name] = true;
115+
return cachedElement;
116+
} else {
117+
return null;
118+
}
119+
} else {
120+
return null;
121+
}
122+
}
123+
108124
for (let index = 0; index < this.elements.length; index++) {
109125
const element = this.elements[index];
110126

@@ -197,6 +213,13 @@ export class OnDemandDocument {
197213
}
198214
}
199215

216+
/**
217+
* Returns the number of elements in this BSON document
218+
*/
219+
public size() {
220+
return this.elements.length;
221+
}
222+
200223
/**
201224
* Checks for the existence of an element by name.
202225
*
@@ -222,16 +245,20 @@ export class OnDemandDocument {
222245
* @param required - whether or not the element is expected to exist, if true this function will throw if it is not present
223246
*/
224247
public get<const T extends keyof JSTypeOf>(
225-
name: string,
248+
name: string | number,
226249
as: T,
227250
required?: false | undefined
228251
): JSTypeOf[T] | null;
229252

230253
/** `required` will make `get` throw if name does not exist or is null/undefined */
231-
public get<const T extends keyof JSTypeOf>(name: string, as: T, required: true): JSTypeOf[T];
254+
public get<const T extends keyof JSTypeOf>(
255+
name: string | number,
256+
as: T,
257+
required: true
258+
): JSTypeOf[T];
232259

233260
public get<const T extends keyof JSTypeOf>(
234-
name: string,
261+
name: string | number,
235262
as: T,
236263
required?: boolean
237264
): JSTypeOf[T] | null {
@@ -303,21 +330,9 @@ export class OnDemandDocument {
303330
});
304331
}
305332

306-
/**
307-
* Iterates through the elements of a document reviving them using the `as` BSONType.
308-
*
309-
* @param as - The type to revive all elements as
310-
*/
311-
public *valuesAs<const T extends keyof JSTypeOf>(as: T): Generator<JSTypeOf[T]> {
312-
if (!this.isArray) {
313-
throw new BSONError('Unexpected conversion of non-array value to array');
314-
}
315-
let counter = 0;
316-
for (const element of this.elements) {
317-
const value = this.toJSValue<T>(element, as);
318-
this.cache[counter] = { element, value };
319-
yield value;
320-
counter += 1;
321-
}
333+
/** Returns this document's bytes only */
334+
toBytes() {
335+
const size = getInt32LE(this.bson, this.offset);
336+
return this.bson.subarray(this.offset, this.offset + size);
322337
}
323338
}

‎src/cmap/wire_protocol/responses.ts

+139-11
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,73 @@
1-
import { type BSONSerializeOptions, BSONType, type Document, type Timestamp } from '../../bson';
1+
import {
2+
type BSONSerializeOptions,
3+
BSONType,
4+
type Document,
5+
Long,
6+
parseToElementsToArray,
7+
type Timestamp
8+
} from '../../bson';
9+
import { MongoUnexpectedServerResponseError } from '../../error';
210
import { type ClusterTime } from '../../sdam/common';
11+
import { type MongoDBNamespace, ns } from '../../utils';
312
import { OnDemandDocument } from './on_demand/document';
413

14+
// eslint-disable-next-line no-restricted-syntax
15+
const enum BSONElementOffset {
16+
type = 0,
17+
nameOffset = 1,
18+
nameLength = 2,
19+
offset = 3,
20+
length = 4
21+
}
22+
/**
23+
* Accepts a BSON payload and checks for na "ok: 0" element.
24+
* This utility is intended to prevent calling response class constructors
25+
* that expect the result to be a success and demand certain properties to exist.
26+
*
27+
* For example, a cursor response always expects a cursor embedded document.
28+
* In order to write the class such that the properties reflect that assertion (non-null)
29+
* we cannot invoke the subclass constructor if the BSON represents an error.
30+
*
31+
* @param bytes - BSON document returned from the server
32+
*/
33+
export function isErrorResponse(bson: Uint8Array): boolean {
34+
const elements = parseToElementsToArray(bson, 0);
35+
for (let eIdx = 0; eIdx < elements.length; eIdx++) {
36+
const element = elements[eIdx];
37+
38+
if (element[BSONElementOffset.nameLength] === 2) {
39+
const nameOffset = element[BSONElementOffset.nameOffset];
40+
41+
// 111 == "o", 107 == "k"
42+
if (bson[nameOffset] === 111 && bson[nameOffset + 1] === 107) {
43+
const valueOffset = element[BSONElementOffset.offset];
44+
const valueLength = element[BSONElementOffset.length];
45+
46+
// If any byte in the length of the ok number (works for any type) is non zero,
47+
// then it is considered "ok: 1"
48+
for (let i = valueOffset; i < valueOffset + valueLength; i++) {
49+
if (bson[i] !== 0x00) return false;
50+
}
51+
52+
return true;
53+
}
54+
}
55+
}
56+
57+
return true;
58+
}
59+
560
/** @internal */
661
export type MongoDBResponseConstructor = {
762
new (bson: Uint8Array, offset?: number, isArray?: boolean): MongoDBResponse;
863
};
964

1065
/** @internal */
1166
export class MongoDBResponse extends OnDemandDocument {
67+
static is(value: unknown): value is MongoDBResponse {
68+
return value instanceof MongoDBResponse;
69+
}
70+
1271
// {ok:1}
1372
static empty = new MongoDBResponse(new Uint8Array([13, 0, 0, 0, 16, 111, 107, 0, 1, 0, 0, 0, 0]));
1473

@@ -83,27 +142,96 @@ export class MongoDBResponse extends OnDemandDocument {
83142
return this.clusterTime ?? null;
84143
}
85144

86-
public override toObject(options: BSONSerializeOptions = {}): Record<string, any> {
145+
public override toObject(options?: BSONSerializeOptions): Record<string, any> {
87146
const exactBSONOptions = {
88-
useBigInt64: options.useBigInt64,
89-
promoteLongs: options.promoteLongs,
90-
promoteValues: options.promoteValues,
91-
promoteBuffers: options.promoteBuffers,
92-
bsonRegExp: options.bsonRegExp,
93-
raw: options.raw ?? false,
94-
fieldsAsRaw: options.fieldsAsRaw ?? {},
147+
useBigInt64: options?.useBigInt64,
148+
promoteLongs: options?.promoteLongs,
149+
promoteValues: options?.promoteValues,
150+
promoteBuffers: options?.promoteBuffers,
151+
bsonRegExp: options?.bsonRegExp,
152+
raw: options?.raw ?? false,
153+
fieldsAsRaw: options?.fieldsAsRaw ?? {},
95154
validation: this.parseBsonSerializationOptions(options)
96155
};
97156
return super.toObject(exactBSONOptions);
98157
}
99158

100-
private parseBsonSerializationOptions({ enableUtf8Validation }: BSONSerializeOptions): {
159+
private parseBsonSerializationOptions(options?: { enableUtf8Validation?: boolean }): {
101160
utf8: { writeErrors: false } | false;
102161
} {
162+
const enableUtf8Validation = options?.enableUtf8Validation;
103163
if (enableUtf8Validation === false) {
104164
return { utf8: false };
105165
}
106-
107166
return { utf8: { writeErrors: false } };
108167
}
109168
}
169+
170+
/** @internal */
171+
export class CursorResponse extends MongoDBResponse {
172+
/**
173+
* This supports a feature of the FindCursor.
174+
* It is an optimization to avoid an extra getMore when the limit has been reached
175+
*/
176+
static emptyGetMore = { id: new Long(0), length: 0, shift: () => null };
177+
178+
static override is(value: unknown): value is CursorResponse {
179+
return value instanceof CursorResponse || value === CursorResponse.emptyGetMore;
180+
}
181+
182+
public id: Long;
183+
public ns: MongoDBNamespace | null = null;
184+
public batchSize = 0;
185+
186+
private batch: OnDemandDocument;
187+
private iterated = 0;
188+
189+
constructor(bytes: Uint8Array, offset?: number, isArray?: boolean) {
190+
super(bytes, offset, isArray);
191+
192+
const cursor = this.get('cursor', BSONType.object, true);
193+
194+
const id = cursor.get('id', BSONType.long, true);
195+
this.id = new Long(Number(id & 0xffff_ffffn), Number((id >> 32n) & 0xffff_ffffn));
196+
197+
const namespace = cursor.get('ns', BSONType.string);
198+
if (namespace != null) this.ns = ns(namespace);
199+
200+
if (cursor.has('firstBatch')) this.batch = cursor.get('firstBatch', BSONType.array, true);
201+
else if (cursor.has('nextBatch')) this.batch = cursor.get('nextBatch', BSONType.array, true);
202+
else throw new MongoUnexpectedServerResponseError('Cursor document did not contain a batch');
203+
204+
this.batchSize = this.batch.size();
205+
}
206+
207+
get length() {
208+
return Math.max(this.batchSize - this.iterated, 0);
209+
}
210+
211+
shift(options?: BSONSerializeOptions): any {
212+
if (this.iterated >= this.batchSize) {
213+
return null;
214+
}
215+
216+
const result = this.batch.get(this.iterated, BSONType.object, true) ?? null;
217+
this.iterated += 1;
218+
219+
if (options?.raw) {
220+
return result.toBytes();
221+
} else {
222+
return result.toObject(options);
223+
}
224+
}
225+
226+
clear() {
227+
this.iterated = this.batchSize;
228+
}
229+
230+
pushMany() {
231+
throw new Error('pushMany Unsupported method');
232+
}
233+
234+
push() {
235+
throw new Error('push Unsupported method');
236+
}
237+
}

0 commit comments

Comments
 (0)