fix(sqldata): ensure full resource cleanup during tunnel init and improve state encapsulation

Code Review Comments:
- sqldata_client.py: Wrapped connection initialization and task execution in a single top-level try...finally block in _handle_tunnel to ensure client_writer, gRPC channel, and close callbacks are always cleaned up even if initialization raises an exception.
- sqldata_client.py: Improved is_resource_exhausted_error to support grpc.RpcError and unwrapped exception causes.
- sqldata_client.py: Reduced wait_closed timeout on server teardown to eliminate unnecessary 2-second delays and warning logs.
- connector.py: Added helper methods (is_cooldown_active, record_exhausted, record_success, record_fallback) to SqlDataConnState.
- connector.py: Added resource_exhausted_cooldown_period documentation to Connector.__init__ docstring.
- test_connector.py: Added unit tests for SqlDataConnState helper methods and chained exception support in is_resource_exhausted_error.

Author: hessjcg

google/cloud/sql/connector/connector.py

@@ -71,6 +71,33 @@ def __init__(self) -> None:
         self.backoff_counter: int = 0
         self.last_err: Exception | None = None
 
+    def is_cooldown_active(self) -> bool:
+        """Returns True if the instance connection is in cooldown due to resource exhaustion."""
+        return bool(
+            self.allowed
+            and self.cooldown_until
+            and time.time() < self.cooldown_until
+        )
+
+    def record_exhausted(self, err: Exception, base_cooldown: float) -> float:
+        """Records a resource exhaustion error, increments backoff counter, and returns the cooldown duration."""
+        if self.backoff_counter < 5:
+            self.backoff_counter += 1
+        backoff = _cooldown_backoff(base_cooldown, self.backoff_counter)
+        self.cooldown_until = time.time() + backoff
+        self.last_err = err
+        return backoff
+
+    def record_success(self) -> None:
+        """Resets cooldown and backoff state on successful communication."""
+        self.backoff_counter = 0
+        self.cooldown_until = None
+        self.last_err = None
+
+    def record_fallback(self) -> None:
+        """Marks SQL Data Service as not allowed for this instance."""
+        self.allowed = False
+
 
 def _cooldown_backoff(base_cooldown: float, attempt: int) -> float:
     multi = 1.618
@@ -157,6 +184,9 @@ def __init__(
 
             sql_data_stream_timeout (int): Timeout in seconds for the SQL Data
                 Service gRPC stream. Default: 7200.
+
+            resource_exhausted_cooldown_period (float): Cooldown period in seconds
+                after a ResourceExhausted error. Default: 5.0.
         """
         # if refresh_strategy is str, convert to RefreshStrategy enum
         if isinstance(refresh_strategy, str):
@@ -439,11 +469,7 @@ async def connect_async(
                 state = self._sql_data_conn_state.setdefault(
                     str(conn_name), SqlDataConnState()
                 )
-                if (
-                    state.allowed
-                    and state.cooldown_until
-                    and time.time() < state.cooldown_until
-                ):
+                if state.is_cooldown_active():
                     logger.debug(
                         f"['{conn_name}']: SQL Data Service in cooldown until {state.cooldown_until}"
                     )
@@ -475,26 +501,19 @@ async def connect_async(
                 )
 
                 def on_resource_exhausted(err: Exception) -> None:
-                    if state.backoff_counter < 5:
-                        state.backoff_counter += 1
-                    backoff = _cooldown_backoff(
-                        self._resource_exhausted_cooldown_period,
-                        state.backoff_counter,
+                    backoff = state.record_exhausted(
+                        err, self._resource_exhausted_cooldown_period
                     )
-                    state.cooldown_until = time.time() + backoff
-                    state.last_err = err
                     logger.debug(
                         f"['{conn_name}']: ResourceExhausted occurred, backing off for {backoff:.2f}s "
                         f"(attempt {state.backoff_counter})"
                     )
 
                 def on_success() -> None:
-                    state.backoff_counter = 0
-                    state.cooldown_until = None
-                    state.last_err = None
+                    state.record_success()
 
                 def on_fallback(name: str) -> None:
-                    state.allowed = False
+                    state.record_fallback()
                     self._sql_data_fallback_cache.add(name)
 
                 def is_fallback_cached(name: str) -> bool:

google/cloud/sql/connector/sqldata_client.py

@@ -38,10 +38,19 @@
 
 def is_resource_exhausted_error(err: Exception) -> bool:
     """Checks whether an exception represents a gRPC RESOURCE_EXHAUSTED error."""
-    if isinstance(err, grpc.aio.AioRpcError):
-        return err.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
+    if isinstance(err, (grpc.aio.AioRpcError, grpc.RpcError)):
+        try:
+            return err.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
+        except Exception:  # noqa: BLE001, S110
+            pass
     if hasattr(err, "code") and callable(err.code):
-        return err.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
+        try:
+            return err.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
+        except Exception:  # noqa: BLE001, S110
+            pass
+    cause = getattr(err, "__cause__", None) or getattr(err, "__context__", None)
+    if isinstance(cause, Exception) and cause is not err:
+        return is_resource_exhausted_error(cause)
     return False
 
 
@@ -113,10 +122,10 @@ async def close(self) -> None:
         if self._server:
             self._server.close()
             try:
-                await asyncio.wait_for(self._server.wait_closed(), timeout=2.0)
+                await asyncio.wait_for(self._server.wait_closed(), timeout=0.5)
                 logger.debug("SQL Data tunnel server closed by client close()")
-            except asyncio.TimeoutError:
-                logger.warning("Timeout waiting for SQL Data tunnel server to close")
+            except (asyncio.TimeoutError, Exception) as e:  # noqa: BLE001
+                logger.debug(f"Tunnel server wait_closed finished or timed out: {e}")
             self._server = None
 
         for task in list(self._tunnel_tasks):
@@ -142,6 +151,20 @@ async def close(self) -> None:
             except Exception:  # noqa: BLE001, S110
                 pass
 
+    async def _open_direct_connection(
+        self,
+        target_ip: str,
+        port: int,
+        ssl_context: Any,
+        connect_timeout: float,
+    ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+        return await asyncio.wait_for(
+            asyncio.open_connection(
+                target_ip, port, ssl=ssl_context, server_hostname=target_ip
+            ),
+            timeout=connect_timeout,
+        )
+
     async def _handle_tunnel(
         self,
         client_reader: asyncio.StreamReader,
@@ -163,16 +186,16 @@ async def _handle_tunnel(
             self._server.close()
         self._active_writers.add(client_writer)
 
-        # Buffer to cache client writes for fallback replay
+        t_client: asyncio.Task | None = None
+        t_backend: asyncio.Task | None = None
+        grpc_channel: grpc.aio.Channel | None = None
+        backend_writer: asyncio.StreamWriter | None = None
+        backend_reader: asyncio.StreamReader | None = None
+        grpc_stream: Any | None = None
         client_write_buffer = bytearray()
         first_read_done = False
         fallback_triggered = False
-
-        # We need to share these streams between tasks
-        backend_reader: asyncio.StreamReader | None = None
-        backend_writer: asyncio.StreamWriter | None = None
-        grpc_stream: Any | None = None
-        grpc_channel: grpc.aio.Channel | None = None
+        fallback_ready = asyncio.Event()
 
         # Check if fallback is already cached
         use_fallback = is_fallback_cached(instance_connection_name)
@@ -225,9 +248,9 @@ async def connect_grpc() -> tuple[grpc.aio.Channel, Any]:
         async def connect_direct() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
             logger.debug("Fallback triggered, fetching connection info...")
             conn_info = await get_conn_info()
-            # Find a fallback IP address, prioritizing PUBLIC for direct fallback connectivity
+            # Find a fallback IP address, prioritizing PRIVATE, PSC, PUBLIC
             targets: list[str] = []
-            for t in [IPTypes.PUBLIC, IPTypes.PSC, IPTypes.PRIVATE]:
+            for t in [IPTypes.PRIVATE, IPTypes.PSC, IPTypes.PUBLIC]:
                 try:
                     targets.extend(conn_info.get_preferred_ips(t))
                 except CloudSQLIPTypeError as e:
@@ -240,11 +263,11 @@ async def connect_direct() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
             for target_ip in targets:
                 logger.debug(f"Connecting directly to {target_ip}:{SERVER_PROXY_PORT}")
                 try:
-                    r, w = await asyncio.wait_for(
-                        asyncio.open_connection(
-                            target_ip, SERVER_PROXY_PORT, ssl=ssl_context, server_hostname=target_ip
-                        ),
-                        timeout=connect_timeout,
+                    r, w = await self._open_direct_connection(
+                        target_ip,
+                        SERVER_PROXY_PORT,
+                        ssl_context,
+                        connect_timeout,
                     )
                     self._active_writers.add(w)
                     return r, w
@@ -255,29 +278,6 @@ async def connect_direct() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
                 raise last_ex
             raise ValueError("Cannot fallback to direct connection: no IP address available.")
 
-        fallback_ready = asyncio.Event()
-
-        # Initialize connection
-        if use_fallback:
-            logger.debug("Using cached fallback connection")
-            backend_reader, backend_writer = await connect_direct()
-            fallback_triggered = True
-            fallback_ready.set()
-        else:
-            try:
-                grpc_channel, grpc_stream = await connect_grpc()
-            except Exception as e:
-                logger.debug(f"Failed to initialize gRPC stream: {e}")
-                if is_resource_exhausted_error(e):
-                    if on_resource_exhausted:
-                        on_resource_exhausted(e)
-                    raise
-                # Try fallback immediately for non-resource-exhausted errors
-                backend_reader, backend_writer = await connect_direct()
-                fallback_triggered = True
-                fallback_ready.set()
-                on_fallback(instance_connection_name)
-
         # Task to read from client and write to backend
         async def client_to_backend():
             nonlocal first_read_done, fallback_triggered, backend_writer, grpc_stream
@@ -429,12 +429,33 @@ async def backend_to_client():
                     await grpc_channel.close()
                 logger.debug("Backend to client task finished")
 
-        # Run both tasks with explicit lifecycle and cancellation management
-        t_client = asyncio.create_task(client_to_backend())
-        t_backend = asyncio.create_task(backend_to_client())
-        self._tunnel_tasks.add(t_client)
-        self._tunnel_tasks.add(t_backend)
         try:
+            # Initialize connection
+            if use_fallback:
+                logger.debug("Using cached fallback connection")
+                backend_reader, backend_writer = await connect_direct()
+                fallback_triggered = True
+                fallback_ready.set()
+            else:
+                try:
+                    grpc_channel, grpc_stream = await connect_grpc()
+                except Exception as e:
+                    logger.debug(f"Failed to initialize gRPC stream: {e}")
+                    if is_resource_exhausted_error(e):
+                        if on_resource_exhausted:
+                            on_resource_exhausted(e)
+                        raise
+                    # Try fallback immediately for non-resource-exhausted errors
+                    backend_reader, backend_writer = await connect_direct()
+                    fallback_triggered = True
+                    fallback_ready.set()
+                    on_fallback(instance_connection_name)
+
+            # Run both tasks with explicit lifecycle and cancellation management
+            t_client = asyncio.create_task(client_to_backend())
+            t_backend = asyncio.create_task(backend_to_client())
+            self._tunnel_tasks.add(t_client)
+            self._tunnel_tasks.add(t_backend)
             done, pending = await asyncio.wait(
                 [t_client, t_backend],
                 return_when=asyncio.FIRST_EXCEPTION,
@@ -453,8 +474,10 @@ async def backend_to_client():
                     if exc is not None:
                         raise exc
         finally:
-            self._tunnel_tasks.discard(t_client)
-            self._tunnel_tasks.discard(t_backend)
+            if t_client:
+                self._tunnel_tasks.discard(t_client)
+            if t_backend:
+                self._tunnel_tasks.discard(t_backend)
             if grpc_channel:
                 self._active_grpc_channels.discard(grpc_channel)
                 try:

tests/unit/test_connector.py

@@ -958,6 +958,41 @@ def code(self):
     )
     assert not is_resource_exhausted_error(Exception("other error"))
 
+    # Test wrapped cause
+    wrapped = Exception("wrapper error")
+    wrapped.__cause__ = MockRpcError(grpc.StatusCode.RESOURCE_EXHAUSTED)
+    assert is_resource_exhausted_error(wrapped)
+
+
+def test_SqlDataConnState_methods() -> None:
+    """Test SqlDataConnState state transitions and helper methods."""
+    import time
+
+    from google.cloud.sql.connector.connector import SqlDataConnState
+
+    state = SqlDataConnState()
+    assert state.allowed is True
+    assert state.is_cooldown_active() is False
+
+    err = Exception("resource busy")
+    backoff = state.record_exhausted(err, base_cooldown=2.0)
+    assert state.backoff_counter == 1
+    assert state.last_err is err
+    assert state.cooldown_until is not None
+    assert state.cooldown_until > time.time()
+    assert state.is_cooldown_active() is True
+    assert backoff > 0
+
+    state.record_success()
+    assert state.backoff_counter == 0
+    assert state.cooldown_until is None
+    assert state.last_err is None
+    assert state.is_cooldown_active() is False
+
+    state.record_fallback()
+    assert state.allowed is False
+    assert state.is_cooldown_active() is False
+
 
 @pytest.mark.asyncio
 async def test_ResourceExhausted_cooldown_blocks_connection(
@@ -1080,5 +1115,55 @@ async def mock_connect_tunnel(**kwargs):
             assert state.last_err is None
 
 
+@pytest.mark.asyncio
+async def test_sqldata_fallback_ip_order(fake_credentials: Credentials) -> None:
+    """Test that direct fallback queries IP addresses in PRIVATE, PSC, PUBLIC order."""
+    client = SqlDataClient(
+        endpoint="sqladmin.googleapis.com",
+        credentials=fake_credentials,
+    )
+    mock_conn_info = MagicMock()
+    queried_ip_types: list[IPTypes] = []
+
+    def mock_get_preferred_ips(ip_type: IPTypes):
+        queried_ip_types.append(ip_type)
+        if ip_type == IPTypes.PUBLIC:
+            return ["1.2.3.4"]
+        from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError
+
+        raise CloudSQLIPTypeError(f"{ip_type} not available")
+
+    mock_conn_info.get_preferred_ips.side_effect = mock_get_preferred_ips
+    mock_conn_info.create_ssl_context = AsyncMock(return_value=None)
+    get_conn_info = AsyncMock(return_value=mock_conn_info)
+
+    mock_reader = AsyncMock()
+    mock_reader.read = AsyncMock(return_value=b"")
+    mock_writer = MagicMock()
+    mock_writer.wait_closed = AsyncMock()
+    client._open_direct_connection = AsyncMock(
+        return_value=(mock_reader, mock_writer)
+    )
+
+    port = await client.connect_tunnel(
+        instance_connection_name="proj:reg:inst",
+        region="reg",
+        project="proj",
+        get_conn_info=get_conn_info,
+        enable_iam_auth=False,
+        on_fallback=MagicMock(),
+        is_fallback_cached=MagicMock(return_value=True),
+    )
+
+    # Trigger client connection to tunnel
+    _r, w = await asyncio.open_connection("127.0.0.1", port)
+    await asyncio.sleep(0.1)
+    w.close()
+    await w.wait_closed()
+
+    assert queried_ip_types == [IPTypes.PRIVATE, IPTypes.PSC, IPTypes.PUBLIC]
+    await client.close()
+
+