Skip to content

Commit a04a119

Browse files
schmidt-sebastianschmidt-sebastian
committed
feat: retry BatchGetDocuments RPCs that fail with errors
1 parent 2406f6a commit a04a119

File tree

4 files changed

+265
-148
lines changed

4 files changed

+265
-148
lines changed

‎dev/src/document-reader.ts

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
/*!
2+
* Copyright 2021 Google LLC. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import {DocumentSnapshot, DocumentSnapshotBuilder} from './document';
18+
import {DocumentReference} from './reference';
19+
import {FieldPath} from './path';
20+
import {isPermanentRpcError} from './util';
21+
import {google} from '../protos/firestore_v1_proto_api';
22+
import {logger} from './logger';
23+
import {Firestore} from './index';
24+
25+
import api = google.firestore.v1;
26+
27+
/**
28+
* A wrapper around BatchGetDocumentsRequest that retries request upon stream
29+
* failure and returns ordered results.
30+
*
31+
* @private
32+
*/
33+
export class DocumentReader<T> {
34+
/** An optional field mask to apply to this read. */
35+
fieldMask?: FieldPath[];
36+
/** An optional transaction ID to use for this read. */
37+
transactionId?: Uint8Array;
38+
39+
private remainingDocuments = new Set<string>();
40+
private retrievedDocuments = new Map<string, DocumentSnapshot>();
41+
42+
/**
43+
* Internal method to retrieve multiple documents from Firestore, optionally
44+
* as part of a transaction.
45+
*
46+
* @param firestore The Firestore instance to use.
47+
* @param allDocuments The documents to receive.
48+
* @returns A Promise that contains an array with the resulting documents.
49+
*/
50+
constructor(
51+
private firestore: Firestore,
52+
private allDocuments: Array<DocumentReference<T>>
53+
) {
54+
for (const docRef of this.allDocuments) {
55+
this.remainingDocuments.add(docRef.formattedName);
56+
}
57+
}
58+
59+
/**
60+
* Invokes the BatchGetDocuments RPC and returns the results.
61+
*
62+
* @param requestTag A unique client-assigned identifier for this request.
63+
*/
64+
async get(requestTag: string): Promise<Array<DocumentSnapshot<T>>> {
65+
await this.fetchAllDocuments(requestTag);
66+
67+
// BatchGetDocuments doesn't preserve document order. We use the request
68+
// order to sort the resulting documents.
69+
const orderedDocuments: Array<DocumentSnapshot<T>> = [];
70+
71+
for (const docRef of this.allDocuments) {
72+
const document = this.retrievedDocuments.get(docRef.formattedName);
73+
if (document !== undefined) {
74+
// Recreate the DocumentSnapshot with the DocumentReference
75+
// containing the original converter.
76+
const finalDoc = new DocumentSnapshotBuilder(
77+
docRef as DocumentReference<T>
78+
);
79+
finalDoc.fieldsProto = document._fieldsProto;
80+
finalDoc.readTime = document.readTime;
81+
finalDoc.createTime = document.createTime;
82+
finalDoc.updateTime = document.updateTime;
83+
orderedDocuments.push(finalDoc.build());
84+
} else {
85+
throw new Error(`Did not receive document for "${docRef.path}".`);
86+
}
87+
}
88+
89+
return orderedDocuments;
90+
}
91+
92+
private async fetchAllDocuments(requestTag: string): Promise<void> {
93+
while (this.remainingDocuments.size > 0) {
94+
try {
95+
return await this.fetchMoreDocuments(requestTag);
96+
} catch (err) {
97+
// If a non-transactional read failed, attempt to restart.
98+
// Transactional reads are retried via the transaction runner.
99+
if (
100+
this.transactionId ||
101+
isPermanentRpcError(err, 'batchGetDocuments')
102+
) {
103+
logger(
104+
'DocumentReader.fetchAllDocuments',
105+
requestTag,
106+
'BatchGetDocuments failed with non-retryable stream error:',
107+
err
108+
);
109+
throw err;
110+
} else {
111+
logger(
112+
'DocumentReader.fetchAllDocuments',
113+
requestTag,
114+
'BatchGetDocuments failed with retryable stream error:',
115+
err
116+
);
117+
}
118+
}
119+
}
120+
}
121+
122+
private fetchMoreDocuments(requestTag: string): Promise<void> {
123+
const request: api.IBatchGetDocumentsRequest = {
124+
database: this.firestore.formattedName,
125+
transaction: this.transactionId,
126+
documents: Array.from(this.remainingDocuments),
127+
};
128+
129+
if (this.fieldMask) {
130+
const fieldPaths = this.fieldMask.map(
131+
fieldPath => (fieldPath as FieldPath).formattedName
132+
);
133+
request.mask = {fieldPaths};
134+
}
135+
136+
let resultCount = 0;
137+
138+
return this.firestore
139+
.requestStream('batchGetDocuments', request, requestTag)
140+
.then(stream => {
141+
return new Promise<void>((resolve, reject) => {
142+
stream
143+
.on('error', err => reject(err))
144+
.on('data', (response: api.IBatchGetDocumentsResponse) => {
145+
try {
146+
let document;
147+
148+
if (response.found) {
149+
logger(
150+
'DocumentReader.fetchMoreDocuments',
151+
requestTag,
152+
'Received document: %s',
153+
response.found.name!
154+
);
155+
document = this.firestore.snapshot_(
156+
response.found,
157+
response.readTime!
158+
);
159+
} else {
160+
logger(
161+
'DocumentReader.fetchMoreDocuments',
162+
requestTag,
163+
'Document missing: %s',
164+
response.missing!
165+
);
166+
document = this.firestore.snapshot_(
167+
response.missing!,
168+
response.readTime!
169+
);
170+
}
171+
172+
const path = document.ref.formattedName;
173+
this.remainingDocuments.delete(path);
174+
this.retrievedDocuments.set(path, document);
175+
++resultCount;
176+
} catch (err) {
177+
reject(err);
178+
}
179+
})
180+
.on('end', () => {
181+
logger(
182+
'DocumentReader.fetchMoreDocuments',
183+
requestTag,
184+
'Received %d results',
185+
resultCount
186+
);
187+
resolve();
188+
});
189+
stream.resume();
190+
});
191+
});
192+
}
193+
}

‎dev/src/index.ts

Lines changed: 7 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import {ExponentialBackoff, ExponentialBackoffSetting} from './backoff';
2626
import {BulkWriter} from './bulk-writer';
2727
import {BundleBuilder} from './bundle';
2828
import {fieldsFromJson, timestampFromJson} from './convert';
29+
import {DocumentReader} from './document-reader';
2930
import {
3031
DocumentSnapshot,
3132
DocumentSnapshotBuilder,
@@ -34,14 +35,12 @@ import {
3435
import {logger, setLibVersion} from './logger';
3536
import {
3637
DEFAULT_DATABASE_ID,
37-
FieldPath,
3838
QualifiedResourcePath,
3939
ResourcePath,
4040
validateResourcePath,
4141
} from './path';
4242
import {ClientPool} from './pool';
43-
import {CollectionReference} from './reference';
44-
import {DocumentReference} from './reference';
43+
import {CollectionReference, DocumentReference} from './reference';
4544
import {Serializer} from './serializer';
4645
import {Timestamp} from './timestamp';
4746
import {parseGetAllArguments, Transaction} from './transaction';
@@ -1077,138 +1076,16 @@ export class Firestore implements firestore.Firestore {
10771076
const stack = Error().stack!;
10781077

10791078
return this.initializeIfNeeded(tag)
1080-
.then(() => this.getAll_(documents, fieldMask, tag))
1079+
.then(() => {
1080+
const reader = new DocumentReader(this, documents);
1081+
reader.fieldMask = fieldMask || undefined;
1082+
return reader.get(tag);
1083+
})
10811084
.catch(err => {
10821085
throw wrapError(err, stack);
10831086
});
10841087
}
10851088

1086-
/**
1087-
* Internal method to retrieve multiple documents from Firestore, optionally
1088-
* as part of a transaction.
1089-
*
1090-
* @private
1091-
* @param docRefs The documents to receive.
1092-
* @param fieldMask An optional field mask to apply to this read.
1093-
* @param requestTag A unique client-assigned identifier for this request.
1094-
* @param transactionId The transaction ID to use for this read.
1095-
* @returns A Promise that contains an array with the resulting documents.
1096-
*/
1097-
getAll_<T>(
1098-
docRefs: Array<firestore.DocumentReference<T>>,
1099-
fieldMask: firestore.FieldPath[] | null,
1100-
requestTag: string,
1101-
transactionId?: Uint8Array
1102-
): Promise<Array<DocumentSnapshot<T>>> {
1103-
const requestedDocuments = new Set<string>();
1104-
const retrievedDocuments = new Map<string, DocumentSnapshot>();
1105-
1106-
for (const docRef of docRefs) {
1107-
requestedDocuments.add((docRef as DocumentReference<T>).formattedName);
1108-
}
1109-
1110-
const request: api.IBatchGetDocumentsRequest = {
1111-
database: this.formattedName,
1112-
transaction: transactionId,
1113-
documents: Array.from(requestedDocuments),
1114-
};
1115-
1116-
if (fieldMask) {
1117-
const fieldPaths = fieldMask.map(
1118-
fieldPath => (fieldPath as FieldPath).formattedName
1119-
);
1120-
request.mask = {fieldPaths};
1121-
}
1122-
1123-
return this.requestStream('batchGetDocuments', request, requestTag).then(
1124-
stream => {
1125-
return new Promise<Array<DocumentSnapshot<T>>>((resolve, reject) => {
1126-
stream
1127-
.on('error', err => {
1128-
logger(
1129-
'Firestore.getAll_',
1130-
requestTag,
1131-
'GetAll failed with error:',
1132-
err
1133-
);
1134-
reject(err);
1135-
})
1136-
.on('data', (response: api.IBatchGetDocumentsResponse) => {
1137-
try {
1138-
let document;
1139-
1140-
if (response.found) {
1141-
logger(
1142-
'Firestore.getAll_',
1143-
requestTag,
1144-
'Received document: %s',
1145-
response.found.name!
1146-
);
1147-
document = this.snapshot_(response.found, response.readTime!);
1148-
} else {
1149-
logger(
1150-
'Firestore.getAll_',
1151-
requestTag,
1152-
'Document missing: %s',
1153-
response.missing!
1154-
);
1155-
document = this.snapshot_(
1156-
response.missing!,
1157-
response.readTime!
1158-
);
1159-
}
1160-
1161-
const path = document.ref.path;
1162-
retrievedDocuments.set(path, document);
1163-
} catch (err) {
1164-
logger(
1165-
'Firestore.getAll_',
1166-
requestTag,
1167-
'GetAll failed with exception:',
1168-
err
1169-
);
1170-
reject(err);
1171-
}
1172-
})
1173-
.on('end', () => {
1174-
logger(
1175-
'Firestore.getAll_',
1176-
requestTag,
1177-
'Received %d results',
1178-
retrievedDocuments.size
1179-
);
1180-
1181-
// BatchGetDocuments doesn't preserve document order. We use
1182-
// the request order to sort the resulting documents.
1183-
const orderedDocuments: Array<DocumentSnapshot<T>> = [];
1184-
1185-
for (const docRef of docRefs) {
1186-
const document = retrievedDocuments.get(docRef.path);
1187-
if (document !== undefined) {
1188-
// Recreate the DocumentSnapshot with the DocumentReference
1189-
// containing the original converter.
1190-
const finalDoc = new DocumentSnapshotBuilder(
1191-
docRef as DocumentReference<T>
1192-
);
1193-
finalDoc.fieldsProto = document._fieldsProto;
1194-
finalDoc.readTime = document.readTime;
1195-
finalDoc.createTime = document.createTime;
1196-
finalDoc.updateTime = document.updateTime;
1197-
orderedDocuments.push(finalDoc.build());
1198-
} else {
1199-
reject(
1200-
new Error(`Did not receive document for "${docRef.path}".`)
1201-
);
1202-
}
1203-
}
1204-
resolve(orderedDocuments);
1205-
});
1206-
stream.resume();
1207-
});
1208-
}
1209-
);
1210-
}
1211-
12121089
/**
12131090
* Registers a listener on this client, incrementing the listener count. This
12141091
* is used to verify that all listeners are unsubscribed when terminate() is

0 commit comments

Comments
 (0)