Skip to content

Commit b39dd3c

Browse files
schmidt-sebastianschmidt-sebastian
authored
feat: retry BatchGetDocuments RPCs that fail with errors (#1544)
1 parent 7473343 commit b39dd3c

File tree

4 files changed

+245
-187
lines changed

4 files changed

+245
-187
lines changed

‎dev/src/document-reader.ts

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
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+
import {DocumentData} from '@google-cloud/firestore';
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 outstandingDocuments = new Set<string>();
40+
private retrievedDocuments = new Map<string, DocumentSnapshot>();
41+
42+
/**
43+
* Creates a new DocumentReader that fetches the provided documents (via
44+
* `get()`).
45+
*
46+
* @param firestore The Firestore instance to use.
47+
* @param allDocuments The documents to get.
48+
*/
49+
constructor(
50+
private firestore: Firestore,
51+
private allDocuments: Array<DocumentReference<T>>
52+
) {
53+
for (const docRef of this.allDocuments) {
54+
this.outstandingDocuments.add(docRef.formattedName);
55+
}
56+
}
57+
58+
/**
59+
* Invokes the BatchGetDocuments RPC and returns the results.
60+
*
61+
* @param requestTag A unique client-assigned identifier for this request.
62+
*/
63+
async get(requestTag: string): Promise<Array<DocumentSnapshot<T>>> {
64+
await this.fetchDocuments(requestTag);
65+
66+
// BatchGetDocuments doesn't preserve document order. We use the request
67+
// order to sort the resulting documents.
68+
const orderedDocuments: Array<DocumentSnapshot<T>> = [];
69+
70+
for (const docRef of this.allDocuments) {
71+
const document = this.retrievedDocuments.get(docRef.formattedName);
72+
if (document !== undefined) {
73+
// Recreate the DocumentSnapshot with the DocumentReference
74+
// containing the original converter.
75+
const finalDoc = new DocumentSnapshotBuilder(
76+
docRef as DocumentReference<T>
77+
);
78+
finalDoc.fieldsProto = document._fieldsProto;
79+
finalDoc.readTime = document.readTime;
80+
finalDoc.createTime = document.createTime;
81+
finalDoc.updateTime = document.updateTime;
82+
orderedDocuments.push(finalDoc.build());
83+
} else {
84+
throw new Error(`Did not receive document for "${docRef.path}".`);
85+
}
86+
}
87+
88+
return orderedDocuments;
89+
}
90+
91+
private async fetchDocuments(requestTag: string): Promise<void> {
92+
if (!this.outstandingDocuments.size) {
93+
return;
94+
}
95+
96+
const request: api.IBatchGetDocumentsRequest = {
97+
database: this.firestore.formattedName,
98+
transaction: this.transactionId,
99+
documents: Array.from(this.outstandingDocuments),
100+
};
101+
102+
if (this.fieldMask) {
103+
const fieldPaths = this.fieldMask.map(
104+
fieldPath => fieldPath.formattedName
105+
);
106+
request.mask = {fieldPaths};
107+
}
108+
109+
let resultCount = 0;
110+
111+
try {
112+
const stream = await this.firestore.requestStream(
113+
'batchGetDocuments',
114+
request,
115+
requestTag
116+
);
117+
stream.resume();
118+
119+
for await (const response of stream) {
120+
let snapshot: DocumentSnapshot<DocumentData>;
121+
122+
if (response.found) {
123+
logger(
124+
'DocumentReader.fetchDocuments',
125+
requestTag,
126+
'Received document: %s',
127+
response.found.name!
128+
);
129+
snapshot = this.firestore.snapshot_(
130+
response.found,
131+
response.readTime!
132+
);
133+
} else {
134+
logger(
135+
'DocumentReader.fetchDocuments',
136+
requestTag,
137+
'Document missing: %s',
138+
response.missing!
139+
);
140+
snapshot = this.firestore.snapshot_(
141+
response.missing!,
142+
response.readTime!
143+
);
144+
}
145+
146+
const path = snapshot.ref.formattedName;
147+
this.outstandingDocuments.delete(path);
148+
this.retrievedDocuments.set(path, snapshot);
149+
++resultCount;
150+
}
151+
} catch (error) {
152+
const shouldRetry =
153+
// Transactional reads are retried via the transaction runner.
154+
!this.transactionId &&
155+
// Only retry if we made progress.
156+
resultCount > 0 &&
157+
// Don't retry permanent errors.
158+
error.code !== undefined &&
159+
!isPermanentRpcError(error, 'batchGetDocuments');
160+
161+
logger(
162+
'DocumentReader.fetchDocuments',
163+
requestTag,
164+
'BatchGetDocuments failed with error: %s. Retrying: %s',
165+
error,
166+
shouldRetry
167+
);
168+
if (shouldRetry) {
169+
return this.fetchDocuments(requestTag);
170+
} else {
171+
throw error;
172+
}
173+
} finally {
174+
logger(
175+
'DocumentReader.fetchDocuments',
176+
requestTag,
177+
'Received %d results',
178+
resultCount
179+
);
180+
}
181+
}
182+
}

‎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';
@@ -1080,138 +1079,16 @@ export class Firestore implements firestore.Firestore {
10801079
const stack = Error().stack!;
10811080

10821081
return this.initializeIfNeeded(tag)
1083-
.then(() => this.getAll_(documents, fieldMask, tag))
1082+
.then(() => {
1083+
const reader = new DocumentReader(this, documents);
1084+
reader.fieldMask = fieldMask || undefined;
1085+
return reader.get(tag);
1086+
})
10841087
.catch(err => {
10851088
throw wrapError(err, stack);
10861089
});
10871090
}
10881091

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

‎dev/src/transaction.ts

Lines changed: 8 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ import {
3838
validateMinNumberOfArguments,
3939
validateOptional,
4040
} from './validate';
41-
41+
import {DocumentReader} from './document-reader';
4242
import api = proto.google.firestore.v1;
4343

4444
/*!
@@ -125,16 +125,9 @@ export class Transaction implements firestore.Transaction {
125125
}
126126

127127
if (refOrQuery instanceof DocumentReference) {
128-
return this._firestore
129-
.getAll_(
130-
[refOrQuery],
131-
/* fieldMask= */ null,
132-
this._requestTag,
133-
this._transactionId
134-
)
135-
.then(res => {
136-
return Promise.resolve(res[0]);
137-
});
128+
const documentReader = new DocumentReader(this._firestore, [refOrQuery]);
129+
documentReader.transactionId = this._transactionId;
130+
return documentReader.get(this._requestTag).then(([res]) => res);
138131
}
139132

140133
if (refOrQuery instanceof Query) {
@@ -191,12 +184,10 @@ export class Transaction implements firestore.Transaction {
191184
documentRefsOrReadOptions
192185
);
193186

194-
return this._firestore.getAll_(
195-
documents,
196-
fieldMask,
197-
this._requestTag,
198-
this._transactionId
199-
);
187+
const documentReader = new DocumentReader(this._firestore, documents);
188+
documentReader.fieldMask = fieldMask || undefined;
189+
documentReader.transactionId = this._transactionId;
190+
return documentReader.get(this._requestTag);
200191
}
201192

202193
/**

0 commit comments

Comments
 (0)