Handle abrupt closing in mapAsyncIterator() (#870)

This adds proper behavior when the mapping function throws an error (As discussed in #868).

src/subscription/__tests__/mapAsyncIterator-test.js

@@ -5,6 +5,8 @@
  *  This source code is licensed under the BSD-style license found in the
  *  LICENSE file in the root directory of this source tree. An additional grant
  *  of patent rights can be found in the PATENTS file in the same directory.
+ *
+ * @flow
  */
 
 import { expect } from 'chai';
@@ -43,7 +45,9 @@ describe('mapAsyncIterator', () => {
       yield 3;
     }
 
-    const doubles = mapAsyncIterator(source(), async x => await x + x);
+    // Flow test: this is *not* AsyncIterator<Promise<number>>
+    const doubles: AsyncIterator<number> =
+      mapAsyncIterator(source(), async x => await x + x);
 
     expect(
       await doubles.next()
@@ -190,4 +194,61 @@ describe('mapAsyncIterator', () => {
     ).to.deep.equal({ value: undefined, done: true });
   });
 
+  async function testClosesSourceWithMapper(mapper) {
+    let didVisitFinally = false;
+
+    async function* source() {
+      try {
+        yield 1;
+        yield 2;
+        yield 3;
+      } finally {
+        didVisitFinally = true;
+        yield 1000;
+      }
+    }
+
+    const throwOver1 = mapAsyncIterator(source(), mapper);
+
+    expect(
+      await throwOver1.next()
+    ).to.deep.equal({ value: 1, done: false });
+
+    let expectedError;
+    try {
+      await throwOver1.next();
+    } catch (error) {
+      expectedError = error;
+    }
+
+    expect(expectedError).to.be.an('error');
+    if (expectedError) {
+      expect(expectedError.message).to.equal('Cannot count to 2');
+    }
+
+    expect(
+      await throwOver1.next()
+    ).to.deep.equal({ value: undefined, done: true });
+
+    expect(didVisitFinally).to.equal(true);
+  }
+
+  it('closes source if mapper throws an error', async () => {
+    await testClosesSourceWithMapper(x => {
+      if (x > 1) {
+        throw new Error('Cannot count to ' + x);
+      }
+      return x;
+    });
+  });
+
+  it('closes source if mapper rejects', async () => {
+    await testClosesSourceWithMapper(async x => {
+      if (x > 1) {
+        throw new Error('Cannot count to ' + x);
+      }
+      return x;
+    });
+  });
+
 });

src/subscription/mapAsyncIterator.js

@@ -17,38 +17,53 @@ import { $$asyncIterator, getAsyncIterator } from 'iterall';
  */
 export default function mapAsyncIterator<T, U>(
   iterable: AsyncIterable<T>,
-  callback: (value: T) => U
-): AsyncIterator<U> {
-  // Fixes a temporary issue with Regenerator/Babel
-  // https://github.com/facebook/regenerator/pull/290
-  const iterator = iterable.next ? (iterable: any) : getAsyncIterator(iterable);
+  callback: (value: T) => Promise<U> | U
+): AsyncGenerator<U, void, void> {
+  const iterator = getAsyncIterator(iterable);
+  let $return;
+  let abruptClose;
+  if (typeof iterator.return === 'function') {
+    $return = iterator.return;
+    abruptClose = error => {
+      const rethrow = () => Promise.reject(error);
+      return $return.call(iterator).then(rethrow, rethrow);
+    };
+  }
 
   function mapResult(result) {
     return result.done ?
       result :
-      Promise.resolve(callback(result.value)).then(
-        mapped => ({ value: mapped, done: false })
-      );
+      asyncMapValue(result.value, callback).then(iteratorResult, abruptClose);
   }
 
   return {
     next() {
       return iterator.next().then(mapResult);
     },
     return() {
-      if (typeof iterator.return === 'function') {
-        return iterator.return().then(mapResult);
-      }
-      return Promise.resolve({ value: undefined, done: true });
+      return $return ?
+        $return.call(iterator).then(mapResult) :
+        Promise.resolve({ value: undefined, done: true });
     },
     throw(error) {
       if (typeof iterator.throw === 'function') {
         return iterator.throw(error).then(mapResult);
       }
-      return Promise.reject(error);
+      return Promise.reject(error).catch(abruptClose);
     },
     [$$asyncIterator]() {
       return this;
     },
   };
 }
+
+function asyncMapValue<T, U>(
+  value: T,
+  callback: (T) => Promise<U> | U
+): Promise<U> {
+  return new Promise(resolve => resolve(callback(value)));
+}
+
+function iteratorResult<T>(value: T): IteratorResult<T, void> {
+  return { value, done: false };
+}