Add a helper for chaining async iterators. (#13528)

This will be used in upcoming locator-related changes.

Author: ericsnowcurrently

src/client/common/utils/async.ts

@@ -110,3 +110,86 @@ export function createDeferredFromPromise<T>(promise: Promise<T>): Deferred<T> {
     promise.then(deferred.resolve.bind(deferred)).catch(deferred.reject.bind(deferred));
     return deferred;
 }
+
+//================================
+// iterators
+
+/**
+ * An iterator that yields nothing.
+ */
+export function iterEmpty<T, R = void>(): AsyncIterator<T, R> {
+    // tslint:disable-next-line:no-empty
+    return ((async function* () {})() as unknown) as AsyncIterator<T, R>;
+}
+
+type NextResult<T, R = void> = { index: number } & (
+    | { result: IteratorResult<T, R>; err: null }
+    | { result: null; err: Error }
+);
+async function getNext<T, R = void>(it: AsyncIterator<T, R>, indexMaybe?: number): Promise<NextResult<T, R>> {
+    const index = indexMaybe === undefined ? -1 : indexMaybe;
+    try {
+        const result = await it.next();
+        return { index, result, err: null };
+    } catch (err) {
+        return { index, err, result: null };
+    }
+}
+
+// tslint:disable-next-line:promise-must-complete no-empty
+const NEVER: Promise<unknown> = new Promise(() => {});
+
+/**
+ * Yield everything produced by the given iterators as soon as each is ready.
+ *
+ * When one of the iterators has something to yield then it gets yielded
+ * right away, regardless of where the iterator is located in the array
+ * of iterators.
+ *
+ * @param iterators - the async iterators from which to yield items
+ * @param onError - called/awaited once for each iterator that fails
+ */
+export async function* chain<T, R = void>(
+    iterators: AsyncIterator<T | void, R>[],
+    onError?: (err: Error, index: number) => Promise<void>
+    // Ultimately we may also want to support cancellation.
+): AsyncIterator<T | R | void, void> {
+    const promises = iterators.map(getNext);
+    let numRunning = iterators.length;
+    while (numRunning > 0) {
+        const { index, result, err } = await Promise.race(promises);
+        if (err !== null) {
+            promises[index] = NEVER as Promise<NextResult<T, R>>;
+            numRunning -= 1;
+            if (onError !== undefined) {
+                await onError(err, index);
+            }
+            // XXX Log the error.
+        } else if (result!.done) {
+            promises[index] = NEVER as Promise<NextResult<T, R>>;
+            numRunning -= 1;
+            // If R is void then result.value will be undefined.
+            if (result!.value !== undefined) {
+                yield result!.value;
+            }
+        } else {
+            promises[index] = getNext(iterators[index], index);
+            yield result!.value;
+        }
+    }
+}
+
+/**
+ * Get everything yielded by the iterator.
+ */
+export async function flattenIterator<T>(iterator: AsyncIterator<T, void>): Promise<T[]> {
+    const results: T[] = [];
+    // We are dealing with an iterator, not an iterable, so we have
+    // to iterate manually rather than with a for-await loop.
+    let result = await iterator.next();
+    while (!result.done) {
+        results.push(result.value);
+        result = await iterator.next();
+    }
+    return results;
+}

src/test/common/utils/async.unit.test.ts

@@ -4,7 +4,7 @@
 'use strict';
 
 import * as assert from 'assert';
-import { createDeferred } from '../../../client/common/utils/async';
+import { chain, createDeferred, flattenIterator } from '../../../client/common/utils/async';
 
 suite('Deferred', () => {
     test('Resolve', (done) => {
@@ -53,3 +53,264 @@ suite('Deferred', () => {
         assert.equal(def.completed, true, 'Promise is not completed even when it should not be');
     });
 });
+
+suite('chain async iterators', () => {
+    const flatten = flattenIterator;
+
+    test('no iterators', async () => {
+        const expected: string[] = [];
+
+        const results = await flatten(chain([]));
+
+        assert.deepEqual(results, expected);
+    });
+
+    test('one iterator, one item', async () => {
+        const expected = ['foo'];
+        const it = (async function* () {
+            yield 'foo';
+        })();
+
+        const results = await flatten(chain([it]));
+
+        assert.deepEqual(results, expected);
+    });
+
+    test('one iterator, many items', async () => {
+        const expected = ['foo', 'bar', 'baz'];
+        const it = (async function* () {
+            yield* expected;
+        })();
+
+        const results = await flatten(chain([it]));
+
+        assert.deepEqual(results, expected);
+    });
+
+    test('one iterator, no items', async () => {
+        const deferred = createDeferred<void>();
+        const it = (async function* () {
+            deferred.resolve();
+        })();
+
+        const results = await flatten(chain([it]));
+
+        assert.deepEqual(results, []);
+        // Make sure chain() actually used up the iterator,
+        // even through it didn't yield anything.
+        assert.ok(deferred.resolved);
+    });
+
+    test('many iterators, one item each', async () => {
+        // For deterministic results we must control when each iterator starts.
+        const deferred12 = createDeferred<void>();
+        const deferred23 = createDeferred<void>();
+        const expected = ['a', 'b', 'c'];
+        const it1 = (async function* () {
+            yield 'a';
+            deferred12.resolve();
+        })();
+        const it2 = (async function* () {
+            await deferred12.promise;
+            yield 'b';
+            deferred23.resolve();
+        })();
+        const it3 = (async function* () {
+            await deferred23.promise;
+            yield 'c';
+        })();
+
+        const results = await flatten(chain([it1, it2, it3]));
+
+        assert.deepEqual(results, expected);
+    });
+
+    test('many iterators, many items each', async () => {
+        // For deterministic results we must control when each iterator starts.
+        const deferred12 = createDeferred<void>();
+        const deferred23 = createDeferred<void>();
+        const expected = ['a1', 'a2', 'a3', 'b1', 'b2', 'b3', 'c1', 'c2', 'c3'];
+        const it1 = (async function* () {
+            yield 'a1';
+            yield 'a2';
+            yield 'a3';
+            deferred12.resolve();
+        })();
+        const it2 = (async function* () {
+            await deferred12.promise;
+            yield 'b1';
+            yield 'b2';
+            yield 'b3';
+            deferred23.resolve();
+        })();
+        const it3 = (async function* () {
+            await deferred23.promise;
+            yield 'c1';
+            yield 'c2';
+            yield 'c3';
+        })();
+
+        const results = await flatten(chain([it1, it2, it3]));
+
+        assert.deepEqual(results, expected);
+    });
+
+    test('many iterators, one empty', async () => {
+        // For deterministic results we must control when each iterator starts.
+        const deferred12 = createDeferred<void>();
+        const deferred23 = createDeferred<void>();
+        const expected = ['a', 'c'];
+        const it1 = (async function* () {
+            yield 'a';
+            deferred12.resolve();
+        })();
+        const it2 = (async function* () {
+            await deferred12.promise;
+            // We do not yield anything.
+            deferred23.resolve();
+        })();
+        const it3 = (async function* () {
+            await deferred23.promise;
+            yield 'c';
+        })();
+        const empty = it2;
+
+        const results = await flatten(chain([it1, empty, it3]));
+
+        assert.deepEqual(results, expected);
+    });
+
+    test('Results are yielded as soon as ready, regardless of source iterator.', async () => {
+        // For deterministic results we must control when each iterator starts.
+        const deferred24 = createDeferred<void>();
+        const deferred41 = createDeferred<void>();
+        const deferred13 = createDeferred<void>();
+        const deferred35 = createDeferred<void>();
+        const deferred56 = createDeferred<void>();
+        const expected = ['b', 'd', 'a', 'c', 'e', 'f'];
+        const it1 = (async function* () {
+            await deferred41.promise;
+            yield 'a';
+            deferred13.resolve();
+        })();
+        const it2 = (async function* () {
+            yield 'b';
+            deferred24.resolve();
+        })();
+        const it3 = (async function* () {
+            await deferred13.promise;
+            yield 'c';
+            deferred35.resolve();
+        })();
+        const it4 = (async function* () {
+            await deferred24.promise;
+            yield 'd';
+            deferred41.resolve();
+        })();
+        const it5 = (async function* () {
+            await deferred35.promise;
+            yield 'e';
+            deferred56.resolve();
+        })();
+        const it6 = (async function* () {
+            await deferred56.promise;
+            yield 'f';
+        })();
+
+        const results = await flatten(chain([it1, it2, it3, it4, it5, it6]));
+
+        assert.deepEqual(results, expected);
+    });
+
+    test('A failed iterator does not block the others, with onError.', async () => {
+        // For deterministic results we must control when each iterator starts.
+        const deferred12 = createDeferred<void>();
+        const deferred23 = createDeferred<void>();
+        const expected = ['a', 'b', 'c'];
+        const it1 = (async function* () {
+            yield 'a';
+            deferred12.resolve();
+        })();
+        const failure = new Error('uh-oh!');
+        const it2 = (async function* () {
+            await deferred12.promise;
+            yield 'b';
+            throw failure;
+        })();
+        const it3 = (async function* () {
+            await deferred23.promise;
+            yield 'c';
+        })();
+        const fails = it2;
+        let gotErr: { err: Error; index: number } | undefined;
+        async function onError(err: Error, index: number) {
+            gotErr = { err, index };
+            deferred23.resolve();
+        }
+
+        const results = await flatten(chain([it1, fails, it3], onError));
+
+        assert.deepEqual(results, expected);
+        assert.deepEqual(gotErr, { err: failure, index: 1 });
+    });
+
+    test('A failed iterator does not block the others, without onError.', async () => {
+        // If this test fails then it will likely fail intermittently.
+        // For (mostly) deterministic results we must control when each iterator starts.
+        const deferred12 = createDeferred<void>();
+        const deferred23 = createDeferred<void>();
+        const expected = ['a', 'b', 'c'];
+        const it1 = (async function* () {
+            yield 'a';
+            deferred12.resolve();
+        })();
+        const failure = new Error('uh-oh!');
+        const it2 = (async function* () {
+            await deferred12.promise;
+            yield 'b';
+            deferred23.resolve();
+            // This is ignored by chain() since we did not provide onError().
+            throw failure;
+        })();
+        const it3 = (async function* () {
+            await deferred23.promise;
+            yield 'c';
+        })();
+        const fails = it2;
+
+        const results = await flatten(chain([it1, fails, it3]));
+
+        assert.deepEqual(results, expected);
+    });
+
+    test('int results', async () => {
+        const expected = [42, 7, 11, 13];
+        const it = (async function* () {
+            yield 42;
+            yield* [7, 11, 13];
+        })();
+
+        const results = await flatten(chain([it]));
+
+        assert.deepEqual(results, expected);
+    });
+
+    test('object results', async () => {
+        type Result = {
+            value: string;
+        };
+        const expected: Result[] = [
+            // We don't need anything special here.
+            { value: 'foo' },
+            { value: 'bar' },
+            { value: 'baz' }
+        ];
+        const it = (async function* () {
+            yield* expected;
+        })();
+
+        const results = await flatten(chain([it]));
+
+        assert.deepEqual(results, expected);
+    });
+});