fix: capture init task ref + cancel on close for Neo4jDriver and FalkorDriver#1636
fix: capture init task ref + cancel on close for Neo4jDriver and FalkorDriver#1636nandanadileep wants to merge 3 commits into
Conversation
…orDriver Both Neo4jDriver and FalkorDriver __init__ scheduled build_indices_and_constraints() via loop.create_task() without capturing the task reference. The orphan task could: 1. Race with explicit build_indices_and_constraints() calls 2. Produce 'Task exception was never retrieved' warnings on failure 3. Leave the connection in a mid-transaction state when the driver is closed concurrently Now captures the task in self._init_task and cancels it in close(), with suppress(CancelledError) to handle the cancellation cleanly. Closes getzep#1513
|
All contributors have signed the CLA ✍️ ✅ |
|
I have read the CLA Document and I hereby sign the CLA behalf on myself, e-mail: nandanadileep29@gmail.com |
|
@nandanadileep Thanks for catching this! One gap: close() skips done() tasks, so a failed init task still leaks its exception. Calling task.exception() on done tasks too would complete it. |
Previously, close() only cancelled and awaited tasks that were still running, but skipped already-done tasks. If a done task failed with an exception, the exception was never retrieved, causing 'Task exception was never retrieved' warnings. Now calls task.exception() on done tasks to consume any pending exception and complete the task's lifecycle.
prasmussen15
left a comment
There was a problem hiding this comment.
Please update close() to retrieve exceptions from completed init tasks too; skipping done() tasks still leaves a failed init task exception unobserved.
Task.exception() raises CancelledError if the task finished in a cancelled state. If the init task was cancelled externally before close() ran, the else branch would re-raise CancelledError out of close(), breaking driver cleanup. Skip exception retrieval for cancelled tasks; only call exception() on tasks that completed normally or failed.
|
@prasmussen15 Thanks for the review! The requested change was already landed in d3fc190 (pushed Jul 20) — if self._init_task is not None:
if not self._init_task.done():
self._init_task.cancel()
with suppress(asyncio.CancelledError):
await self._init_task
elif not self._init_task.cancelled():
self._init_task.exception()The review may have been against a stale diff. That said, it surfaced a related edge case: |
Summary
Both
Neo4jDriver.__init__andFalkorDriver.__init__schedulebuild_indices_and_constraints()vialoop.create_task()without capturing the task reference. This creates a fire-and-forget background task whose lifecycle is not managed.Problem
await driver.build_indices_and_constraints()race with the constructor-fired task, producingEquivalentSchemaRuleAlreadyExistserrors.Fix
self._init_task: asyncio.Task | Noneclose()\, cancel the task if still running and await it withsuppress(asyncio.CancelledError)import asyncioto the top ofneo4j_driver.pyTesting
make lintpassesmake typecheckpassesDISABLE_FALKORDB=1 DISABLE_KUZU=1 DISABLE_NEPTUNE=1 uv run pytest -m 'not integration'— 165 passed, 2 skipped (FalkorDB integration), 1 pre-existing Neo4j auth error (no local Neo4j)Closes #1513