Optimize run to avoid quadratic map cloning

[jridgewell]

Currently, every time run() is called, we clone the underlying storage Map. This is done so that the the modification will not change a map which is held by a wrap's snapshot. I believe "wrapping" (either explicitly with AsyncContext.wrap() or implicitly through Promises) will be performed 10000x more times in a normal program than run(), so I'm doing everything possible to make the wrap() method as fast as possible. Further, that get() will be called at least at much and probably more as run(). So my priorities are:

• wrap() must be O(1) (and ideally require nothing more than a pointer reference)
• get() should be as fast as possible
• run() may be slow, but let's try to make it fast

These priorities lead to code that causes quadratic cloning of the underlying Map as we add more keys in nesting run() calls. That's not ideal.

// clones map with 0 keys
a.run(1, () => {
  // clones map with 1 key
  b.run(2, () => {
    // clones map with 2 keys
    c.run(3, () => {
      a.get();
      b.get();
      c.get();
    });
  });
});

We can fix this if we know whether a snapshot is held by anyone. If there is, then we must treat the Map as immutable (because someone has a direct access to it). If not, then we can directly mutate without worrying.

// no clone, direct mutate
a.run(1, () => {
  // no clone, direct mutate
  b.run(2, () => {
    // no clone, direct mutate
    c.run(3, () => {
      a.get();
      b.get();
      c.get();
    });
  });
});

Because no wrapping was done (and there's no await/promises), we didn't have to clone anything.

There are some interesting cases with the new code. Once a snapshot has been taken, further mutations will clone the underlying Map and mutate that clone. This comes up when we enter a run from a frozen state, and when we exit a run that has taken a snapshot:

// Entering this state will not clone
a.run(1, () => {
  // Entering this state will not clone
  b.run(2, () => {

    // Take a snapshot of the current { a: 1, b: 2 } state,
    // which will freeze the mappings.
    AsyncContext.wrap(cb);

    // Because the current state is frozen, entering this run will require
    // a clone to add the `c: 3` data.
    c.run(3, () => {
      // The state inside this run isn't frozen, it's a brand new clone of
      // the old state with the `c: 3` data.

      // Entering this run will not clone.
      d.run(4, () => {
        // Exiting this run will not clone.
      });

      // Exiting this run will not clone.
    });

    // We're back to the frozen { a: 1, b: 2 } state.
    
    // Because we're frozen, we cannot remove the `b: 2` data.
    // Exiting will cause a clone.
  });

  // We're back to the `{ a: 1 }` state.
  // This state isn't frozen.
 
  // Exiting this run will not clone.
});

We can prove that the new code clones less overall than the old code does. In the old code, every enter would clone. In the new code, we only clone:

• on enter if the current state is frozen
  • Meaning entering is no worse, and likely much better
• on exit if the the current state is frozen (the exit case is a littler harder to reason about)
  1. Current state is unfrozen before entering, and we snapshot the new state before exit:
     1. we did not clone to enter
     2. we freeze new state
     3. we clone on exit to remove the newly added key
     • This is no worse than cloning to enter in the old code.
  2. Current state is frozen before entering, and we snapshot the new state before exit:
     1. we clone to enter
     2. we freeze new state
     3. we restore the old (known-frozen-at-start) state without modifying the current state
     • The trick is the way we restore!
     • Instead of trying to delete the newly added key, we just set the global state back to the old frozen state.

So, we're more optimal in all cases:

• wrap(): O(1)
  • allocate a struct with only a single pointer, which can optimize to a single pointer in memory
• get(): O(1)
  • just a Map.p.get()
• run(): O(1) (if not wrapped) or O(n) (if wrapped)
  • No worse than before, and sometimes better.

There are alternative algorithms using purely-functional immutable data structures, which trade make run() and get() operate at O(1) and O(n), but I don't think that's optimal.

[legendecas]

Thank you for wrapping this awesome work up!

I find that the mapping Storage is similar to a conceptual tree:

• AsyncContext.wrap saves the active head of the tree and marks it as immutable.
  • Subsequent modification (AsyncContext.prototype.run) to the head of the tree would create a new leaf and make the new leaf the active head.
• When the wrapped function is called,
  • Storage switches to the saved leaf and switch back to the current tip after fn is invoked.
• When AsyncContext.prototype.run is called,
  • If the active tip of the tree is immutable, Storage creates a new leaf from it and switches to the new leaf as the active head. After the fn is invoked, the head of the tree is switched back to the previous one.
  • If the active head of the tree is mutable, Storage saves the current value corresponding to the AsyncContext instance at the head and resets the value after the fn is invoked.
• AsyncContext.prototype.get retrieves the value from the active head corresponding to the AsyncContext instance.

In this sense, I'm wondering if a renaming would describe these operations in a more clear way (See https://github.com/legendecas/proposal-async-context/blob/2220629/src/storage.ts#L11 for more details):

export class AsyncContext<T> {
  static wrap<F extends AnyFunc<any>>(fn: F): F {
    // No clones in commit.
    const wrapHead = Storage.commit();

    function wrap(this: ThisType<F>, ...args: Parameters<F>): ReturnType<F> {
      const head = Storage.switch(wrapHead);
      try {
        return fn.apply(this, args);
      } finally {
        Storage.switch(head);
      }
    }

    return wrap as unknown as F;
  }

  run<F extends AnyFunc<null>>(
    value: T,
    fn: F,
    ...args: Parameters<F>
  ): ReturnType<F> {
    const head = Storage.stage(this);
    Storage.set(this, value);
    try {
      return fn.apply(null, args);
    } finally {
      Storage.switch(head);
    }
  }

  get(): T | undefined {
    return Storage.get(this);
  }
}

[legendecas]

I believe this method can be merged with Storage.restore. They all restore the global storage state to state at the time of the fork.

[jridgewell]

Not quite. join and restore (now restore and switch respectively) have different assumptions about what can happen to the current global state.

For join, we assume the the current mappings will be modified (or if it's frozen, reallocated then modified). For restore, we want to switch back to the state of a wrap's snapshot and return a FrozenFork so the wrapper can restore the prior state after running. It wouldn't be valid for join to return a fork.

[legendecas]

Thanks, the naming is more intuitive to me now.