[RNE Rewrite] fix(ts): release native resources when a create<Task> fails - #1386
[RNE Rewrite] fix(ts): release native resources when a create<Task> fails#1386msluszniak wants to merge 1 commit into
Conversation
A factory allocates as it goes: it loads a model, maybe a tokenizer, a phonemizer, an LLM runner or a nested pipeline, validates the schema, pre-allocates its tensors, and only at the end hands back a `dispose`. Anything that threw in between left the caller with no reference to what was already allocated, and native memory is not garbage collected, so it stayed alive for the rest of the process. `useModel` re-runs its factory whenever the config changes, so an app pointed at a mismatched model leaked a full resource set per attempt. `createResourceScope` gives a factory one teardown path for both outcomes: it tracks each resource as it is created and releases them in reverse order, and the same function becomes the pipeline's `dispose`. Every factory now allocates through a scope and wraps its body in try/catch. Kokoro already did this with a local array and moves to the shared helper, which also closes a smaller hole it had: its two models were loaded before the try block, so one load rejecting stranded the other. Parallel loads elsewhere track inside each promise for the same reason. `createTokenizer` is unchanged. It loads a tokenizer and returns with nothing in between that can throw, so it has no window to leak through. The helper is internal, not exported from `src/index.ts`, so the public surface is unchanged. Test-side, `constructionFailure.test.ts` was recording the leak as current behavior. It now asserts the opposite and covers all fifteen factories rather than five, checking every kind of handle so a failure names the factory and the resource. The per-pipeline suites and `hooks/taskHooks.test.ts` drop the `allowNativeLeaks()` calls they needed for the same reason: nothing in the suite leaks any more, so the global leak check now asserts this on every construction-failure test for free.
| @@ -0,0 +1,74 @@ | |||
| /** | |||
| * Construction-time ownership of native resources. | |||
There was a problem hiding this comment.
Please change the docs so that they don't mention create<Task> at all. The extensions depend on core no the other way around so nothing in core should reference extensions.
| * one `dispose` releases whatever exists at the moment it is called. The same | ||
| * function is the factory's own `dispose`, so there is a single teardown path | ||
| * rather than one for failure and one for success. | ||
| * @module Core/Lifetime |
There was a problem hiding this comment.
| * @module Core/Lifetime |
These are not used in docs actually.
| readonly track: <D extends Disposable>(resource: D) => D; | ||
| /** |
There was a problem hiding this comment.
| readonly track: <D extends Disposable>(resource: D) => D; | |
| /** | |
| readonly track: <D extends Disposable>(resource: D) => D; | |
| /** |
| * Creates a {@link ResourceScope} for a factory to allocate into. | ||
| * | ||
| * Wrap the factory body in `try`/`catch`, `track` each resource as it is | ||
| * created, and return `scope.dispose` as the pipeline's `dispose`: | ||
| * ```typescript | ||
| * const scope = createResourceScope(); | ||
| * const dispose = scope.dispose; | ||
| * try { | ||
| * const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); | ||
| * const { dims } = validateSpec(model.schema, { ... }); // may throw | ||
| * const tensors = [scope.track(tensor('float32', shape))]; | ||
| * return { run, dispose }; | ||
| * } catch (error) { | ||
| * dispose(); | ||
| * throw error; | ||
| * } |
There was a problem hiding this comment.
Same as above. Let's not mention extension context at all here, it's just for semi-automatic lifetime management.
There was a problem hiding this comment.
This should also be added to public API, so users can follow this pattern in their custom pipelines.
| const [model, tokenizer] = await Promise.all([ | ||
| wrapAsync(loadModel, runtime)(modelPath).then(scope.track), | ||
| wrapAsync(loadTokenizer, runtime)(tokenizerPath).then(scope.track), | ||
| ]); |
There was a problem hiding this comment.
Using Promise.all can still leak if I'm not mistaken, because it will reject on the first reject, but won't cancel the other enqueued loads. Safer to just rewrite this to sequantial await calls.
| * ```typescript | ||
| * const scope = createResourceScope(); | ||
| * const dispose = scope.dispose; | ||
| * try { | ||
| * const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); | ||
| * const { dims } = validateSpec(model.schema, { ... }); // may throw | ||
| * const tensors = [scope.track(tensor('float32', shape))]; | ||
| * return { run, dispose }; | ||
| * } catch (error) { | ||
| * dispose(); | ||
| * throw error; | ||
| * } |
There was a problem hiding this comment.
Put this example at the bottom under @example tag.
Description
Fixes the finding recorded in #1355: a
create<Task>that throws part-way through construction abandoned everything it had already allocated.A factory allocates as it goes, and only hands back a
disposeat the very end. Anything that threw in between left the caller with no reference to what was already there, and native memory is not garbage collected, so it stayed alive for the rest of the process. Depending on the task that is a model, a tokenizer, a phonemizer, an LLM runner, or a whole nested pipeline (Whisper owns both a tokenizer and a VAD).useModelre-runs its factory whenever the config changes, so an app pointed at a mismatched model leaked a full resource set per attempt.src/core/lifetime.tsaddscreateResourceScope, which gives a factory one teardown path for both outcomes: it tracks each resource as it is created, releases them in reverse order, and the same function becomes the pipeline'sdispose. Every factory now allocates through a scope and wraps its body intry/catch:createKokoroTextToSpeechalready did this with a local array and moves to the shared helper. That also closes a smaller hole it had: its two models were loaded before thetryblock, so one load rejecting stranded the other. Parallel loads elsewhere now track inside each promise for the same reason.createTokenizeris unchanged. It loads a tokenizer and returns with nothing in between that can throw, so it has no window to leak through.The helper is internal, so the public export surface is unchanged.
Introduces a breaking change?
Type of change
Tested on
Testing instructions
yarn workspace react-native-executorch test yarn typecheck yarn lintExpected: 28 suites, 3117 tests, 4 snapshots, 0 skipped; typecheck and lint clean.
To see the tests bite, stash
src/and re-run: every construction-failure case fails.Screenshots
N/A
Related issues
Follows up the finding recorded in #1355.
Checklist
Additional notes
Stacked on
@ms/api-tests(#1355), because that is where the suite recording this behavior lives. Merge #1355 first and this retargets torne-rewritecleanly.tasks/constructionFailure.test.tswas written to record the leak as current behavior so it would fail loudly the day a factory started cleaning up. That is this PR, so it now asserts the opposite, and covers all fifteen factories rather than five. It checks every kind of handle, so a failure names the factory and the resource rather than just reporting that something leaked.The per-pipeline suites and
hooks/taskHooks.test.tsdrop theallowNativeLeaks()calls they needed for the same reason. Nothing in the suite leaks any more, so the setup file's global leak check now asserts this on every construction-failure test for free.The diff is large mostly because indenting a factory body inside
trytouches every line of it. The behavioral change per file is the scope, thetrackcalls and thecatch.