DRIVERS-3218 Avoid clearing the connection pool when the server connection rate limiter triggers

Author: blink1073

pymongo/asynchronous/pool.py

@@ -37,7 +37,7 @@
 from bson import DEFAULT_CODEC_OPTIONS
 from pymongo import _csot, helpers_shared
 from pymongo.asynchronous.client_session import _validate_session_write_concern
-from pymongo.asynchronous.helpers import _handle_reauth
+from pymongo.asynchronous.helpers import _backoff, _handle_reauth
 from pymongo.asynchronous.network import command
 from pymongo.common import (
     MAX_BSON_SIZE,
@@ -791,6 +791,7 @@ def __init__(
         self._max_connecting = self.opts.max_connecting
         self._pending = 0
         self._client_id = client_id
+        self._backoff = 0
         if self.enabled_for_cmap:
             assert self.opts._event_listeners is not None
             self.opts._event_listeners.publish_pool_created(
@@ -846,6 +847,8 @@ async def _reset(
         async with self.size_cond:
             if self.closed:
                 return
+            # Clear the backoff state.
+            self._backoff = 0
             if self.opts.pause_enabled and pause and not self.opts.load_balanced:
                 old_state, self.state = self.state, PoolState.PAUSED
             self.gen.inc(service_id)
@@ -937,6 +940,12 @@ async def update_is_writable(self, is_writable: Optional[bool]) -> None:
             for _socket in self.conns:
                 _socket.update_is_writable(self.is_writable)  # type: ignore[arg-type]
 
+    async def backoff(self, service_id: Optional[ObjectId] = None) -> None:
+        # Mark the pool as in backoff.
+        # TODO: how to handle load balancers?
+        self._backoff += 1
+        # TODO: emit a message.
+
     async def reset(
         self, service_id: Optional[ObjectId] = None, interrupt_connections: bool = False
     ) -> None:
@@ -994,7 +1003,8 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                 async with self._max_connecting_cond:
                     # If maxConnecting connections are already being created
                     # by this pool then try again later instead of waiting.
-                    if self._pending >= self._max_connecting:
+                    max_connecting = 1 if self._backoff else self._max_connecting
+                    if self._pending >= max_connecting:
                         return
                     self._pending += 1
                     incremented = True
@@ -1051,6 +1061,10 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
                 driverConnectionId=conn_id,
             )
 
+        # Apply backoff if applicable.
+        if self._backoff:
+            await asyncio.sleep(_backoff(self._backoff))
+
         try:
             networking_interface = await _configured_protocol_interface(self.address, self.opts)
         # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
@@ -1103,6 +1117,8 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
         if handler:
             await handler.client._topology.receive_cluster_time(conn._cluster_time)
 
+        # Clear the backoff state.
+        self._backoff = 0
         return conn
 
     @contextlib.asynccontextmanager
@@ -1279,12 +1295,13 @@ async def _get_conn(
                 # to be checked back into the pool.
                 async with self._max_connecting_cond:
                     self._raise_if_not_ready(checkout_started_time, emit_event=False)
-                    while not (self.conns or self._pending < self._max_connecting):
+                    max_connecting = 1 if self._backoff else self._max_connecting
+                    while not (self.conns or self._pending < max_connecting):
                         timeout = deadline - time.monotonic() if deadline else None
                         if not await _async_cond_wait(self._max_connecting_cond, timeout):
                             # Timed out, notify the next thread to ensure a
                             # timeout doesn't consume the condition.
-                            if self.conns or self._pending < self._max_connecting:
+                            if self.conns or self._pending < max_connecting:
                                 self._max_connecting_cond.notify()
                             emitted_event = True
                             self._raise_wait_queue_timeout(checkout_started_time)
@@ -1395,6 +1412,20 @@ async def checkin(self, conn: AsyncConnection) -> None:
                     # Pool.reset().
                     if self.stale_generation(conn.generation, conn.service_id):
                         close_conn = True
+                    # If in backoff state, check the conn's readiness.
+                    elif self._backoff:
+                        # Set a 1ms read deadline and attempt to read 1 byte from the connection.
+                        # Expect it to block for 1ms then return a deadline exceeded error. If it
+                        # returns any other error, the connection is not usable, so return false.
+                        # If it doesn't return an error and actually reads data, the connection is
+                        # also not usable, so return false.
+                        conn.conn.get_conn.settimeout(0.001)
+                        close_conn = True
+                        try:
+                            conn.conn.get_conn.read()
+                        except Exception as _:
+                            # TODO: verify the exception
+                            close_conn = False
                     else:
                         conn.update_last_checkin_time()
                         conn.update_is_writable(bool(self.is_writable))

pymongo/asynchronous/server.py

@@ -91,6 +91,10 @@ async def reset(self, service_id: Optional[ObjectId] = None) -> None:
         """Clear the connection pool."""
         await self.pool.reset(service_id)
 
+    async def backoff(self, service_id: Optional[ObjectId] = None) -> None:
+        """Set the connection pool in backoff mode."""
+        await self.pool.backoff(service_id)
+
     async def close(self) -> None:
         """Clear the connection pool and stop the monitor.
 

pymongo/asynchronous/topology.py

@@ -896,12 +896,18 @@ async def _handle_error(self, address: _Address, err_ctx: _ErrorContext) -> None
             # ... MUST NOT request an immediate check of the server."
             if not self._settings.load_balanced:
                 await self._process_change(ServerDescription(address, error=error))
-            # Clear the pool.
-            await server.reset(service_id)
-            # "When a client marks a server Unknown from `Network error when
-            # reading or writing`_, clients MUST cancel the hello check on
-            # that server and close the current monitoring connection."
-            server._monitor.cancel_check()
+
+            if err_ctx.completed_handshake:
+                # Clear the pool.
+                await server.reset(service_id)
+                # "When a client marks a server Unknown from `Network error when
+                # reading or writing`_, clients MUST cancel the hello check on
+                # that server and close the current monitoring connection."
+                server._monitor.cancel_check()
+                return
+
+            # Set the pool into backoff mode.
+            await server.backoff(service_id)
 
     async def handle_error(self, address: _Address, err_ctx: _ErrorContext) -> None:
         """Handle an application error.

pymongo/synchronous/pool.py

@@ -84,7 +84,7 @@
 from pymongo.server_type import SERVER_TYPE
 from pymongo.socket_checker import SocketChecker
 from pymongo.synchronous.client_session import _validate_session_write_concern
-from pymongo.synchronous.helpers import _handle_reauth
+from pymongo.synchronous.helpers import _backoff, _handle_reauth
 from pymongo.synchronous.network import command
 
 if TYPE_CHECKING:
@@ -789,6 +789,7 @@ def __init__(
         self._max_connecting = self.opts.max_connecting
         self._pending = 0
         self._client_id = client_id
+        self._backoff = 0
         if self.enabled_for_cmap:
             assert self.opts._event_listeners is not None
             self.opts._event_listeners.publish_pool_created(
@@ -844,6 +845,8 @@ def _reset(
         with self.size_cond:
             if self.closed:
                 return
+            # Clear the backoff state.
+            self._backoff = 0
             if self.opts.pause_enabled and pause and not self.opts.load_balanced:
                 old_state, self.state = self.state, PoolState.PAUSED
             self.gen.inc(service_id)
@@ -935,6 +938,12 @@ def update_is_writable(self, is_writable: Optional[bool]) -> None:
             for _socket in self.conns:
                 _socket.update_is_writable(self.is_writable)  # type: ignore[arg-type]
 
+    def backoff(self, service_id: Optional[ObjectId] = None) -> None:
+        # Mark the pool as in backoff.
+        # TODO: how to handle load balancers?
+        self._backoff += 1
+        # TODO: emit a message.
+
     def reset(
         self, service_id: Optional[ObjectId] = None, interrupt_connections: bool = False
     ) -> None:
@@ -990,7 +999,8 @@ def remove_stale_sockets(self, reference_generation: int) -> None:
                 with self._max_connecting_cond:
                     # If maxConnecting connections are already being created
                     # by this pool then try again later instead of waiting.
-                    if self._pending >= self._max_connecting:
+                    max_connecting = 1 if self._backoff else self._max_connecting
+                    if self._pending >= max_connecting:
                         return
                     self._pending += 1
                     incremented = True
@@ -1047,6 +1057,10 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
                 driverConnectionId=conn_id,
             )
 
+        # Apply backoff if applicable.
+        if self._backoff:
+            asyncio.sleep(_backoff(self._backoff))
+
         try:
             networking_interface = _configured_socket_interface(self.address, self.opts)
         # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
@@ -1099,6 +1113,8 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
         if handler:
             handler.client._topology.receive_cluster_time(conn._cluster_time)
 
+        # Clear the backoff state.
+        self._backoff = 0
         return conn
 
     @contextlib.contextmanager
@@ -1275,12 +1291,13 @@ def _get_conn(
                 # to be checked back into the pool.
                 with self._max_connecting_cond:
                     self._raise_if_not_ready(checkout_started_time, emit_event=False)
-                    while not (self.conns or self._pending < self._max_connecting):
+                    max_connecting = 1 if self._backoff else self._max_connecting
+                    while not (self.conns or self._pending < max_connecting):
                         timeout = deadline - time.monotonic() if deadline else None
                         if not _cond_wait(self._max_connecting_cond, timeout):
                             # Timed out, notify the next thread to ensure a
                             # timeout doesn't consume the condition.
-                            if self.conns or self._pending < self._max_connecting:
+                            if self.conns or self._pending < max_connecting:
                                 self._max_connecting_cond.notify()
                             emitted_event = True
                             self._raise_wait_queue_timeout(checkout_started_time)
@@ -1391,6 +1408,20 @@ def checkin(self, conn: Connection) -> None:
                     # Pool.reset().
                     if self.stale_generation(conn.generation, conn.service_id):
                         close_conn = True
+                    # If in backoff state, check the conn's readiness.
+                    elif self._backoff:
+                        # Set a 1ms read deadline and attempt to read 1 byte from the connection.
+                        # Expect it to block for 1ms then return a deadline exceeded error. If it
+                        # returns any other error, the connection is not usable, so return false.
+                        # If it doesn't return an error and actually reads data, the connection is
+                        # also not usable, so return false.
+                        conn.conn.get_conn.settimeout(0.001)
+                        close_conn = True
+                        try:
+                            conn.conn.get_conn.read()
+                        except Exception as _:
+                            # TODO: verify the exception
+                            close_conn = False
                     else:
                         conn.update_last_checkin_time()
                         conn.update_is_writable(bool(self.is_writable))

pymongo/synchronous/server.py

@@ -91,6 +91,10 @@ def reset(self, service_id: Optional[ObjectId] = None) -> None:
         """Clear the connection pool."""
         self.pool.reset(service_id)
 
+    def backoff(self, service_id: Optional[ObjectId] = None) -> None:
+        """Set the connection pool in backoff mode."""
+        self.pool.backoff(service_id)
+
     def close(self) -> None:
         """Clear the connection pool and stop the monitor.
 

pymongo/synchronous/topology.py

@@ -894,12 +894,18 @@ def _handle_error(self, address: _Address, err_ctx: _ErrorContext) -> None:
             # ... MUST NOT request an immediate check of the server."
             if not self._settings.load_balanced:
                 self._process_change(ServerDescription(address, error=error))
-            # Clear the pool.
-            server.reset(service_id)
-            # "When a client marks a server Unknown from `Network error when
-            # reading or writing`_, clients MUST cancel the hello check on
-            # that server and close the current monitoring connection."
-            server._monitor.cancel_check()
+
+            if err_ctx.completed_handshake:
+                # Clear the pool.
+                server.reset(service_id)
+                # "When a client marks a server Unknown from `Network error when
+                # reading or writing`_, clients MUST cancel the hello check on
+                # that server and close the current monitoring connection."
+                server._monitor.cancel_check()
+                return
+
+            # Set the pool into backoff mode.
+            server.backoff(service_id)
 
     def handle_error(self, address: _Address, err_ctx: _ErrorContext) -> None:
         """Handle an application error.