Improve back-pressure control in Result[Symbol.asyncIterator] (#864)

The previous implemenation only controls the flow while is iterating over the records, but if the iteration
is done in a really low pace, the control flow method won't be called in time for pausing the flow.

Changing the flow control for observing the queue size instead of check on each time it iterates solve this issue.

Author: bigmontz

packages/core/src/result.ts

@@ -241,19 +241,28 @@ class Result implements Promise<QueryResult> {
       summary?: ResultSummary,
     } = { paused: true, firstRun: true, finished: false }
 
-
-    const controlFlow = async  () => {
-      if (state.queuedObserver === undefined) {
-        state.queuedObserver = this._createQueuedResultObserver()
-        state.streaming = await this._subscribe(state.queuedObserver, true).catch(() => undefined)
+    const controlFlow = () => {
+      if (!state.streaming) {
+        return;
       }
-      if (state.queuedObserver.size >= this._watermarks.high && !state.paused) {
+      const queueSizeIsOverHighOrEqualWatermark = state.queuedObserver!.size >= this._watermarks.high 
+      const queueSizeIsBellowOrEqualLowWatermark = state.queuedObserver!.size <= this._watermarks.low
+
+      if (queueSizeIsOverHighOrEqualWatermark && !state.paused) {
         state.paused = true
-        state.streaming?.pause()
-      } else if (state.queuedObserver.size <= this._watermarks.low && state.paused || state.firstRun) {
+        state.streaming.pause()
+      } else if (queueSizeIsBellowOrEqualLowWatermark && state.paused || state.firstRun && !queueSizeIsOverHighOrEqualWatermark ) {
         state.firstRun = false
         state.paused = false
-        state.streaming?.resume()
+        state.streaming.resume()
+      }
+    }
+
+    const initializeObserver = async  () => {
+      if (state.queuedObserver === undefined) {
+        state.queuedObserver = this._createQueuedResultObserver(controlFlow)
+        state.streaming = await this._subscribe(state.queuedObserver, true).catch(() => undefined)
+        controlFlow()
       }
     }
 
@@ -262,7 +271,7 @@ class Result implements Promise<QueryResult> {
         if (state.finished) {
           return { done: true, value: state.summary! }
         }
-        await controlFlow()
+        await initializeObserver()
         const next = await state.queuedObserver!.dequeue()
         if (next.done) {
           state.finished = next.done
@@ -280,7 +289,7 @@ class Result implements Promise<QueryResult> {
         if (state.finished) {
           return { done: true, value: state.summary! }
         }
-        await controlFlow()
+        await initializeObserver()
         return await state.queuedObserver!.head()
       }
     }
@@ -461,7 +470,7 @@ class Result implements Promise<QueryResult> {
   /**
    * @access private
    */
-   private _createQueuedResultObserver (): QueuedResultObserver {
+   private _createQueuedResultObserver (onQueueSizeChanged: () => void): QueuedResultObserver {
     interface ResolvablePromise<T> {
       promise: Promise<T>
       resolve: (arg: T) => any | undefined
@@ -509,11 +518,13 @@ class Result implements Promise<QueryResult> {
           }
         } else {
           buffer.push(element)
+          onQueueSizeChanged()
         }
       },
       dequeue: async () => {
         if (buffer.length > 0) {
           const element = buffer.shift()!
+          onQueueSizeChanged()
           if (isError(element)) {
               throw element
           }
@@ -538,6 +549,8 @@ class Result implements Promise<QueryResult> {
         } catch (error) {
           buffer.unshift(error)
           throw error
+        } finally {
+          onQueueSizeChanged()
         }
       },
       get size (): number {

packages/core/test/result.test.ts

@@ -598,16 +598,39 @@ describe('Result', () => {
 
       it('should pause the stream if queue is bigger than high watermark', async () => {
         const pause = jest.spyOn(streamObserverMock, 'pause')
+        const resume = jest.spyOn(streamObserverMock, 'resume')
         streamObserverMock.onKeys(['a'])
 
-        for (let i = 0; i <= watermarks.high; i++) {
+        for (let i = 0; i < watermarks.high + 3; i++) {
           streamObserverMock.onNext([i])
         }
 
         const it = result[Symbol.asyncIterator]()
         await it.next()
 
         expect(pause).toBeCalledTimes(1)
+        expect(resume).toBeCalledTimes(0)
+      })
+
+      it('should pause the stream if queue is bigger than high watermark and not iteraction with the stream', async () => {
+        const pause = jest.spyOn(streamObserverMock, 'pause')
+        const resume = jest.spyOn(streamObserverMock, 'resume')
+        streamObserverMock.onKeys(['a'])
+
+        streamObserverMock.onNext([-1])
+
+        const it = result[Symbol.asyncIterator]()
+        await it.next()
+
+        expect(pause).toBeCalledTimes(1)
+        expect(resume).toBeCalledTimes(1)
+
+        for (let i = 0; i <= watermarks.high + 1; i++) {
+          streamObserverMock.onNext([i])
+        }
+
+        expect(pause).toBeCalledTimes(2)
+        expect(resume).toBeCalledTimes(1)
       })
 
       it('should call resume if queue is smaller than low watermark', async () => {