Fix race condition in async connection initialization

Add asyncio.Lock to AsyncConnection._ensure_connection() to prevent
concurrent callers from getting a half-initialized connection. Set
_async_conn only after connect() succeeds. Add threading.Lock to
Connection._ensure_loop() for the sync wrapper.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: antoineleclair

src/dqlitedbapi/aio/connection.py

@@ -1,5 +1,6 @@
 """Async connection implementation for dqlite."""
 
+import asyncio
 from typing import Any
 
 from dqliteclient import DqliteConnection
@@ -29,24 +30,33 @@ def __init__(
         self._timeout = timeout
         self._async_conn: DqliteConnection | None = None
         self._closed = False
+        self._connect_lock = asyncio.Lock()
 
     async def _ensure_connection(self) -> DqliteConnection:
         """Ensure the underlying connection is established."""
         if self._closed:
             raise InterfaceError("Connection is closed")
 
-        if self._async_conn is None:
-            self._async_conn = DqliteConnection(
+        if self._async_conn is not None:
+            return self._async_conn
+
+        async with self._connect_lock:
+            # Double-check after acquiring lock
+            if self._async_conn is not None:
+                return self._async_conn
+
+            conn = DqliteConnection(
                 self._address,
                 database=self._database,
                 timeout=self._timeout,
             )
             try:
-                await self._async_conn.connect()
+                await conn.connect()
             except Exception as e:
-                self._async_conn = None
                 raise OperationalError(f"Failed to connect: {e}") from e
 
+            self._async_conn = conn
+
         return self._async_conn
 
     async def connect(self) -> None:

src/dqlitedbapi/connection.py

@@ -33,17 +33,21 @@ def __init__(
         self._closed = False
         self._loop: asyncio.AbstractEventLoop | None = None
         self._thread: threading.Thread | None = None
+        self._loop_lock = threading.Lock()
 
     def _ensure_loop(self) -> asyncio.AbstractEventLoop:
         """Ensure a dedicated event loop is running in a background thread.
 
         This allows sync methods to work even when called from within
         an already-running async context (e.g. uvicorn).
         """
-        if self._loop is None or self._loop.is_closed():
-            self._loop = asyncio.new_event_loop()
-            self._thread = threading.Thread(target=self._loop.run_forever, daemon=True)
-            self._thread.start()
+        if self._loop is not None and not self._loop.is_closed():
+            return self._loop
+        with self._loop_lock:
+            if self._loop is None or self._loop.is_closed():
+                self._loop = asyncio.new_event_loop()
+                self._thread = threading.Thread(target=self._loop.run_forever, daemon=True)
+                self._thread.start()
         return self._loop
 
     def _run_sync(self, coro: Any) -> Any:
@@ -62,16 +66,16 @@ async def _get_async_connection(self) -> DqliteConnection:
             raise InterfaceError("Connection is closed")
 
         if self._async_conn is None:
-            self._async_conn = DqliteConnection(
+            conn = DqliteConnection(
                 self._address,
                 database=self._database,
                 timeout=self._timeout,
             )
             try:
-                await self._async_conn.connect()
+                await conn.connect()
             except Exception as e:
-                self._async_conn = None
                 raise OperationalError(f"Failed to connect: {e}") from e
+            self._async_conn = conn
 
         return self._async_conn
 

tests/test_async_connection_race.py

@@ -0,0 +1,47 @@
+"""Tests for race condition in async connection initialization."""
+
+import asyncio
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from dqlitedbapi.aio.connection import AsyncConnection
+
+
+class TestAsyncConnectionRace:
+    @pytest.mark.asyncio
+    async def test_concurrent_ensure_connection_waits_for_connect(self) -> None:
+        """A second _ensure_connection() call must not return before connect() finishes."""
+        conn = AsyncConnection("localhost:9001")
+
+        connect_started = asyncio.Event()
+        connect_finished = asyncio.Event()
+
+        async def slow_connect() -> None:
+            connect_started.set()
+            await asyncio.sleep(0.1)  # Simulate slow TCP handshake
+            connect_finished.set()
+
+        with patch("dqlitedbapi.aio.connection.DqliteConnection") as MockDqliteConn:
+            mock_instance = AsyncMock()
+            mock_instance.connect = slow_connect
+            mock_instance._protocol = AsyncMock()
+            mock_instance._db_id = 0
+            MockDqliteConn.return_value = mock_instance
+
+            async def second_caller() -> bool:
+                # Wait for first caller to start connecting
+                await connect_started.wait()
+                # Now call _ensure_connection — it should wait for connect to finish
+                await conn._ensure_connection()
+                # At this point, connect must have finished
+                return connect_finished.is_set()
+
+            first_task = asyncio.create_task(conn._ensure_connection())
+            second_task = asyncio.create_task(second_caller())
+
+            await first_task
+            connect_was_finished = await second_task
+
+            # The second caller must have seen connect_finished=True
+            assert connect_was_finished, "Second caller got connection before connect() finished"