Skip to content

Commit b78c76d

Browse files
yaacovCRyaacovCR
committed
extract withCleanup from mapAsyncIterable
1 parent 9b3d5b6 commit b78c76d

File tree

5 files changed

+223
-76
lines changed

5 files changed

+223
-76
lines changed

‎src/execution/__tests__/mapAsyncIterable-test.ts‎

Lines changed: 0 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -89,54 +89,6 @@ describe('mapAsyncIterable', () => {
8989
});
9090
});
9191

92-
it('calls done when completes', async () => {
93-
async function* source() {
94-
yield 1;
95-
yield 2;
96-
yield 3;
97-
}
98-
99-
let done = false;
100-
const doubles = mapAsyncIterable(
101-
source(),
102-
(x) => Promise.resolve(x + x),
103-
() => {
104-
done = true;
105-
},
106-
);
107-
108-
expect(await doubles.next()).to.deep.equal({ value: 2, done: false });
109-
expect(await doubles.next()).to.deep.equal({ value: 4, done: false });
110-
expect(await doubles.next()).to.deep.equal({ value: 6, done: false });
111-
expect(done).to.equal(false);
112-
expect(await doubles.next()).to.deep.equal({
113-
value: undefined,
114-
done: true,
115-
});
116-
expect(done).to.equal(true);
117-
});
118-
119-
it('calls done when completes with error', async () => {
120-
async function* source() {
121-
yield 1;
122-
throw new Error('Oops');
123-
}
124-
125-
let done = false;
126-
const doubles = mapAsyncIterable(
127-
source(),
128-
(x) => Promise.resolve(x + x),
129-
() => {
130-
done = true;
131-
},
132-
);
133-
134-
expect(await doubles.next()).to.deep.equal({ value: 2, done: false });
135-
expect(done).to.equal(false);
136-
await expectPromise(doubles.next()).toRejectWith('Oops');
137-
expect(done).to.equal(true);
138-
});
139-
14092
it('allows returning early from mapped async generator', async () => {
14193
async function* source() {
14294
try {

‎src/execution/__tests__/withCleanup-test.ts‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
import { expect } from 'chai';
2+
import { describe, it } from 'mocha';
3+
4+
import { expectPromise } from '../../__testUtils__/expectPromise.js';
5+
6+
import { withCleanup } from '../withCleanup.js';
7+
8+
/* eslint-disable @typescript-eslint/require-await */
9+
describe('withCleanup', () => {
10+
it('calls cleanup function when completes', async () => {
11+
async function* source() {
12+
yield 1;
13+
}
14+
15+
let done = false;
16+
const generator = withCleanup(source(), () => {
17+
done = true;
18+
});
19+
20+
expect(await generator.next()).to.deep.equal({ value: 1, done: false });
21+
expect(done).to.equal(false);
22+
expect(await generator.next()).to.deep.equal({
23+
value: undefined,
24+
done: true,
25+
});
26+
expect(done).to.equal(true);
27+
});
28+
29+
it('calls cleanup function when completes with error', async () => {
30+
async function* source() {
31+
yield 1;
32+
throw new Error('Oops');
33+
}
34+
35+
let done = false;
36+
const generator = withCleanup(source(), () => {
37+
done = true;
38+
});
39+
40+
expect(await generator.next()).to.deep.equal({ value: 1, done: false });
41+
expect(done).to.equal(false);
42+
await expectPromise(generator.next()).toRejectWith('Oops');
43+
expect(done).to.equal(true);
44+
});
45+
46+
it('calls cleanup function when returned', async () => {
47+
async function* source() {
48+
yield 1;
49+
}
50+
51+
let done = false;
52+
const generator = withCleanup(source(), () => {
53+
done = true;
54+
});
55+
56+
expect(await generator.next()).to.deep.equal({ value: 1, done: false });
57+
expect(done).to.equal(false);
58+
expect(await generator.return()).to.deep.equal({
59+
value: undefined,
60+
done: true,
61+
});
62+
expect(done).to.equal(true);
63+
});
64+
65+
it('calls cleanup function when thrown', async () => {
66+
async function* source() {
67+
yield 1;
68+
}
69+
70+
let done = false;
71+
const generator = withCleanup(source(), () => {
72+
done = true;
73+
});
74+
75+
expect(await generator.next()).to.deep.equal({ value: 1, done: false });
76+
expect(done).to.equal(false);
77+
await expectPromise(generator.throw(new Error('Oops'))).toRejectWith(
78+
'Oops',
79+
);
80+
expect(done).to.equal(true);
81+
});
82+
83+
it('calls cleanup function when disposed', async () => {
84+
let returned = false;
85+
86+
const items = [1, 2, 3];
87+
const source: AsyncGenerator<number, void, void> = {
88+
[Symbol.asyncIterator]() {
89+
return this;
90+
},
91+
next(): Promise<IteratorResult<number, void>> {
92+
const value = items.shift();
93+
if (value !== undefined) {
94+
return Promise.resolve({ done: false, value });
95+
}
96+
97+
return Promise.resolve({ done: true, value: undefined });
98+
},
99+
return(): Promise<IteratorResult<number, void>> {
100+
returned = true;
101+
return Promise.resolve({ done: true, value: undefined });
102+
},
103+
throw(): Promise<IteratorResult<number, void>> {
104+
returned = true;
105+
return Promise.reject(new Error());
106+
},
107+
async [Symbol.asyncDispose]() {
108+
await this.return();
109+
},
110+
};
111+
112+
let cleanedUp = false;
113+
{
114+
await using generator = withCleanup(source, () => {
115+
cleanedUp = true;
116+
});
117+
118+
expect(await generator.next()).to.deep.equal({ value: 1, done: false });
119+
expect(await generator.next()).to.deep.equal({ value: 2, done: false });
120+
}
121+
122+
expect(cleanedUp).to.equal(true);
123+
expect(returned).to.equal(true);
124+
});
125+
126+
it('returns the generator itself when the `Symbol.asyncIterator` method is called', async () => {
127+
async function* source() {
128+
yield 1;
129+
}
130+
131+
const generator = withCleanup(source(), () => {
132+
/* noop */
133+
});
134+
135+
expect(generator[Symbol.asyncIterator]()).to.equal(generator);
136+
});
137+
});

‎src/execution/execute.ts‎

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ import {
8989
getDirectiveValues,
9090
getVariableValues,
9191
} from './values.js';
92+
import { withCleanup } from './withCleanup.js';
9293

9394
/* eslint-disable max-params */
9495
// This file contains a lot of such errors but we plan to refactor it anyway
@@ -2265,19 +2266,23 @@ function mapSourceToResponse(
22652266
// GraphQL `execute` function, with `payload` as the rootValue.
22662267
// This implements the "MapSourceToResponseEvent" algorithm described in
22672268
// the GraphQL specification..
2268-
return mapAsyncIterable(
2269-
abortSignalListener
2270-
? cancellableIterable(resultOrStream, abortSignalListener)
2271-
: resultOrStream,
2272-
(payload: unknown) => {
2273-
const perEventExecutionArgs: ValidatedExecutionArgs = {
2274-
...validatedExecutionArgs,
2275-
rootValue: payload,
2276-
};
2277-
return validatedExecutionArgs.perEventExecutor(perEventExecutionArgs);
2278-
},
2279-
() => abortSignalListener?.disconnect(),
2280-
);
2269+
function mapFn(payload: unknown): PromiseOrValue<ExecutionResult> {
2270+
const perEventExecutionArgs: ValidatedExecutionArgs = {
2271+
...validatedExecutionArgs,
2272+
rootValue: payload,
2273+
};
2274+
return validatedExecutionArgs.perEventExecutor(perEventExecutionArgs);
2275+
}
2276+
2277+
return abortSignalListener
2278+
? withCleanup(
2279+
mapAsyncIterable(
2280+
cancellableIterable(resultOrStream, abortSignalListener),
2281+
mapFn,
2282+
),
2283+
() => abortSignalListener.disconnect(),
2284+
)
2285+
: mapAsyncIterable(resultOrStream, mapFn);
22812286
}
22822287

22832288
export function executeSubscriptionEvent(

‎src/execution/mapAsyncIterable.ts‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,26 +7,18 @@ import type { PromiseOrValue } from '../jsutils/PromiseOrValue.js';
77
export function mapAsyncIterable<T, U, R = undefined>(
88
iterable: AsyncGenerator<T, R, void> | AsyncIterable<T>,
99
callback: (value: T) => PromiseOrValue<U>,
10-
onDone?: () => void,
1110
): AsyncGenerator<U, R, void> {
1211
const iterator = iterable[Symbol.asyncIterator]();
1312

1413
async function mapResult(
1514
promise: Promise<IteratorResult<T, R>>,
1615
): Promise<IteratorResult<U, R>> {
17-
let value: T;
18-
try {
19-
const result = await promise;
20-
if (result.done) {
21-
onDone?.();
22-
return result;
23-
}
24-
value = result.value;
25-
} catch (error) {
26-
onDone?.();
27-
throw error;
16+
const result = await promise;
17+
if (result.done) {
18+
return result;
2819
}
2920

21+
const value = result.value;
3022
try {
3123
return { value: await callback(value), done: false };
3224
} catch (error) {
@@ -46,6 +38,10 @@ export function mapAsyncIterable<T, U, R = undefined>(
4638
}
4739
}
4840

41+
const asyncDispose: typeof Symbol.asyncDispose =
42+
Symbol.asyncDispose /* c8 ignore start */ ??
43+
Symbol.for('Symbol.asyncDispose'); /* c8 ignore stop */
44+
4945
return {
5046
async next() {
5147
return mapResult(iterator.next());
@@ -70,13 +66,13 @@ export function mapAsyncIterable<T, U, R = undefined>(
7066
[Symbol.asyncIterator]() {
7167
return this;
7268
},
73-
async [Symbol.asyncDispose]() {
69+
async [asyncDispose]() {
7470
await this.return(undefined as R);
7571
if (
76-
typeof (iterable as AsyncGenerator<T, R, void>)[Symbol.asyncDispose] ===
72+
typeof (iterable as AsyncGenerator<T, R, void>)[asyncDispose] ===
7773
'function'
7874
) {
79-
await (iterable as AsyncGenerator<T, R, void>)[Symbol.asyncDispose]();
75+
await (iterable as AsyncGenerator<T, R, void>)[asyncDispose]();
8076
}
8177
},
8278
};

‎src/execution/withCleanup.ts‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import type { PromiseOrValue } from '../jsutils/PromiseOrValue.js';
2+
3+
/**
4+
* Given an AsyncGenerator and a cleanup function, return an AsyncGenerator
5+
* which calls the given function when the generator closes.
6+
*
7+
* This is useful for ensuring cleanup logic is called immediately when the
8+
* generator's `return()` method is called, even if the generator is currently
9+
* paused, e.g. if a `await` is pending within the generator's `next()` method.
10+
*/
11+
export function withCleanup<T>(
12+
generator: AsyncGenerator<T, void, void>,
13+
onDone: () => PromiseOrValue<void>,
14+
): AsyncGenerator<T, void, void> {
15+
let finished = false;
16+
const finish = async () => {
17+
if (!finished) {
18+
finished = true;
19+
await onDone();
20+
}
21+
};
22+
23+
const asyncDispose: typeof Symbol.asyncDispose =
24+
Symbol.asyncDispose /* c8 ignore start */ ??
25+
Symbol.for('Symbol.asyncDispose'); /* c8 ignore stop */
26+
27+
return {
28+
[Symbol.asyncIterator]() {
29+
return this;
30+
},
31+
async next() {
32+
try {
33+
const result = await generator.next();
34+
if (result.done) {
35+
await finish();
36+
return result;
37+
}
38+
return { value: await result.value, done: false };
39+
} catch (error) {
40+
await finish();
41+
throw error;
42+
}
43+
},
44+
async return(): Promise<IteratorResult<T>> {
45+
await finish();
46+
return generator.return();
47+
},
48+
async throw(error?: unknown) {
49+
await finish();
50+
return generator.throw(error);
51+
},
52+
async [asyncDispose]() {
53+
await finish();
54+
await generator[asyncDispose]();
55+
},
56+
};
57+
}

0 commit comments

Comments
 (0)