Fix case where pending items flowing will not be decremented after an error

Author: Pedro Kehl

package-lock.json

@@ -60,6 +60,7 @@
   "license": "ISC",
   "devDependencies": {
     "@types/jest": "^29.5.12",
+    "@types/node": "^20.16.0",
     "@typescript-eslint/eslint-plugin": "^6.21.0",
     "@typescript-eslint/parser": "^6.21.0",
     "eslint": "^8.57.0",
@@ -72,4 +73,4 @@
   "dependencies": {
     "rxjs": "^7.8.1"
   }
-}
+}

package.json

@@ -10,16 +10,23 @@ import { parallel } from './operators/parallel'
 import { reduce, type ReduceParams } from './operators/reduce'
 import { filter, type FilterPredicate } from './operators/filter'
 
-import { applyOperator, isBatch, type OperatorApplier } from './operators/helpers/operatorHelpers'
+import {
+  applyOperator,
+  isBatch,
+  type OperatorApplier,
+  type OperatorApplierWithRunId,
+} from './operators/helpers/operatorHelpers'
 import { type PendingDataControl, PendingDataControlInMemory } from './utils/PendingDataControl'
 
 import { getOnStepFinished } from './utils/onStepFinished'
 import { getOnStepStarted } from './utils/onStepStarted'
+import { generateId } from './utils/generateId'
+
+type Generator = (initialBag: ValueBag, runId: string) => AsyncGenerator<ValueBag>
 
 export class Caminho implements CaminhoInterface {
-  private generator: (initialBag: ValueBag) => AsyncGenerator<ValueBag>
-  private operators: OperatorApplier[] = []
-  private finalStep?: OperatorApplier
+  private generator: Generator
+  private operators: OperatorApplierWithRunId[] = []
   private pendingDataControl?: PendingDataControl
 
   constructor(generatorParams: FromGeneratorParams, private options?: CaminhoOptions) {
@@ -40,14 +47,14 @@ export class Caminho implements CaminhoInterface {
 
   public pipe(params: PipeGenericParams): this {
     const operatorApplier = this.getApplierForPipeOrBatch(params)
-    this.addOperatorApplier(operatorApplier)
+    this.addOperatorApplier(() => operatorApplier)
     return this
   }
 
   public parallel(params: PipeGenericParams[]): this {
     const operatorAppliers: OperatorApplier[] = params.map(this.getApplierForPipeOrBatch)
     const operatorApplier = parallel(params, operatorAppliers)
-    this.addOperatorApplier(operatorApplier)
+    this.addOperatorApplier(() => operatorApplier)
     return this
   }
 
@@ -64,34 +71,32 @@ export class Caminho implements CaminhoInterface {
   }
 
   public async run(initialBag?: ValueBag): Promise<ValueBag> {
-    const observable$ = this.buildObservable(initialBag)
-    return lastValueFrom(observable$, { defaultValue: initialBag })
+    const runId = generateId()
+    const initial$ = from(this.generator({ ...initialBag }, runId))
+    const observable$ = this.operators.reduce((acc, operator) => applyOperator(acc, operator, runId), initial$)
+
+    const finalObservable$ = this.options?.maxItemsFlowing
+      ? observable$.pipe(tap(() => (this.pendingDataControl as PendingDataControl).decrement(runId)))
+      : observable$
+
+    try {
+      return await lastValueFrom(finalObservable$, { defaultValue: initialBag })
+    } finally {
+      this.pendingDataControl?.destroyBucket(runId)
+    }
   }
 
-  private getGenerator(generatorParams: FromGeneratorParams): (initialBag: ValueBag) => AsyncGenerator<ValueBag> {
+  private getGenerator(generatorParams: FromGeneratorParams): Generator {
     const loggers = this.getLoggers(generatorParams)
     if (this.options?.maxItemsFlowing) {
       const pendingDataControl = this.pendingDataControl as PendingDataControl
-      this.finalStep = tap({ next: () => pendingDataControl.decrement(), error: () => pendingDataControl.decrement() })
       return wrapGeneratorWithBackPressure(generatorParams, this.options.maxItemsFlowing, pendingDataControl, loggers)
     }
 
     return wrapGenerator(generatorParams, loggers)
   }
 
-  private buildObservable(initialBag: ValueBag = {}) {
-    const initialObservable$ = from(this.generator({ ...initialBag }))
-    const observable$ = this.operators.reduce(applyOperator, initialObservable$)
-
-    if (this.finalStep) {
-      return observable$
-        .pipe(this.finalStep)
-    }
-
-    return observable$
-  }
-
-  private addOperatorApplier(operatorApplier: OperatorApplier) {
+  private addOperatorApplier(operatorApplier: OperatorApplierWithRunId) {
     this.operators.push(operatorApplier)
   }
 

src/Caminho.ts

@@ -1,22 +1,22 @@
 import { filter as filterRxJs } from 'rxjs'
 import { type PendingDataControl } from '../utils/PendingDataControl'
 import { type Loggers, type ValueBag } from '../types'
-import { type OperatorApplier } from './helpers/operatorHelpers'
+import { type OperatorApplierWithRunId } from './helpers/operatorHelpers'
 
 export type FilterPredicate = (valueBag: ValueBag, index: number) => boolean
 
 export function filter(
   predicate: FilterPredicate,
   loggers: Loggers,
   pendingDataControl?: PendingDataControl,
-): OperatorApplier {
-  function wrappedFilter(valueBag: ValueBag, index: number): boolean {
+): OperatorApplierWithRunId {
+  function wrappedFilter(valueBag: ValueBag, index: number, runId: string): boolean {
     const startedAt = new Date()
     loggers.onStepStarted([valueBag])
     try {
       const filterResult = predicate(valueBag, index)
       if (!filterResult) {
-        pendingDataControl?.decrement()
+        pendingDataControl?.decrement(runId)
       }
       loggers.onStepFinished([valueBag], startedAt)
       return filterResult
@@ -26,5 +26,7 @@ export function filter(
     }
   }
 
-  return filterRxJs(wrappedFilter)
+  return function filterOperatorWithRunId(runId: string) {
+    return filterRxJs((valueBag, index) => wrappedFilter(valueBag, index, runId))
+  }
 }

src/operators/filter.ts

@@ -38,9 +38,9 @@ export function wrapGeneratorWithBackPressure(
   loggers: Loggers,
 ) {
   const wrappedGenerator = wrapGenerator(generatorParams, loggers)
-  return async function* wrappedGeneratorWithBackPressure(initialBag: ValueBag) {
-    for await (const value of wrappedGenerator(initialBag)) {
-      pendingDataControl.increment()
+  return async function* wrappedGeneratorWithBackPressure(initialBag: ValueBag, runId: string) {
+    for await (const value of wrappedGenerator({ ...initialBag })) {
+      pendingDataControl.increment(runId)
       yield value
       if (needsToWaitForBackpressure(pendingDataControl, maxItemsFlowing)) {
         await waitOnBackpressure(maxItemsFlowing, pendingDataControl)

src/operators/generator.ts

@@ -5,14 +5,16 @@ import { type BatchParams } from '../batch'
 import { type PipeParams } from '../pipe'
 
 export type OperatorApplier = (observable: Observable<ValueBag>) => Observable<ValueBag>
+export type OperatorApplierWithRunId = (runId: string) => OperatorApplier
 
 export function isBatch(params: PipeParams | BatchParams): params is BatchParams {
   return !!(params as BatchParams)?.batch
 }
 
 export function applyOperator(
   observable: Observable<ValueBag>,
-  operatorApplier: OperatorApplier,
+  operatorApplier: OperatorApplierWithRunId,
+  runId: string,
 ): Observable<ValueBag> {
-  return operatorApplier(observable)
+  return operatorApplier(runId)(observable)
 }

src/operators/helpers/operatorHelpers.ts

@@ -1,7 +1,7 @@
 import { type Observable, map, reduce as reduceRxJs } from 'rxjs'
 import type { Loggers, ValueBag } from '../types'
 import { type PendingDataControl } from '../utils/PendingDataControl'
-import { type OperatorApplier } from './helpers/operatorHelpers'
+import { type OperatorApplierWithRunId } from './helpers/operatorHelpers'
 import { getNewValueBag } from '../utils/valueBag'
 import { pick } from '../utils/pick'
 
@@ -18,12 +18,12 @@ export function reduce<T>(
   reduceParams: ReduceParams<T>,
   loggers: Loggers,
   pendingDataControl?: PendingDataControl,
-): OperatorApplier {
+): OperatorApplierWithRunId {
   const { provides, keep, seed } = reduceParams
   const immutableSeed = Object.freeze(seed)
   let lastBag: ValueBag = {}
 
-  function wrappedReduce(acc: T, valueBag: ValueBag, index: number): T {
+  function wrappedReduce(acc: T, valueBag: ValueBag, index: number, runId: string): T {
     const startedAt = new Date()
     loggers.onStepStarted([valueBag])
     // RxJs doesn't create a structureClone from the seed parameter when start processing.
@@ -39,16 +39,18 @@ export function reduce<T>(
       loggers.onStepFinished([valueBag], startedAt, err as Error)
       throw err
     } finally {
-      pendingDataControl?.decrement()
+      pendingDataControl?.decrement(runId)
     }
   }
 
-  return function operatorApplier(observable: Observable<ValueBag>) {
-    return observable
-      .pipe(reduceRxJs(wrappedReduce, seed))
-      .pipe(map((reduceResult: T) => {
-        pendingDataControl?.increment()
-        return getNewValueBag(pick(lastBag, keep ?? []), provides, reduceResult)
-      }))
+  return function operatorApplierWithRunId(runId: string) {
+    return function operatorApplier(observable: Observable<ValueBag>) {
+      return observable
+        .pipe(reduceRxJs((acc, valueBag, index) => wrappedReduce(acc, valueBag, index, runId), seed))
+        .pipe(map((reduceResult: T) => {
+          pendingDataControl?.increment(runId)
+          return getNewValueBag(pick(lastBag, keep ?? []), provides, reduceResult)
+        }))
+    }
   }
 }

src/operators/reduce.ts

@@ -1,17 +1,29 @@
 export class PendingDataControlInMemory implements PendingDataControl {
   public size = 0
+  private buckets = new Map()
 
-  increment(value = 1): void {
+  increment(bucketId: string, value = 1): void {
     this.size += value
+    const current = this.buckets.get(bucketId) ?? 0
+    this.buckets.set(bucketId, current + value)
   }
 
-  decrement(value = 1): void {
+  decrement(bucketId: string, value = 1): void {
     this.size -= value
+    const current = this.buckets.get(bucketId)
+    this.buckets.set(bucketId, current - value)
+  }
+
+  destroyBucket(bucketId: string) {
+    const inBucket = this.buckets.get(bucketId) ?? 0
+    this.size -= inBucket
+    this.buckets.delete(bucketId)
   }
 }
 
 export type PendingDataControl = {
   size: number
-  increment: (value?: number) => void
-  decrement: (value?: number) => void
+  increment: (bucketId: string, value?: number) => void
+  decrement: (bucketId: string, value?: number) => void
+  destroyBucket: (bucketId: string) => void
 }

src/utils/PendingDataControl.ts

@@ -0,0 +1,3 @@
+export function generateId() {
+  return `${Date.now().toString(36)}-${Math.random().toString(36).substr(2, 9)}`
+}

src/utils/generateId.ts

@@ -1,4 +1,6 @@
 import { fromGenerator } from '../../src'
+import { sleep } from '../../src/utils/sleep'
+import { getNumberedArray } from '../mocks/array.mock'
 import { getMockedGenerator, getThrowingGenerator } from '../mocks/generator.mock'
 
 describe('Error Handling', () => {
@@ -65,4 +67,20 @@ describe('Error Handling', () => {
     await expect(caminho.run()).rejects.toMatchObject({ message: 'Operator error' })
     expect(caminho.getNumberOfItemsFlowing()).toBe(0)
   })
+
+  test('Should not keep account for the pending items in the flow after an error', async () => {
+    const operator = jest.fn()
+      .mockImplementationOnce(() => sleep(1))
+      .mockImplementationOnce(() => sleep(1))
+      .mockImplementationOnce(() => sleep(1))
+      .mockRejectedValueOnce(new Error('Operator error'))
+
+    const generator = getMockedGenerator(getNumberedArray(10))
+
+    const caminho = fromGenerator({ fn: generator, provides: 'number' }, { maxItemsFlowing: 5 })
+      .pipe({ fn: operator })
+
+    await expect(caminho.run()).rejects.toMatchObject({ message: 'Operator error' })
+    expect(caminho.getNumberOfItemsFlowing()).toBe(0)
+  })
 })