Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add Symbol.asyncInterator to Query.stream() #843

Merged
merged 6 commits into from
Dec 31, 2019
Merged

feat: add Symbol.asyncInterator to Query.stream() #843

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5df7fea
feat: add Symbol.asyncInterator to Query.stream()
schmidt-sebastian Dec 30, 2019
e264664
add dependency on readable streams
schmidt-sebastian Dec 30, 2019
2dca083
Merge branch 'master' into mrschmidt/removebun
schmidt-sebastian Dec 30, 2019
305f0a0
Review feedback
schmidt-sebastian Dec 31, 2019
935c113
Update Copyright year
schmidt-sebastian Dec 31, 2019
865981e
Update index.ts
schmidt-sebastian Dec 31, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions dev/src/external-modules.d.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,18 @@
// TODO(mrschmidt): Come up with actual definitions for these modules.
/*!
* Copyright 2019 Google Inc. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

declare module 'bun';
declare module 'functional-red-black-tree'
// TODO(mrschmidt): Come up with actual definitions for these modules.
declare module 'functional-red-black-tree'
260 changes: 78 additions & 182 deletions dev/src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
* limitations under the License.
*/

import * as bun from 'bun';
import {CallOptions} from 'google-gax';
import {Duplex, PassThrough} from 'stream';
import * as through2 from 'through2';
import {URL} from 'url';

Expand Down Expand Up @@ -950,7 +950,7 @@ export class Firestore {
const self = this;

return self
.readStream('batchGetDocuments', request, requestTag, true)
.requestStream('batchGetDocuments', 'unidirectional', request, requestTag)
.then(stream => {
return new Promise<DocumentSnapshot[]>((resolve, reject) => {
stream
Expand Down Expand Up @@ -1166,144 +1166,100 @@ export class Firestore {
}

/**
* Opens the provided stream and waits for it to become healthy. If an error
* occurs before the first byte is read, the method rejects the returned
* Promise.
* Waits for the provided stream to become active and returns a paused but
* healthy stream. If an error occurs before the first byte is read, the
* method rejects the returned Promise.
*
* @private
* @param resultStream The Node stream to monitor.
* @param backendStream The Node stream to monitor.
* @param requestTag A unique client-assigned identifier for this request.
* @param request If specified, the request that should be written to the
* stream after it opened.
* @returns The given Stream once it is considered healthy.
* stream after opening.
* @returns A guaranteed healthy stream that should be used instead of
* `backendStream`.
*/
private _initializeStream(
resultStream: NodeJS.ReadableStream,
requestTag: string
): Promise<void>;
private _initializeStream(
resultStream: NodeJS.ReadWriteStream,
requestTag: string,
request: {}
): Promise<void>;
private _initializeStream(
resultStream: NodeJS.ReadableStream | NodeJS.ReadWriteStream,
backendStream: Duplex,
requestTag: string,
request?: {}
): Promise<void> {
/** The last error we received and have not forwarded yet. */
let errorReceived: Error | null = null;
): Promise<Duplex> {
const resultStream = new PassThrough({objectMode: true});
resultStream.pause();

/**
* Whether we have resolved the Promise and returned the stream to the
* caller.
*/
let streamInitialized = false;

/**
* Whether the stream end has been reached. This has to be forwarded to the
* caller..
*/
let endCalled = false;

return new Promise((resolve, reject) => {
const streamReady = () => {
if (errorReceived) {
logger(
'Firestore._initializeStream',
requestTag,
'Emit error:',
errorReceived
);
resultStream.emit('error', errorReceived);
errorReceived = null;
} else if (!streamInitialized) {
logger('Firestore._initializeStream', requestTag, 'Releasing stream');
return new Promise<Duplex>((resolve, reject) => {
function streamReady() {
if (!streamInitialized) {
streamInitialized = true;
resultStream.pause();

// Calling 'stream.pause()' only holds up 'data' events and not the
// 'end' event we intend to forward here. We therefore need to wait
// until the API consumer registers their listeners (in the .then()
// call) before emitting any further events.
resolve();

// We execute the forwarding of the 'end' event via setTimeout() as
// V8 guarantees that the above the Promise chain is resolved before
// any calls invoked via setTimeout().
setTimeout(() => {
if (endCalled) {
logger(
'Firestore._initializeStream',
requestTag,
'Forwarding stream close'
);
resultStream.emit('end');
}
}, 0);
logger('Firestore._initializeStream', requestTag, 'Releasing stream');
resolve(resultStream);
}
};

// We capture any errors received and buffer them until the caller has
// registered a listener. We register our event handler as early as
// possible to avoid the default stream behavior (which is just to log and
// continue).
resultStream.on('readable', () => {
streamReady();
});
}

resultStream.on('end', () => {
function streamEnded() {
logger(
'Firestore._initializeStream',
requestTag,
'Received stream end'
);
endCalled = true;
streamReady();
});
resultStream.unpipe(backendStream);
}

resultStream.on('error', err => {
logger(
'Firestore._initializeStream',
requestTag,
'Received stream error:',
err
);
// If we receive an error before we were able to receive any data,
// reject this stream.
backendStream.on('data', () => streamReady());
backendStream.on('end', () => streamEnded());
backendStream.on('close', () => streamEnded());

backendStream.on('error', err => {
if (!streamInitialized) {
// If we receive an error before we were able to receive any data,
// reject this stream.
logger(
'Firestore._initializeStream',
requestTag,
'Received initial error:',
err
);
streamInitialized = true;
reject(err);
} else {
errorReceived = err;
logger(
'Firestore._initializeStream',
requestTag,
'Received stream error:',
err
);
// We execute the forwarding of the 'error' event via setImmediate() as
// V8 guarantees that the Promise chain returned from this method
// is resolved before any code executed via setImmediate(). This
// allows the caller to attach an error handler.
setImmediate(() => {
resultStream.emit('error', err);
}, 0);
}
});

backendStream.pipe(resultStream);

if (request) {
logger(
'Firestore._initializeStream',
requestTag,
'Sending request: %j',
request
);
(resultStream as NodeJS.WritableStream)
// The stream returned by the Gapic library accepts Protobuf
// messages, but the type information does not expose this.
// tslint:disable-next-line no-any
.write(request as any, 'utf-8', () => {
logger(
'Firestore._initializeStream',
requestTag,
'Marking stream as healthy'
);
streamReady();
});
backendStream.write(request, 'utf-8', () => {
logger(
'Firestore._initializeStream',
requestTag,
'Marking stream as healthy'
);
streamReady();
});
}
});
}
Expand Down Expand Up @@ -1363,130 +1319,70 @@ export class Firestore {
}

/**
* A funnel for read-only streaming API requests, assigning a project ID where
* necessary within the request options.
* A funnel for streaming API requests, assigning a project ID where necessary
* within the request options.
*
* The stream is returned in paused state and needs to be resumed once all
* listeners are attached.
*
* @private
* @param methodName Name of the streaming Veneer API endpoint that
* takes a request and GAX options.
* @param mode Whether this a unidirectional or bidirectional call.
* @param request The Protobuf request to send.
* @param requestTag A unique client-assigned identifier for this request.
* @param allowRetries Whether this is an idempotent request that can be
* retried.
* @returns A Promise with the resulting read-only stream.
*/
readStream(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it a breaking change?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not part of of our public API (we define our own types). I can use an underscore prefix here, but in general the client is pretty inconsistent here.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK. Since it's TypeScript, we might as well make those methods private.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unfortunately, these are "module private" and used in other internal classes.

requestStream(
methodName: string,
mode: 'unidirectional' | 'bidirectional',
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if it's an externally visible function, I would vote for an enum (but up to you)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be private.

request: {},
requestTag: string,
allowRetries: boolean
): Promise<NodeJS.ReadableStream> {
const attempts = allowRetries ? MAX_REQUEST_RETRIES : 1;
requestTag: string
): Promise<Duplex> {
const callOptions = this.createCallOptions();

const result = new Deferred<NodeJS.ReadableStream>();
const result = new Deferred<Duplex>();

this._clientPool.run(requestTag, gapicClient => {
// While we return the stream to the callee early, we don't want to
// release the GAPIC client until the callee has finished processing the
// stream.
const lifetime = new Deferred<void>();

this._retry(attempts, requestTag, async () => {
this._retry(MAX_REQUEST_RETRIES, requestTag, async () => {
logger(
'Firestore.readStream',
'Firestore.requestStream',
requestTag,
'Sending request: %j',
request
);
const stream = gapicClient[methodName](request, callOptions);
const stream: Duplex =
mode === 'unidirectional'
? gapicClient[methodName](request, callOptions)
: gapicClient[methodName](callOptions);
const logStream = through2.obj(function(this, chunk, enc, callback) {
logger(
'Firestore.readStream',
'Firestore.requestStream',
requestTag,
'Received response: %j',
chunk
);
this.push(chunk);
callback();
});

const resultStream = bun([stream, logStream]);
resultStream.on('close', lifetime.resolve);
resultStream.on('end', lifetime.resolve);
resultStream.on('error', lifetime.resolve);

await this._initializeStream(resultStream, requestTag);
result.resolve(resultStream);
}).catch(err => {
lifetime.resolve();
result.reject(err);
});

return lifetime.promise;
});

return result.promise;
}
stream.pipe(logStream);
stream.on('close', lifetime.resolve);
stream.on('end', lifetime.resolve);
stream.on('finish', lifetime.resolve);
stream.on('error', lifetime.resolve);

/**
* A funnel for read-write streaming API requests, assigning a project ID
* where necessary for all writes.
*
* The stream is returned in paused state and needs to be resumed once all
* listeners are attached.
*
* @private
* @param methodName Name of the streaming Veneer API endpoint that takes
* GAX options.
* @param request The Protobuf request to send as the first stream message.
* @param requestTag A unique client-assigned identifier for this request.
* @param allowRetries Whether this is an idempotent request that can be
* retried.
* @returns A Promise with the resulting read/write stream.
*/
readWriteStream(
methodName: string,
request: {},
requestTag: string,
allowRetries: boolean
): Promise<NodeJS.ReadWriteStream> {
const attempts = allowRetries ? MAX_REQUEST_RETRIES : 1;
const callOptions = this.createCallOptions();

const result = new Deferred<NodeJS.ReadWriteStream>();

this._clientPool.run(requestTag, gapicClient => {
// While we return the stream to the callee early, we don't want to
// release the GAPIC client until the callee has finished processing the
// stream.
const lifetime = new Deferred<void>();

this._retry(attempts, requestTag, async () => {
logger('Firestore.readWriteStream', requestTag, 'Opening stream');
const requestStream = gapicClient[methodName](callOptions);

const logStream = through2.obj(function(this, chunk, enc, callback) {
logger(
'Firestore.readWriteStream',
requestTag,
'Received response: %j',
chunk
);
this.push(chunk);
callback();
});

const resultStream = bun([requestStream, logStream]);
resultStream.on('close', lifetime.resolve);
resultStream.on('finish', lifetime.resolve);
resultStream.on('end', lifetime.resolve);
resultStream.on('error', lifetime.resolve);
const resultStream = await this._initializeStream(
stream,
requestTag,
mode === 'bidirectional' ? request : undefined
);

await this._initializeStream(resultStream, requestTag, request);
resultStream.on('end', () => stream.end());
result.resolve(resultStream);
}).catch(err => {
lifetime.resolve();
Expand Down
Loading