Sync tests all passing

Author: NoahStapp

pymongo/asynchronous/encryption.py

@@ -63,7 +63,6 @@
 from pymongo.asynchronous.cursor import AsyncCursor
 from pymongo.asynchronous.database import AsyncDatabase
 from pymongo.asynchronous.mongo_client import AsyncMongoClient
-from pymongo.asynchronous.pool import _configured_socket, _raise_connection_failure
 from pymongo.common import CONNECT_TIMEOUT
 from pymongo.daemon import _spawn_daemon
 from pymongo.encryption_options import AutoEncryptionOpts, RangeOpts
@@ -75,12 +74,13 @@
     PyMongoError,
     ServerSelectionTimeoutError,
 )
-from pymongo.network_layer import BLOCKING_IO_ERRORS, async_sendall
+from pymongo.network_layer import async_sendall
 from pymongo.operations import UpdateOne
 from pymongo.pool_options import PoolOptions
+from pymongo.pool_shared import _configured_socket, _raise_connection_failure
 from pymongo.read_concern import ReadConcern
 from pymongo.results import BulkWriteResult, DeleteResult
-from pymongo.ssl_support import get_ssl_context
+from pymongo.ssl_support import BLOCKING_IO_ERRORS, get_ssl_context
 from pymongo.typings import _DocumentType, _DocumentTypeArg
 from pymongo.uri_parser import parse_host
 from pymongo.write_concern import WriteConcern

pymongo/asynchronous/network.py

@@ -189,7 +189,7 @@ async def command(
         )
 
     try:
-        await async_sendall(conn.conn.writer, msg)
+        await async_sendall(conn.conn.get_conn, msg)
         if use_op_msg and unacknowledged:
             # Unacknowledged, fake a successful command response.
             reply = None

pymongo/asynchronous/pool.py

@@ -17,11 +17,8 @@
 import asyncio
 import collections
 import contextlib
-import functools
 import logging
 import os
-import socket
-import ssl
 import sys
 import time
 import weakref
@@ -52,16 +49,13 @@
 from pymongo.errors import (  # type:ignore[attr-defined]
     AutoReconnect,
     ConfigurationError,
-    ConnectionFailure,
     DocumentTooLarge,
     ExecutionTimeout,
     InvalidOperation,
-    NetworkTimeout,
     NotPrimaryError,
     OperationFailure,
     PyMongoError,
     WaitQueueTimeoutError,
-    _CertificateError,
 )
 from pymongo.hello import Hello, HelloCompat
 from pymongo.lock import (
@@ -79,10 +73,15 @@
     ConnectionCheckOutFailedReason,
     ConnectionClosedReason,
 )
-from pymongo.network_layer import async_receive_message, async_sendall, AsyncNetworkingInterface
+from pymongo.network_layer import AsyncNetworkingInterface, async_receive_message, async_sendall
 from pymongo.pool_options import PoolOptions
-from pymongo.pool_shared import _configured_protocol, _CancellationContext, _get_timeout_details, format_timeout_details, \
-    _raise_connection_failure
+from pymongo.pool_shared import (
+    _CancellationContext,
+    _configured_protocol,
+    _get_timeout_details,
+    _raise_connection_failure,
+    format_timeout_details,
+)
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_api import _add_to_command
 from pymongo.server_type import SERVER_TYPE
@@ -101,7 +100,6 @@
         ZstdContext,
     )
     from pymongo.message import _OpMsg, _OpReply
-    from pymongo.pyopenssl_context import _sslConn
     from pymongo.read_concern import ReadConcern
     from pymongo.read_preferences import _ServerMode
     from pymongo.typings import ClusterTime, _Address, _CollationIn
@@ -195,6 +193,7 @@ def set_conn_timeout(self, timeout: Optional[float]) -> None:
         if timeout == self.last_timeout:
             return
         self.last_timeout = timeout
+        self.conn.get_conn.settimeout(timeout)
 
     def apply_timeout(
         self, client: AsyncMongoClient, cmd: Optional[MutableMapping[str, Any]]
@@ -453,7 +452,7 @@ async def send_message(self, message: bytes, max_doc_size: int) -> None:
             )
 
         try:
-            await async_sendall(self.conn.writer, message)
+            await async_sendall(self.conn.get_conn, message)
         except BaseException as error:
             self._raise_connection_failure(error)
 
@@ -589,7 +588,10 @@ def _close_conn(self) -> None:
 
     def conn_closed(self) -> bool:
         """Return True if we know socket has been closed, False otherwise."""
-        return self.conn.is_closing()
+        if _IS_SYNC:
+            return self.socket_checker.socket_closed(self.conn.get_conn)
+        else:
+            return self.conn.is_closing()
 
     def send_cluster_time(
         self,
@@ -977,9 +979,7 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                     self.requests -= 1
                     self.size_cond.notify()
 
-    async def connect(
-        self, handler: Optional[_MongoClientErrorHandler] = None
-    ) -> AsyncConnection:
+    async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> AsyncConnection:
         """Connect to Mongo and return a new AsyncConnection.
 
         Can raise ConnectionFailure.

pymongo/network_layer.py

@@ -26,7 +26,7 @@
     Union,
 )
 
-from pymongo import _csot
+from pymongo import _csot, ssl_support
 from pymongo._asyncio_task import create_task
 from pymongo.common import MAX_MESSAGE_SIZE
 from pymongo.compression_support import decompress
@@ -59,7 +59,9 @@
 
 
 class NetworkingInterfaceBase:
-    def __init__(self, conn: Union[socket.socket, _sslConn] | tuple[asyncio.BaseTransport, PyMongoProtocol]):
+    def __init__(
+        self, conn: Union[socket.socket, _sslConn] | tuple[asyncio.BaseTransport, PyMongoProtocol]
+    ):
         self.conn = conn
 
     def gettimeout(self):
@@ -74,10 +76,7 @@ def close(self):
     def is_closing(self) -> bool:
         raise NotImplementedError
 
-    def writer(self):
-        raise NotImplementedError
-
-    def reader(self):
+    def get_conn(self):
         raise NotImplementedError
 
 
@@ -99,11 +98,7 @@ def is_closing(self):
         self.conn[0].is_closing()
 
     @property
-    def writer(self) -> PyMongoProtocol:
-        return self.conn[1]
-
-    @property
-    def reader(self) -> PyMongoProtocol:
+    def get_conn(self) -> PyMongoProtocol:
         return self.conn[1]
 
 
@@ -124,11 +119,7 @@ def is_closing(self):
         self.conn.is_closing()
 
     @property
-    def writer(self):
-        return self.conn
-
-    @property
-    def reader(self):
+    def get_conn(self):
         return self.conn
 
 
@@ -333,35 +324,8 @@ async def _poll_cancellation(conn: AsyncConnection) -> None:
         await asyncio.sleep(_POLL_TIMEOUT)
 
 
-async def async_receive_data(
-    conn: AsyncConnection,
-    deadline: Optional[float],
-    request_id: Optional[int],
-    max_message_size: int,
-) -> memoryview:
-    conn_timeout = conn.conn.gettimeout
-    timeout: Optional[Union[float, int]]
-    if deadline:
-        # When the timeout has expired perform one final check to
-        # see if the socket is readable. This helps avoid spurious
-        # timeouts on AWS Lambda and other FaaS environments.
-        timeout = max(deadline - time.monotonic(), 0)
-    else:
-        timeout = conn_timeout
-
-    cancellation_task = create_task(_poll_cancellation(conn))
-    read_task = create_task(conn.conn.reader.read(request_id, max_message_size))
-    tasks = [read_task, cancellation_task]
-    done, pending = await asyncio.wait(tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED)
-    for task in pending:
-        task.cancel()
-    if pending:
-        await asyncio.wait(pending)
-    if len(done) == 0:
-        raise socket.timeout("timed out")
-    if read_task in done:
-        return read_task.result()
-    raise _OperationCancelled("operation cancelled")
+# Errors raised by sockets (and TLS sockets) when in non-blocking mode.
+BLOCKING_IO_ERRORS = (BlockingIOError, *ssl_support.BLOCKING_IO_ERRORS)
 
 
 def receive_data(conn: Connection, length: int, deadline: Optional[float]) -> memoryview:
@@ -384,12 +348,12 @@ def receive_data(conn: Connection, length: int, deadline: Optional[float]) -> me
                 short_timeout = _POLL_TIMEOUT
             conn.set_conn_timeout(short_timeout)
             try:
-                chunk_length = conn.conn.recv_into(mv[bytes_read:])
-            # except BLOCKING_IO_ERRORS:
-            #     if conn.cancel_context.cancelled:
-            #         raise _OperationCancelled("operation cancelled") from None
-            #     # We reached the true deadline.
-            #     raise socket.timeout("timed out") from None
+                chunk_length = conn.conn.get_conn.recv_into(mv[bytes_read:])
+            except BLOCKING_IO_ERRORS:
+                if conn.cancel_context.cancelled:
+                    raise _OperationCancelled("operation cancelled") from None
+                # We reached the true deadline.
+                raise socket.timeout("timed out") from None
             except socket.timeout:
                 if conn.cancel_context.cancelled:
                     raise _OperationCancelled("operation cancelled") from None
@@ -416,22 +380,42 @@ async def async_receive_message(
     max_message_size: int = MAX_MESSAGE_SIZE,
 ) -> Union[_OpReply, _OpMsg]:
     """Receive a raw BSON message or raise socket.error."""
+    timeout: Optional[Union[float, int]]
     if _csot.get_timeout():
         deadline = _csot.get_deadline()
     else:
-        timeout = conn.conn.reader.gettimeout
+        timeout = conn.conn.get_conn.gettimeout
         if timeout:
             deadline = time.monotonic() + timeout
         else:
             deadline = None
-    data, op_code = await async_receive_data(conn, deadline, request_id, max_message_size)
-    try:
-        unpack_reply = _UNPACK_REPLY[op_code]
-    except KeyError:
-        raise ProtocolError(
-            f"Got opcode {op_code!r} but expected {_UNPACK_REPLY.keys()!r}"
-        ) from None
-    return unpack_reply(data)
+    if deadline:
+        # When the timeout has expired perform one final check to
+        # see if the socket is readable. This helps avoid spurious
+        # timeouts on AWS Lambda and other FaaS environments.
+        timeout = max(deadline - time.monotonic(), 0)
+
+    cancellation_task = create_task(_poll_cancellation(conn))
+    read_task = create_task(conn.conn.get_conn.read(request_id, max_message_size))
+    tasks = [read_task, cancellation_task]
+    done, pending = await asyncio.wait(tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED)
+    for task in pending:
+        task.cancel()
+    if pending:
+        await asyncio.wait(pending)
+    if len(done) == 0:
+        raise socket.timeout("timed out")
+    if read_task in done:
+        data, op_code = read_task.result()
+
+        try:
+            unpack_reply = _UNPACK_REPLY[op_code]
+        except KeyError:
+            raise ProtocolError(
+                f"Got opcode {op_code!r} but expected {_UNPACK_REPLY.keys()!r}"
+            ) from None
+        return unpack_reply(data)
+    raise _OperationCancelled("operation cancelled")
 
 
 def receive_message(

pymongo/pool_shared.py

@@ -9,30 +9,20 @@
     Any,
     NoReturn,
     Optional,
-    Union,
 )
 
 from pymongo import _csot
 from pymongo.errors import (  # type:ignore[attr-defined]
     AutoReconnect,
-    ConfigurationError,
     ConnectionFailure,
-    DocumentTooLarge,
-    ExecutionTimeout,
-    InvalidOperation,
     NetworkTimeout,
-    NotPrimaryError,
-    OperationFailure,
-    PyMongoError,
-    WaitQueueTimeoutError,
     _CertificateError,
 )
-from pymongo.network_layer import PyMongoProtocol, AsyncNetworkingInterface
+from pymongo.network_layer import AsyncNetworkingInterface, NetworkingInterface, PyMongoProtocol
 from pymongo.pool_options import PoolOptions
 from pymongo.ssl_support import HAS_SNI, SSLError
 
 if TYPE_CHECKING:
-    from pymongo.pyopenssl_context import _sslConn
     from pymongo.typings import _Address
 
 try:
@@ -50,6 +40,7 @@ def _set_non_inheritable_non_atomic(fd: int) -> None:
     def _set_non_inheritable_non_atomic(fd: int) -> None:  # noqa: ARG001
         """Dummy function for platforms that don't provide fcntl."""
 
+
 _MAX_TCP_KEEPIDLE = 120
 _MAX_TCP_KEEPINTVL = 10
 _MAX_TCP_KEEPCNT = 9
@@ -249,9 +240,7 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
         raise OSError("getaddrinfo failed")
 
 
-async def _configured_protocol(
-    address: _Address, options: PoolOptions
-) -> AsyncNetworkingInterface:
+async def _configured_protocol(address: _Address, options: PoolOptions) -> AsyncNetworkingInterface:
     """Given (host, port) and PoolOptions, return a configured transport, protocol pair.
 
     Can raise socket.error, ConnectionFailure, or _CertificateError.
@@ -263,9 +252,11 @@ async def _configured_protocol(
     timeout = sock.gettimeout()
 
     if ssl_context is None:
-        return AsyncNetworkingInterface(await asyncio.get_running_loop().create_connection(
-            lambda: PyMongoProtocol(timeout=timeout, buffer_size=2**16), sock=sock
-        ))
+        return AsyncNetworkingInterface(
+            await asyncio.get_running_loop().create_connection(
+                lambda: PyMongoProtocol(timeout=timeout, buffer_size=2**16), sock=sock
+            )
+        )
 
     host = address[0]
     try:
@@ -303,9 +294,7 @@ async def _configured_protocol(
     return AsyncNetworkingInterface((transport, protocol))
 
 
-def _configured_socket(
-    address: _Address, options: PoolOptions
-) -> Union[socket.socket, _sslConn]:
+def _configured_socket(address: _Address, options: PoolOptions) -> NetworkingInterface:
     """Given (host, port) and PoolOptions, return a configured socket.
 
     Can raise socket.error, ConnectionFailure, or _CertificateError.
@@ -317,7 +306,7 @@ def _configured_socket(
 
     if ssl_context is None:
         sock.settimeout(options.socket_timeout)
-        return sock
+        return NetworkingInterface(sock)
 
     host = address[0]
     try:
@@ -351,4 +340,4 @@ def _configured_socket(
             raise
 
     ssl_sock.settimeout(options.socket_timeout)
-    return ssl_sock
+    return NetworkingInterface(ssl_sock)