wip

Author: blink1073

pymongo/asynchronous/pool.py

@@ -1029,6 +1029,16 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                     self.requests -= 1
                     self.size_cond.notify()
 
+    def _handle_connection_error(self, error: Exception, phase: str) -> None:
+        # Handle system overload condition.  When the base AutoReconnect is
+        # raised and we are not an sdam pool, add to backoff and add the
+        # appropriate error label.
+        if not self.is_sdam and type(error) == AutoReconnect:
+            self._backoff += 1
+            error._add_error_label("SystemOverloaded")
+            error._add_error_label("Retryable")
+            print(f"Setting backoff in {phase}:", self._backoff)  # noqa: T201
+
     async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> AsyncConnection:
         """Connect to Mongo and return a new AsyncConnection.
 
@@ -1066,7 +1076,6 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
             networking_interface = await _configured_protocol_interface(self.address, self.opts)
         # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
-            print("Got the TLS handshake error")  # noqa: T201
             async with self.lock:
                 self.active_contexts.discard(tmp_context)
             if self.enabled_for_cmap:
@@ -1085,10 +1094,10 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
                     reason=_verbose_connection_error_reason(ConnectionClosedReason.ERROR),
                     error=ConnectionClosedReason.ERROR,
                 )
+            self._handle_connection_error(error, "handshake")
             if isinstance(error, (IOError, OSError, *SSLErrors)):
                 details = _get_timeout_details(self.opts)
                 _raise_connection_failure(self.address, error, timeout_details=details)
-
             raise
 
         conn = AsyncConnection(networking_interface, self, self.address, conn_id, self.is_sdam)  # type: ignore[arg-type]
@@ -1109,14 +1118,7 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
         except BaseException as e:
             async with self.lock:
                 self.active_contexts.discard(conn.cancel_context)
-            # Handle system overload condition.  When the base AutoReconnect is
-            # raised and we are not an sdam pool, add to backoff and add the
-            # appropriate error label.
-            if not self.is_sdam and type(e) == AutoReconnect:
-                self._backoff += 1
-                e._add_error_label("SystemOverloaded")
-                e._add_error_label("Retryable")
-                print("Setting backoff:", self._backoff)  # noqa: T201
+            self._handle_connection_error(e, "hello")
             await conn.close_conn(ConnectionClosedReason.ERROR)
             raise
 

pymongo/network_layer.py

@@ -254,7 +254,7 @@ class PyMongoBaseProtocol(Protocol):
     def __init__(self, timeout: Optional[float] = None):
         self.transport: Transport = None  # type: ignore[assignment]
         self._timeout = timeout
-        self._connection_made = asyncio.get_running_loop().create_future()
+        self._closing_error = asyncio.get_running_loop().create_future()
         self._closed = asyncio.get_running_loop().create_future()
         self._connection_lost = False
 
@@ -271,13 +271,9 @@ def close(self, exc: Optional[Exception] = None) -> None:
         self._resolve_pending(exc)
         self._connection_lost = True
 
-    def connection_made(self, transport: BaseTransport) -> None:
-        super().connection_made(transport)
-        self._connection_made.set_result(None)
-
     def connection_lost(self, exc: Optional[Exception] = None) -> None:
-        if exc is not None and not self._connection_made.done():
-            self._connection_made.set_exception(exc)
+        if exc is not None and not self._closing_error.done():
+            self._closing_error.set_exception(exc)
         self._resolve_pending(exc)
         if not self._closed.done():
             self._closed.set_result(None)
@@ -344,7 +340,7 @@ async def read(self, request_id: Optional[int], max_message_size: int) -> tuple[
             message = await self._done_messages.popleft()
         else:
             if self.transport and self.transport.is_closing():
-                raise OSError("connection is already closed")
+                return await self._closing_error
             read_waiter = asyncio.get_running_loop().create_future()
             self._pending_messages.append(read_waiter)
             try:

pymongo/pool_shared.py

@@ -278,7 +278,6 @@ async def _configured_protocol_interface(
             server_hostname=host,
             ssl=ssl_context,
         )
-        await protocol._connection_made
     except _CertificateError:
         # Raise _CertificateError directly like we do after match_hostname
         # below.

pymongo/synchronous/pool.py

@@ -1025,6 +1025,16 @@ def remove_stale_sockets(self, reference_generation: int) -> None:
                     self.requests -= 1
                     self.size_cond.notify()
 
+    def _handle_connection_error(self, error: Exception, phase: str) -> None:
+        # Handle system overload condition.  When the base AutoReconnect is
+        # raised and we are not an sdam pool, add to backoff and add the
+        # appropriate error label.
+        if not self.is_sdam and type(error) == AutoReconnect:
+            self._backoff += 1
+            error._add_error_label("SystemOverloaded")
+            error._add_error_label("Retryable")
+            print(f"Setting backoff in {phase}:", self._backoff)  # noqa: T201
+
     def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connection:
         """Connect to Mongo and return a new Connection.
 
@@ -1062,7 +1072,6 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
             networking_interface = _configured_socket_interface(self.address, self.opts)
         # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
-            print("Got the TLS handshake error")  # noqa: T201
             with self.lock:
                 self.active_contexts.discard(tmp_context)
             if self.enabled_for_cmap:
@@ -1081,10 +1090,10 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
                     reason=_verbose_connection_error_reason(ConnectionClosedReason.ERROR),
                     error=ConnectionClosedReason.ERROR,
                 )
+            self._handle_connection_error(error, "handshake")
             if isinstance(error, (IOError, OSError, *SSLErrors)):
                 details = _get_timeout_details(self.opts)
                 _raise_connection_failure(self.address, error, timeout_details=details)
-
             raise
 
         conn = Connection(networking_interface, self, self.address, conn_id, self.is_sdam)  # type: ignore[arg-type]
@@ -1105,14 +1114,7 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
         except BaseException as e:
             with self.lock:
                 self.active_contexts.discard(conn.cancel_context)
-            # Handle system overload condition.  When the base AutoReconnect is
-            # raised and we are not an sdam pool, add to backoff and add the
-            # appropriate error label.
-            if not self.is_sdam and type(e) == AutoReconnect:
-                self._backoff += 1
-                e._add_error_label("SystemOverloaded")
-                e._add_error_label("Retryable")
-                print("Setting backoff:", self._backoff)  # noqa: T201
+            self._handle_connection_error(e, "hello")
             conn.close_conn(ConnectionClosedReason.ERROR)
             raise