withAsyncContext race condition

[skirtles-code]

Version

3.1.3

Reproduction link

sfc.vuejs.org

Steps to reproduce

Note the rendered text.

What is expected?

It should say 'first - first / second - second'.

What is actually happening?

It says 'first - second / second - second'. The 'second' has jumped to the first component.

getCurrentInstance is returning the second component in both cases.

---

As withAsyncContext is new in 3.1.3 it isn't necessarily clear whether what I'm doing is allowed. However...

This failure is quite subtle. In 3.1.2 it would have failed obviously and explosively, whereas in 3.1.3 it appears to work and mysteriously yields incorrect values in certain cases.

I haven't stepped through in the debugger but from looking at the code I believe the problem is that the compiled code using withAsyncContext is assuming that the microtask queue is only dealing with a single component at once. My example creates two promises that resolve at the same time, so their microtasks get weaved together in the queue and the currentInstance gets set twice before the setup code resumes.

[jods4]

My fault.
The design I proposed for withAsyncContext is flawed.

The problem is that withAsyncContext itself returns a promise. It is awaited. Continuation code doesn't run synchronously after it.

What happens in this bug report is:

• Comp 1 runs before
• Comp 1 saves its context and awaits 7
• Comp 2 runs before
• Comp 2 saves its context and awaits 7
• <then>
• Comp 1 restores its state
• Comp 1 awaits withAsyncContext itself
• <then>
• Comp 2 restores its state
• Comp 2 awaits withAsyncContext itself
• <then>
• Comp 1 runs after, it has the second, incorrect, context
• <then>
• Comp 2 runs after, it has the second context, by chance!
• <then>
• Vue setup for component 1 completed, it removes context
• <then>
• Vue setup for component 2 completed, it removes context

As can be seen, the problem is deeper than just withAsyncContext, it also touches on how Vue removes the component after setup completes.

Fundamentally, we would like to run code synchronously, i.e. immediately, when some promises settle.
This is possible on other platforms (e.g. .net) but AFAIK impossible in JS.

I suppose the only solution is to insert those synchronous calls manually in the right places.
For generated code such as script setup, that's no biggie.
But usability of classical script or even async composable functions will suffer a lot.

The high-level concept is that we want this:

const awaitable = ...;
const ctx = saveCtx(); // happens right before await
const result = await awaitable;
ctx.restore(); // must not happen in a Promise, must be right here in sync code
...
clearCtx(); // must not happen in a continuation of this promise, must be right here in sync code

(I'm gonna continue with a revised API proposal)

[jods4]

I'm sharing one failed idea, just in case it inspires someone:

I thought I could subclass Promise into a VuePromise, or just create a VueThenable.
(Subclassing Promise is super-tricky, but doable)

My idea was to set/remove context in .then() just so that it happens synchronously with the continuation itself.
That would have worked wonders, the api could work as-is in 3.1.3.

And it does work in .then()-style code.
But sadly, aync/await is doing crazy shenanigans, you'll never guess the exact sequence in which everything happen.
Long story short, the resolve callback that is passed to .then() is not synchronously executing user code; it enqueues a micro-task.
So attempting to set/remove context around this callback is not working.

Given this, it seems there has to be some calls in user code to set/remove context at boundary points.
Trying to think what the nicest API to do so could be, but nothing really satisfying comes to mind.

[jods4]

While working on the new proposal I noticed another bug.

If you use the current withAsyncContext in an asynchronous composable function, e.g.:

async function useRemoteSum() {
  const data = await withAsyncContext(fetch("/data"));
  return computed(() => data.reduce((a, b) => a + b, 0));
}

Then the caller of this function (e.g. setup()) will have already lost context when its own withAsyncContext is called.

[jods4]

Here's a new API proposal.
Async code is really tricky, I hope I thought of everything, even if I didn't spell it out fully -- and that's already a very long post. 🤞
I am not a huge fan but I don't have a better idea right now, I tried to balance ease of use and correctness as much as possible...

Far future

Basically what we want is an AsyncLocal from the Async Context proposal (not even stage 1).
If this lands, none of this would be required and everything can be hidden inside Vue.

Until then we need a workaround. Quote from that proposal:

While monkey-patching is quite straightforward solution to track async tasks, there is no way to patch JavaScript features like async/await. Also, monkey-patching only works if all third-party libraries with custom scheduling call a corresponding task awareness registration like AsyncResource.runInAsyncScope. Furthermore, for those custom scheduling third-party libraries, we need to get library owners to think in terms of async context propagation.

General idea

ctx = new AsyncContext() captures the current async context.
This must happen before nested async functions are called, as they will clear the current context at their await point.

ctx.async(awaitable) unsets context.
It returns an object with a result getter, which will be responsible for restoring context.
The key insight here is that result is read inside the user code.
Doing it this way rather than adding a ctx.restore() call means: one less call; and less likely to be forgotten as you can't use the returned value without it.

At the end of an async component setup, users must explicitely call ctx.end(), which unsets the context.
We can't leave this responsability to Vue itself, because other micro-tasks can run between the setup completion and the Vue continuation.
For convenience in functions that return an expression (see useAsync example below), it can passthrough a parameter ctx.end(42) == 42.

Explicit Resource Management is a stage 2 proposal that could make this better. If it lands, ctx.end() would not be necessary anymore if you declare it with using const ctx = new AsyncContext() instead.

About effectScope()

effectScope() is not finalized so things will have to be adjusted to take it into account if an effectScope can be asynchronous.

Basically, the effectScope callback should follow the same rules as setup function.

API

class AsyncContext {
  #ctx = getCurrentInstance();
  
  async<T>(awaitable: T | PromiseLike<T>) {
    setCurrentInstance(null);
    return Promise.resolve(awaitable).then(
      value => ({ get result() { setCurrentInstance(this.#ctx); return value; } }),
      error => ({ get result() { setCurrentInstance(this.#ctx); throw error; } })
    );
  }
  
  end<T>(result?: T) {
    setCurrentInstance(null);
    return result;
  }
  
  // If one day Explicit Resource Management lands, we could add the following:
  [Symbol.dispose]() { this.end() }
} 

Example usage

This is what is required in hand-written code.
The nice thing about script setup is that all of this can be automatic.

async setup() {
  const ctx = new AsyncContext();
  // ... do stuff
  const { result: data } = await ctx.async(fetch("/data"));
  // ... do more stuff with data
  ctx.end();
}

// Async composable function

// 1. If no Vue function is called after await, then nothing special needs to be done.
//    This is a tricky situation, though. 
//    You must be sure that neither your code, nor any transitive function you call, needs Vue context.
useAsync() {
  // Vue API before async call is ok.
  const result = computed(() => 42);
  // No need for context preservation if context isn't used after the async call
  await delay(10);
  return result;
}

// 2. If Vue context is required after await, AsyncContext must be used.
useAsync() {
  const ctx = new AsyncContext();
  const a = computed(() => 42);
  const { result: _ } = await ctx.async(delay(10));
  // Do more Vue stuff
  const b = computed(() => a + 1);
  return ctx.end(b);
}  

Pitfalls and safeguards

1. Any async code that uses context-aware Vue functions needs to create an AsyncContext.
   setup and effectScope could use a flag to check if an AsyncContext has been created when a Promise is returned, and warn in dev. This isn't perfect and induces:
   (a) False negative for nested async function calls, we can only check if at least one AsyncContext was created. Docs and linters can help here.
   (b) False positive for the case where an async setup doesn't require context after its first await. That's technically ok but it might be good to force context usage to be on the safe side anyway.

2. ctx.end() must be called.
   Assuming 1st created context during setup or effectScope is captured, we could check when the promise settles if end() was called and warn in dev if it wasn't.
   Impossible to check in composable async functions, but docs and linter could help.
   setup and effectScope should still remove the current context when the they first return, and when the promise settles, for extra safety or the case where an async function doesn't need to preserve its context.

3. ctx.async() must be called at async points. This can be validated in the same way as 4 below.

4. ctx.async().result must be read.
   This is a natural thing to do, except for void functions where you could do this: await ctx.async(delay(10)).
   We could check for that in dev by having the next ctx.async or ctx.end call check if there's a context set.
   If the context is null then something is amiss.

5. This all assumes async setup don't crash in normal operation.
   Otherwise the next ctx.async() or ctx.end() won't be called and the context will not be cleared. This can lead to subtle errors as there's a small window where micro-tasks will see a set context, until Vue clears it because the promise has settled.
   A function that can crash in normal operations should use a try..finally to ensure ctx.end() is called.

[yyx990803]

Thanks for the research into this @jods4 - I think it's ok for the API to be a bit unwieldy as long as it is easy to automatically generate.

It isn't that much a burden for userland composition functions to ensure proper placement of await - in fact I'd say it's good to discourage await in composition functions in general as it makes it too easy to create serial awaits that makes the suspense tree render appear slower.

[jods4]

@yyx990803 03e2684 is not a complete, proper fix.

I noticed:

1. You can't restore context in case of failure inside a .catch().
   That's like restoring context in case of succes inside .then().
   The continuation of .catch() runs in a micro-task so anything could run inbetween.
   Restoring the context has to happen sync inside the user code.

2. I don't see a call to clear the context at the end of setup.
   You can't rely on Vue runtime to do so, because anything can run inbetween the end of setup and the runtime continuation.

3. Capturing the context right before the expression that is awaited can lead to code patterns that mysteriously won't work, such as:

let promise = fetchData();
let data = await promise; // context is already lost here

That's why I chose captured the state once at the beginning of code.

More thoughts on advanced patterns

I realized that for code patterns that manipulate promises without using await (could be as simple as a .then()), transformations can't be fully automated and something needs to be added to script setup to support them fully.
(Of course, await should be encouraged to keep code simple and correct.)

Here's some code that tries to fetch three things in parallel:

const promises = [fetch(1), fetch(2), fetch(3)]
const [data1, data2, data3] = await promises;

Using the API I proposed, it's possible to write the following code manually:

// at beginning of method
const ctx = new VueContext();

const promises = [
  ctx.async(fetch(1), /*stayInContext:*/ true),
  ctx.async(fetch(2), true),
  ctx.async(fetch(3), true),
];
const { result: [data1, data2, data3] } = await ctx.async(promises);

// at the end of method
ctx.end();

But that can't be a fully automated transformation, because of the fetch calls without immediate await.
This kind of advanced async methods would need an API to help the code transformation.

Maybe the compiler should allow and recognize const ctx = new VueContext() in an async setup method and re-use it instead of injecting its own -> this allows advanced cases where users introduce more calls.
The example above would become, in script setup style:

// Not required in basic case, but is picked up by compiler
const ctx = new VueContext();

// Additional, manual calls
const promises = [
  ctx.async(fetch(1), /*stayInContext:*/ true),
  ctx.async(fetch(2), true),
  ctx.async(fetch(3), true),
];

// await is handled by compiler itself
const [data1, data2, data3] = await promises;

// end of setup is handled by compiler

[paraboul]

Thanks for the research into this @jods4 - I think it's ok for the API to be a bit unwieldy as long as it is easy to automatically generate.

It isn't that much a burden for userland composition functions to ensure proper placement of await - in fact I'd say it's good to discourage await in composition functions in general as it makes it too easy to create serial awaits that makes the suspense tree render appear slower.

@yyx990803
While I agree that chain blocking should be discouraged, it's easy to shoot yourself in the foot if you have a watch call placed after an await at setup() level, as it would (rightfully) not warn you being outside of an active instance -- while silently leaking. Having withAsyncContext being some sort of syntactic sugar as an escape plan could be practical in some cases (edit: unless <script setup> become the de facto standard for SFC?).

As you said, this new "fixed" version (03e2684) makes it almost impossible to use for handwritten code.

Although it's not going to be possible to have a simple promise wrapper, seeking for an intermediate readable syntax is not crazy imho.
Just my two cents after being relieved for this new feature of 3.1.3 :-)