PYTHON-5021 - Fix usages of getaddrinfo to be non-blocking

pymongo/_asyncio_executor.py

@@ -0,0 +1,22 @@
+# Copyright 2024-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A separate ThreadPoolExecutor instance used internally to avoid competing for resources with the default asyncio ThreadPoolExecutor
+that user code will use."""
+
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+
+_PYMONGO_EXECUTOR = ThreadPoolExecutor(thread_name_prefix="PYMONGO_EXECUTOR-")

pymongo/asynchronous/auth.py

@@ -38,6 +38,7 @@
     _authenticate_oidc,
     _get_authenticator,
 )
+from pymongo.asynchronous.helpers import getaddrinfo
 from pymongo.auth_shared import (
     MongoCredential,
     _authenticate_scram_start,
@@ -177,15 +178,22 @@ def _auth_key(nonce: str, username: str, password: str) -> str:
     return md5hash.hexdigest()
 
 
-def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
+async def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
     """Canonicalize hostname following MIT-krb5 behavior."""
     # https://github.com/krb5/krb5/blob/d406afa363554097ac48646a29249c04f498c88e/src/util/k5test.py#L505-L520
     if option in [False, "none"]:
         return hostname
 
-    af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
-        hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
-    )[0]
+    af, socktype, proto, canonname, sockaddr = (
+        await getaddrinfo(
+            hostname,
+            None,
+            family=0,
+            type=0,
+            proto=socket.IPPROTO_TCP,
+            flags=socket.AI_CANONNAME,
+        )
+    )[0]  # type: ignore[index]
 
     # For forward just to resolve the cname as dns.lookup() will not return it.
     if option == "forward":
@@ -213,7 +221,7 @@ async def _authenticate_gssapi(credentials: MongoCredential, conn: AsyncConnecti
         # Starting here and continuing through the while loop below - establish
         # the security context. See RFC 4752, Section 3.1, first paragraph.
         host = props.service_host or conn.address[0]
-        host = _canonicalize_hostname(host, props.canonicalize_host_name)
+        host = await _canonicalize_hostname(host, props.canonicalize_host_name)
         service = props.service_name + "@" + host
         if props.service_realm is not None:
             service = service + "@" + props.service_realm

pymongo/asynchronous/helpers.py

@@ -15,7 +15,10 @@
 """Miscellaneous pieces that need to be synchronized."""
 from __future__ import annotations
 
+import asyncio
 import builtins
+import functools
+import socket
 import sys
 from typing import (
     Any,
@@ -24,6 +27,7 @@
     cast,
 )
 
+from pymongo._asyncio_executor import _PYMONGO_EXECUTOR
 from pymongo.errors import (
     OperationFailure,
 )
@@ -68,6 +72,26 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
+async def getaddrinfo(
+    host: Any, port: Any, **kwargs: Any
+) -> list[
+    tuple[
+        socket.AddressFamily,
+        socket.SocketKind,
+        int,
+        str,
+        tuple[str, int] | tuple[str, int, int, int],
+    ]
+]:
+    if not _IS_SYNC:
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(  # type: ignore[return-value]
+            _PYMONGO_EXECUTOR, functools.partial(socket.getaddrinfo, host, port, **kwargs)
+        )
+    else:
+        return socket.getaddrinfo(host, port, **kwargs)
+
+
 if sys.version_info >= (3, 10):
     anext = builtins.anext
     aiter = builtins.aiter

pymongo/asynchronous/pool.py

@@ -39,8 +39,9 @@
 
 from bson import DEFAULT_CODEC_OPTIONS
 from pymongo import _csot, helpers_shared
+from pymongo._asyncio_executor import _PYMONGO_EXECUTOR
 from pymongo.asynchronous.client_session import _validate_session_write_concern
-from pymongo.asynchronous.helpers import _handle_reauth
+from pymongo.asynchronous.helpers import _handle_reauth, getaddrinfo
 from pymongo.asynchronous.network import command, receive_message
 from pymongo.common import (
     MAX_BSON_SIZE,
@@ -783,7 +784,7 @@ def __repr__(self) -> str:
         )
 
 
-def _create_connection(address: _Address, options: PoolOptions) -> socket.socket:
+async def _create_connection(address: _Address, options: PoolOptions) -> socket.socket:
     """Given (host, port) and PoolOptions, connect and return a socket object.
 
     Can raise socket.error.
@@ -814,7 +815,7 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
         family = socket.AF_UNSPEC
 
     err = None
-    for res in socket.getaddrinfo(host, port, family, socket.SOCK_STREAM):
+    for res in await getaddrinfo(host, port, family=family, type=socket.SOCK_STREAM):  # type: ignore[attr-defined]
         af, socktype, proto, dummy, sa = res
         # SOCK_CLOEXEC was new in CPython 3.2, and only available on a limited
         # number of platforms (newer Linux and *BSD). Starting with CPython 3.4
@@ -863,7 +864,7 @@ async def _configured_socket(
 
     Sets socket's SSL and timeout options.
     """
-    sock = _create_connection(address, options)
+    sock = await _create_connection(address, options)
     ssl_context = options._ssl_context
 
     if ssl_context is None:
@@ -883,7 +884,7 @@ async def _configured_socket(
                 else:
                     loop = asyncio.get_running_loop()
                     ssl_sock = await loop.run_in_executor(
-                        None,
+                        _PYMONGO_EXECUTOR,
                         functools.partial(ssl_context.wrap_socket, sock, server_hostname=host),  # type: ignore[assignment, misc]
                     )
         else:
@@ -894,7 +895,9 @@ async def _configured_socket(
                     ssl_sock = await ssl_context.a_wrap_socket(sock)  # type: ignore[assignment, misc]
                 else:
                     loop = asyncio.get_running_loop()
-                    ssl_sock = await loop.run_in_executor(None, ssl_context.wrap_socket, sock)  # type: ignore[assignment, misc]
+                    ssl_sock = await loop.run_in_executor(
+                        _PYMONGO_EXECUTOR, ssl_context.wrap_socket, sock
+                    )  # type: ignore[assignment, misc]
     except _CertificateError:
         sock.close()
         # Raise _CertificateError directly like we do after match_hostname

pymongo/pyopenssl_context.py

@@ -31,6 +31,7 @@
 from OpenSSL import SSL as _SSL
 from OpenSSL import crypto as _crypto
 
+from pymongo._asyncio_executor import _PYMONGO_EXECUTOR
 from pymongo.errors import ConfigurationError as _ConfigurationError
 from pymongo.errors import _CertificateError  # type:ignore[attr-defined]
 from pymongo.ocsp_cache import _OCSPCache
@@ -405,15 +406,15 @@ async def a_wrap_socket(
                 ssl_conn.set_tlsext_host_name(server_hostname.encode("idna"))
             if self.verify_mode != _stdlibssl.CERT_NONE:
                 # Request a stapled OCSP response.
-                await loop.run_in_executor(None, ssl_conn.request_ocsp)
+                await loop.run_in_executor(_PYMONGO_EXECUTOR, ssl_conn.request_ocsp)
             ssl_conn.set_connect_state()
         # If this wasn't true the caller of wrap_socket would call
         # do_handshake()
         if do_handshake_on_connect:
             # XXX: If we do hostname checking in a callback we can get rid
             # of this call to do_handshake() since the handshake
             # will happen automatically later.
-            await loop.run_in_executor(None, ssl_conn.do_handshake)
+            await loop.run_in_executor(_PYMONGO_EXECUTOR, ssl_conn.do_handshake)
             # XXX: Do this in a callback registered with
             # SSLContext.set_info_callback? See Twisted for an example.
             if self.check_hostname and server_hostname is not None:

pymongo/synchronous/auth.py

@@ -45,6 +45,7 @@
     _authenticate_oidc,
     _get_authenticator,
 )
+from pymongo.synchronous.helpers import getaddrinfo
 
 if TYPE_CHECKING:
     from pymongo.hello import Hello
@@ -180,9 +181,16 @@ def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
     if option in [False, "none"]:
         return hostname
 
-    af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
-        hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
-    )[0]
+    af, socktype, proto, canonname, sockaddr = (
+        getaddrinfo(
+            hostname,
+            None,
+            family=0,
+            type=0,
+            proto=socket.IPPROTO_TCP,
+            flags=socket.AI_CANONNAME,
+        )
+    )[0]  # type: ignore[index]
 
     # For forward just to resolve the cname as dns.lookup() will not return it.
     if option == "forward":

pymongo/synchronous/helpers.py

@@ -15,7 +15,10 @@
 """Miscellaneous pieces that need to be synchronized."""
 from __future__ import annotations
 
+import asyncio
 import builtins
+import functools
+import socket
 import sys
 from typing import (
     Any,
@@ -24,6 +27,7 @@
     cast,
 )
 
+from pymongo._asyncio_executor import _PYMONGO_EXECUTOR
 from pymongo.errors import (
     OperationFailure,
 )
@@ -68,6 +72,26 @@ def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
+def getaddrinfo(
+    host: Any, port: Any, **kwargs: Any
+) -> list[
+    tuple[
+        socket.AddressFamily,
+        socket.SocketKind,
+        int,
+        str,
+        tuple[str, int] | tuple[str, int, int, int],
+    ]
+]:
+    if not _IS_SYNC:
+        loop = asyncio.get_running_loop()
+        return loop.run_in_executor(  # type: ignore[return-value]
+            _PYMONGO_EXECUTOR, functools.partial(socket.getaddrinfo, host, port, **kwargs)
+        )
+    else:
+        return socket.getaddrinfo(host, port, **kwargs)
+
+
 if sys.version_info >= (3, 10):
     next = builtins.next
     iter = builtins.iter

pymongo/synchronous/pool.py

@@ -39,6 +39,7 @@
 
 from bson import DEFAULT_CODEC_OPTIONS
 from pymongo import _csot, helpers_shared
+from pymongo._asyncio_executor import _PYMONGO_EXECUTOR
 from pymongo.common import (
     MAX_BSON_SIZE,
     MAX_MESSAGE_SIZE,
@@ -84,7 +85,7 @@
 from pymongo.socket_checker import SocketChecker
 from pymongo.ssl_support import HAS_SNI, SSLError
 from pymongo.synchronous.client_session import _validate_session_write_concern
-from pymongo.synchronous.helpers import _handle_reauth
+from pymongo.synchronous.helpers import _handle_reauth, getaddrinfo
 from pymongo.synchronous.network import command, receive_message
 
 if TYPE_CHECKING:
@@ -812,7 +813,7 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
         family = socket.AF_UNSPEC
 
     err = None
-    for res in socket.getaddrinfo(host, port, family, socket.SOCK_STREAM):
+    for res in getaddrinfo(host, port, family=family, type=socket.SOCK_STREAM):  # type: ignore[attr-defined]
         af, socktype, proto, dummy, sa = res
         # SOCK_CLOEXEC was new in CPython 3.2, and only available on a limited
         # number of platforms (newer Linux and *BSD). Starting with CPython 3.4
@@ -879,7 +880,7 @@ def _configured_socket(address: _Address, options: PoolOptions) -> Union[socket.
                 else:
                     loop = asyncio.get_running_loop()
                     ssl_sock = loop.run_in_executor(
-                        None,
+                        _PYMONGO_EXECUTOR,
                         functools.partial(ssl_context.wrap_socket, sock, server_hostname=host),  # type: ignore[assignment, misc]
                     )
         else:
@@ -890,7 +891,9 @@ def _configured_socket(address: _Address, options: PoolOptions) -> Union[socket.
                     ssl_sock = ssl_context.a_wrap_socket(sock)  # type: ignore[assignment, misc]
                 else:
                     loop = asyncio.get_running_loop()
-                    ssl_sock = loop.run_in_executor(None, ssl_context.wrap_socket, sock)  # type: ignore[assignment, misc]
+                    ssl_sock = loop.run_in_executor(
+                        _PYMONGO_EXECUTOR, ssl_context.wrap_socket, sock
+                    )  # type: ignore[assignment, misc]
     except _CertificateError:
         sock.close()
         # Raise _CertificateError directly like we do after match_hostname

test/asynchronous/test_auth.py

@@ -275,10 +275,10 @@ async def test_gssapi_threaded(self):
     async def test_gssapi_canonicalize_host_name(self):
         # Test the low level method.
         assert GSSAPI_HOST is not None
-        result = _canonicalize_hostname(GSSAPI_HOST, "forward")
+        result = await _canonicalize_hostname(GSSAPI_HOST, "forward")
         if "compute-1.amazonaws.com" not in result:
             self.assertEqual(result, GSSAPI_HOST)
-        result = _canonicalize_hostname(GSSAPI_HOST, "forwardAndReverse")
+        result = await _canonicalize_hostname(GSSAPI_HOST, "forwardAndReverse")
         self.assertEqual(result, GSSAPI_HOST)
 
         # Use the equivalent named CANONICALIZE_HOST_NAME.