Can I prevent an `AsyncGenerator` from yielding after its `return()` method has been invoked?

Question

AsyncGenerator.prototype.return() - JavaScript | MDN states:

The return() method of an async generator acts as if a return statement is inserted in the generator's body at the current suspended position, which finishes the generator and allows the generator to perform any cleanup tasks when combined with a try...finally block.

Why then does the following code print 0–3 rather than only 0–2?

const delay = (ms) => new Promise((resolve) => setTimeout(resolve, ms));

const values = (async function* delayedIntegers() {
  let n = 0;
  while (true) {
    yield n++;
    await delay(100);
  }
})();

await Promise.all([
  (async () => {
    for await (const value of values) console.log(value);
  })(),
  (async () => {
    await delay(250);
    values.return();
  })(),
]);

I tried adding log statements to better understand where the "current suspended position" is and from what I can tell when I call the return() method the AsyncGenerator instance isn't suspended (the body execution isn't at a yield statement) and instead of returning once reaching the yield statement the next value is yielded and then suspended at which point the "return" finally happens.

Is there any way to detect that the return() method has been invoked and not yield afterwards?

---

I can implement the AsyncIterator interface myself but then I lose the yield syntax supported by async generators:

const delay = (ms) => new Promise((resolve) => setTimeout(resolve, ms));

const values = (() => {
  let n = 0;
  let done = false;
  return {
    [Symbol.asyncIterator]() {
      return this;
    },
    async next() {
      if (done) return { done, value: undefined };
      if (n !== 0) {
        await delay(100);
        if (done) return { done, value: undefined };
      }
      return { done, value: n++ };
    },
    async return() {
      done = true;
      return { done, value: undefined };
    },
  };
})();

await Promise.all([
  (async () => {
    for await (const value of values) console.log(value);
  })(),
  (async () => {
    await delay(250);
    values.return();
  })(),
]);

Answer 1

Why does the code print 0–3 rather than only 0–2? From what I can tell, when I call the return() method, the AsyncGenerator instance isn't suspended (the body execution isn't at a yield statement) and instead of returning once reaching the yield statement the next value is yielded and then suspended at which point the "return" finally happens.

Yes, precisely this is what happens. The generator is already running because the for await … of loop did call its .next() method, and so the generator will complete that before considering the .return() call.

All the methods that you invoke on an async generator are queued. (In a sync generator, you'd get a "TypeError: Generator is already running" instead). One can demonstrate this by immediately calling next multiple times:

const values = (async function*() {
  let i=0; while (true) {
    await new Promise(r => { setTimeout(r, 1000); });
    yield i++;
  }
})();
values.next().then(console.log, console.error);
values.next().then(console.log, console.error);
values.next().then(console.log, console.error);
values.return('done').then(console.log, console.error);
values.next().then(console.log, console.error);

Is there any way to detect that the return() method has been invoked and not yield afterwards?

No, not from within the generator. And really you probably still should yield the value if you already expended the effort to produce it.

It sounds like what you want to do is to ignore the produced value when you want the generator to stop. You should do that in your for await … of loop - and you can also use it to stop the generator by using a break statement:

const delay = (ms) => new Promise((resolve) => {
  setTimeout(resolve, ms);
});

async function* delayedIntegers() {
  let n = 0;
  while (true) {
    yield n++;
    await delay(1000);
  }
}

(async function main() {
  const start = Date.now();
  const values = delayedIntegers();
  for await (const value of values) {
    if (Date.now() - start > 2500) {
      console.log('done:', value);
      break;
    }
    console.log(value);
  }
})();

But if you really want to abort the generator from the outside, you need an out-of-band channel to signal the cancellation. You can use an AbortSignal for this:

const delay = (ms, signal) => new Promise((resolve, reject) => {
  function done() {
    resolve();
    signal?.removeEventListener("abort", stop);
  }
  function stop() {
    reject(this.reason);
    clearTimeout(handle);
  }
  signal?.throwIfAborted();
  const handle = setTimeout(done, ms);
  signal?.addEventListener("abort", stop, {once: true});
});

async function* delayedIntegers(signal) {
  let n = 0;
  while (true) {
    yield n++;
    await delay(1000, signal);
  }
}

(async function main() {
  try {
    const values = delayedIntegers(AbortSignal.timeout(2500));
    for await (const value of values) {
      console.log(value);
    }
  } catch(e) {
    if (e.name != "TimeoutError") throw e;
    console.log("done");
  }
})();

This will actually permit to stop the generator during the timeout, not after the full second has elapsed.

Answer 2

Is there a way to prevent this "extra yield" after invoking the return method? If not, are there libraries, patterns, etc. our there that avoid this while still implementing these AsyncIterator interface optional properties?

As @Bergi clearly explained, the extra yield cannot be avoided with the AsyncGenerator.return() method. This is a really interesting case, but I don't think you will find libraries that fix it. @Bergi proposed a clever solution using the AbortSignal, I have tried a different approach with only Promises:

(async function test() {
const delay = (ms) => new Promise((resolve) => setTimeout(resolve, ms));

const wrapIntoStoppable = function (generator) {
    const newGenerator = {
        isGeneratorStopped: false,
        resolveStopPromise: null,
        async *[Symbol.asyncIterator]() {
            let stoppedSymbol = Symbol('stoppedPromise')
            let stoppingPromise

            while (true) {
                if (this.isGeneratorStopped)
                    return
                stoppingPromise = new Promise((resolve, _) => this.resolveStopPromise = resolve)
                    .then(_ => stoppedSymbol)
                nextValuePromise = generator.next()
                const result = await Promise.race([nextValuePromise, stoppingPromise])
                this.resolveStopPromise() // resolve the promise in case it is still pending
                if (result === stoppedSymbol)
                    return
                else
                    yield result.value
            }
        },    
        stop: function() {
            this.resolveStopPromise()
            this.isGeneratorStopped = true
        }
    }

    const handler = {
        get: function(target, prop, receiver) {
            if (['next', 'return', 'throw'].includes(prop))
                return generator[prop].bind(generator)
            else
                return newGenerator[prop].bind(newGenerator)
        }
    }
  
    return new Proxy(newGenerator, handler)
}

const values = wrapIntoStoppable((async function* delayedIntegers() {
    let n = 0;
    while (true) {
        yield n++;
        await delay(100);
    }
})());

await Promise.all([
    (async () => {
        for await (const value of values) {
            console.log(Date.now())
            console.log(value);
        }
        console.log(Date.now())
        // console.log(await values.next())
        // console.log(await values.return())
        // console.log(await values.throw())
    })(),
    (async () => {
        await delay(250);
        values.stop()
    })(),
]);
})();

The idea is that I wrap an async generator with an object that has an async iterator. All the elements yielded by the wrapping generator are yielded by the original generator, but now 2 promises are started:

• nextValuePromise that will return the value to yield
• stoppingPromise that will end the iteration if resolved before the previous one

In this way, if the stop() method (which resolves stoppingPromise) is called before the first promise is resolved, then Promise.race() will immediately return a dummy Symbol. When the result of the race is this symbol, the iterator returns. The stop() function also sets the isGeneratorStopped flag that makes sure the iterator will eventually return if the stop() method is called after the stoppingPromise() is manually resolved.

I have also used a Proxy to make sure that the wrapping object behaves as a true AsyncGenerator. Calls to next(), return() or throw() are simply forwarded to the wrapped generator.

---

Let's see the pros:

1. wrapIntoStoppable can become a util method that just wraps any async generator. This is certainly convenient because you don't have to use signals every time there is a pending Promise ¹
2. Once the stop() method is called on the async generator, the for await...of loop immediately returns. Note: this doesn't mean that pending Promises are aborted

And now the cons:

1. Maybe too much code to maintain? Now the generator has a proxy that wraps another wrapper... I would like to simplify the design at least
2. After the generator is stopped, the nextValuePromise() could be resolved in the meantime, causing some potential side effects. This is the main reason why it is a pretty dangerous library function.

---

1. _{Actually, I think you could even merge @Bergi's and my solution and manage to abort a Promise when the stop() method is called. However, in this case, all the promises need to handle the abort signals.}