fix(checkpoint): fix AsyncBatchedBaseStore getting stuck

This commit fixes #5503

Gist of it is:

- `asyncio.exception.InvalidStateError` were being raised when the
  future was cancelled
- this exception bubbled up and killed the background task
- `AsyncBatchedBaseStore` stopped doing queries because the background
  task wasn't running anymore

This commit adds some "if future is not done" checks to guard against
this.
This commit is contained in:
Darren Clark
2025-07-14 18:03:51 -04:00
parent a5ce13eb45
commit 0d8dfa7bba
2 changed files with 43 additions and 2 deletions
@@ -343,10 +343,14 @@ async def _run(
# set the results of each operation
for fut, result in zip(futs, results):
fut.set_result(result)
# guard against future being done (e.g. cancelled)
if not fut.done():
fut.set_result(result)
except Exception as e:
for fut in futs:
fut.set_exception(e)
# guard against future being done (e.g. cancelled)
if not fut.done():
fut.set_exception(e)
finally:
# remove strong ref to store
del s