Repository navigation
fix: preserve task identity in threaded query execution - #952
Ian Hellen (ianhelle) merged 4 commits into
Conversation
_exec_queries_threaded reconstructed which result belonged to which
query_id by zipping thread_tasks (submission-order keys) against both
an intermediate task iterator and, later, the results list
(completion order, from asyncio.as_completed). Those two sequences
aren't the same order, so the zip silently mispaired query_ids with
results and exceptions whenever tasks didn't complete in submission
order -- the common case under real concurrency.
Two concrete failure modes from this:
- The final concatenated DataFrame's row order ("Sort the results by
the order of the tasks") was actually ordered by sorting already-
mismatched (wrong_id, result) pairs, which only incidentally matched
submission order when alphabetical id order and submission order
happened to agree.
- Worse, which query_id got logged as "failed" and (with retry=True)
actually retried was also determined by this same mismatch, not by
which task genuinely failed. A fast, never-failing query could be
logged and retried under a different id, while the real failure's
result silently never makes it into the output -- a correctness bug
on the retry path, not just an ordering annoyance.
Fix: tag each task with its own query_id before running it through
asyncio.as_completed via a small _run_tagged_task helper, so identity
travels with the outcome instead of being reconstructed afterwards by
position. Results are now collected into a dict keyed by query_id;
final ordering is rebuilt by iterating the original submission-order
keys, and retry correctly targets the query_id that actually failed.
Added two regression tests to test_async_queries.py:
- test_exec_queries_threaded_preserves_identity_out_of_order
- test_exec_queries_threaded_retries_the_task_that_actually_failed
Confirmed via git stash against unmodified main that the second test
fails on the pre-fix code: it retries the wrong query_id (the one
that never failed) and silently drops the one that actually did,
giving a final result of two copies of the never-failing query's data
and zero copies of the one that failed and was "retried".
Ian Hellen (ianhelle)
left a comment
Ian Hellen (ianhelle)
left a comment
There was a problem hiding this comment.
The task-identity fix looks sound, including selecting the correct query for retry. Please also preserve submission order when a retry succeeds and add an assertion covering that case, as described inline. This ordering gap predates the PR; it is not a newly introduced security vulnerability, but it remains inconsistent with the ordering guarantee this change establishes. The current head has only the CLA check, so a full CI run is also needed before merging.
ianhelle's review on microsoft#952 found a residual ordering gap: ordered_results was built from the first-attempt results dict before the retry pass ran, so a retried task's result was always appended at the end regardless of its real submission position. For tasks submitted as a, b where a fails initially and succeeds on retry while b succeeds immediately, the output was b, a instead of a, b. Extracted the per-batch run loop (submit tagged tasks, await via as_completed, bucket into a query_id-keyed results dict plus a failed-id list) into a new _run_tasks_threaded helper, used for both the first attempt and the retry pass. The retry pass's results are merged back into the same results dict via results.update(retry_results) before the final submission-order reconstruction runs, so a query that only succeeded on retry now lands in its correct position instead of after every query that succeeded on the first attempt. This replaces the previous retry path, which recursively called _exec_queries_threaded(retry=False) and appended its already-concatenated return value to the end of ordered_results -- that's exactly what threw away per-query identity for the retried batch and made correct interleaving impossible without this refactor. Updated test_exec_queries_threaded_retries_the_task_that_actually_failed per review: replaced the sorted(result["query_id"]) assertion (which can't distinguish correct order from the bug) with an exact a, b order assertion. Confirmed this fails on the previous commit's code -- order comes back as b, a -- and passes with this fix. tests/data/test_async_queries.py: 6/6 pass. tests/data/: 383 passed, 5 skipped (pre-existing, unrelated). black/pylint/mypy clean on the changed lines (the file has pre-existing lines that reformat under my locally installed black version due to version drift, same as noted in the original PR -- none of them touched by this change).
|
Good catch — fixed in 75fe1a6. Extracted the per-batch run loop into a shared Updated Also merged |
But your original fix was a good catch. Thanks so much for not just finding the bug, but doing something about it :-) |
2 participants
_exec_queries_threadedreconstructed which result belonged to whichquery_idby zippingthread_tasks(submission-order keys) against both an intermediate task iterator and, later, the results list (completion order, fromasyncio.as_completed). Those two sequences aren't the same order, so the zip silently mispairsquery_ids with results and exceptions whenever tasks don't complete in submission order — the common case under real concurrency.Two concrete failure modes from this:
(wrong_id, result)pairs — only incidentally correct when alphabetical id order happened to match submission order.query_idgot logged as "failed" and (withretry=True) actually retried was determined by this same mismatch, not by which task genuinely failed. A fast, never-failing query could get logged and retried under a different id, while the real failure's result silently never makes it into the output.Fix: tag each task with its own
query_idbefore running it throughasyncio.as_completed, via a small_run_tagged_taskhelper — identity travels with the outcome instead of being reconstructed by position afterward. Results are collected into a dict keyed byquery_id; final ordering is rebuilt by iterating the original submission keys, and retry now correctly targets thequery_idthat actually failed.Tests: added
test_exec_queries_threaded_preserves_identity_out_of_orderandtest_exec_queries_threaded_retries_the_task_that_actually_failedtotests/data/test_async_queries.py, using controlled completion-order delays to deterministically force the out-of-order scenario (not relying on real thread-scheduling timing). Confirmed viagit stashagainst unmodifiedmainthat the retry test fails on pre-fix code — it retries the wrongquery_idand silently drops the one that actually failed, producing two copies of the never-failing query's data and zero copies of the one that failed.Full
tests/data/test_async_queries.py: 6/6 pass.black --check,pylint(errors-only),mypy: all clean on the changed file.