Refactor stream (#8376)

* Refactor stream

* Pretty

Author: tom-andersen

packages/firestore/src/remote/persistent_stream.ts

@@ -212,6 +212,11 @@ export abstract class PersistentStream<
     this.backoff = new ExponentialBackoff(queue, connectionTimerId);
   }
 
+  /**
+   * Count of response messages received.
+   */
+  protected responseCount: number = 0;
+
   /**
    * Returns true if start() has been called and no error has occurred. True
    * indicates the stream is open or in the process of opening (which
@@ -246,6 +251,7 @@ export abstract class PersistentStream<
    * When start returns, isStarted() will return true.
    */
   start(): void {
+    this.responseCount = 0;
     if (this.state === PersistentStreamState.Error) {
       this.performBackoff();
       return;
@@ -429,11 +435,18 @@ export abstract class PersistentStream<
   ): Stream<SendType, ReceiveType>;
 
   /**
-   * Called after the stream has received a message. The function will be
-   * called on the right queue and must return a Promise.
+   * Called when the stream receives first message.
+   * The function will be called on the right queue and must return a Promise.
+   * @param message - The message received from the stream.
+   */
+  protected abstract onFirst(message: ReceiveType): Promise<void>;
+
+  /**
+   * Called on subsequent messages after the stream has received first message.
+   * The function will be called on the right queue and must return a Promise.
    * @param message - The message received from the stream.
    */
-  protected abstract onMessage(message: ReceiveType): Promise<void>;
+  protected abstract onNext(message: ReceiveType): Promise<void>;
 
   private auth(): void {
     debugAssert(
@@ -522,7 +535,11 @@ export abstract class PersistentStream<
     });
     this.stream.onMessage((msg: ReceiveType) => {
       dispatchIfNotClosed(() => {
-        return this.onMessage(msg);
+        if (++this.responseCount === 1) {
+          return this.onFirst(msg);
+        } else {
+          return this.onNext(msg);
+        }
       });
     });
   }
@@ -643,7 +660,11 @@ export class PersistentListenStream extends PersistentStream<
     );
   }
 
-  protected onMessage(watchChangeProto: ProtoListenResponse): Promise<void> {
+  protected onFirst(watchChangeProto: ProtoListenResponse): Promise<void> {
+    return this.onNext(watchChangeProto);
+  }
+
+  protected onNext(watchChangeProto: ProtoListenResponse): Promise<void> {
     // A successful response means the stream is healthy
     this.backoff.reset();
 
@@ -723,8 +744,6 @@ export class PersistentWriteStream extends PersistentStream<
   ProtoWriteResponse,
   WriteStreamListener
 > {
-  private handshakeComplete_ = false;
-
   constructor(
     queue: AsyncQueue,
     connection: Connection,
@@ -760,18 +779,17 @@ export class PersistentWriteStream extends PersistentStream<
    * the stream is ready to accept mutations.
    */
   get handshakeComplete(): boolean {
-    return this.handshakeComplete_;
+    return this.responseCount > 0;
   }
 
   // Override of PersistentStream.start
   start(): void {
-    this.handshakeComplete_ = false;
     this.lastStreamToken = undefined;
     super.start();
   }
 
   protected tearDown(): void {
-    if (this.handshakeComplete_) {
+    if (this.handshakeComplete) {
       this.writeMutations([]);
     }
   }
@@ -787,35 +805,41 @@ export class PersistentWriteStream extends PersistentStream<
     );
   }
 
-  protected onMessage(responseProto: ProtoWriteResponse): Promise<void> {
+  protected onFirst(responseProto: ProtoWriteResponse): Promise<void> {
+    // Always capture the last stream token.
+    hardAssert(
+      !!responseProto.streamToken,
+      'Got a write handshake response without a stream token'
+    );
+    this.lastStreamToken = responseProto.streamToken;
+
+    // The first response is always the handshake response
+    hardAssert(
+      !responseProto.writeResults || responseProto.writeResults.length === 0,
+      'Got mutation results for handshake'
+    );
+    return this.listener!.onHandshakeComplete();
+  }
+
+  protected onNext(responseProto: ProtoWriteResponse): Promise<void> {
     // Always capture the last stream token.
     hardAssert(
       !!responseProto.streamToken,
       'Got a write response without a stream token'
     );
     this.lastStreamToken = responseProto.streamToken;
 
-    if (!this.handshakeComplete_) {
-      // The first response is always the handshake response
-      hardAssert(
-        !responseProto.writeResults || responseProto.writeResults.length === 0,
-        'Got mutation results for handshake'
-      );
-      this.handshakeComplete_ = true;
-      return this.listener!.onHandshakeComplete();
-    } else {
-      // A successful first write response means the stream is healthy,
-      // Note, that we could consider a successful handshake healthy, however,
-      // the write itself might be causing an error we want to back off from.
-      this.backoff.reset();
+    // A successful first write response means the stream is healthy,
+    // Note, that we could consider a successful handshake healthy, however,
+    // the write itself might be causing an error we want to back off from.
+    this.backoff.reset();
 
-      const results = fromWriteResults(
-        responseProto.writeResults,
-        responseProto.commitTime
-      );
-      const commitVersion = fromVersion(responseProto.commitTime!);
-      return this.listener!.onMutationResult(commitVersion, results);
-    }
+    const results = fromWriteResults(
+      responseProto.writeResults,
+      responseProto.commitTime
+    );
+    const commitVersion = fromVersion(responseProto.commitTime!);
+    return this.listener!.onMutationResult(commitVersion, results);
   }
 
   /**
@@ -825,7 +849,7 @@ export class PersistentWriteStream extends PersistentStream<
    */
   writeHandshake(): void {
     debugAssert(this.isOpen(), 'Writing handshake requires an opened stream');
-    debugAssert(!this.handshakeComplete_, 'Handshake already completed');
+    debugAssert(!this.handshakeComplete, 'Handshake already completed');
     debugAssert(
       !this.lastStreamToken,
       'Stream token should be empty during handshake'
@@ -841,7 +865,7 @@ export class PersistentWriteStream extends PersistentStream<
   writeMutations(mutations: Mutation[]): void {
     debugAssert(this.isOpen(), 'Writing mutations requires an opened stream');
     debugAssert(
-      this.handshakeComplete_,
+      this.handshakeComplete,
       'Handshake must be complete before writing mutations'
     );
     debugAssert(