Function async_receive_message

pymongo/network_layer.py:695–743

Receive a raw BSON message or raise socket.error.

(
    conn: AsyncConnection,
    request_id: Optional[int],
    max_message_size: int = MAX_MESSAGE_SIZE,
)

async def async_receive_message(
    conn: AsyncConnection,
    request_id: Optional[int],
    max_message_size: int = MAX_MESSAGE_SIZE,
) -> Union[_OpReply, _OpMsg]:
    """Receive a raw BSON message or raise socket.error."""
    timeout: Optional[Union[float, int]]
    timeout = conn.conn.gettimeout
    if _csot.get_timeout():
        deadline = _csot.get_deadline()
    else:
        if timeout:
            deadline = time.monotonic() + timeout
        else:
            deadline = None
    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]
    try:
        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")
    except asyncio.CancelledError:
        for task in tasks:
            task.cancel()
        await asyncio.wait(tasks)
        raise


def receive_message(