WebSocket Connection¶

async def connect(self) -> None:
    """Establish the WebSocket connection.

    If already connected, this method does nothing. Otherwise, it attempts
    to establish a new connection and starts the idle monitor task.

    Raises:
        ConnectionUnavailableError: If the connection has been closed or
            if connection establishment fails.
    """
    async with self._lock:
        if self.is_connected:
            return

        if self._closed:
            raise ConnectionUnavailableError("Connection has been closed")

        # Clean up old monitor task if done
        if self._monitor_task is not None and self._monitor_task.done():
            try:
                await self._monitor_task
            finally:
                self._monitor_task = None

    # Establish connection outside the lock to avoid blocking
    try:
        websocket = await connect(self.uri, additional_headers=self.headers, **self._params)

        async with self._lock:
            self.websocket = websocket
            self._update_activity()

            # Start monitor task if not already running
            if self._monitor_task is None or self._monitor_task.done():
                self._monitor_task = aio.create_task(self._monitor_idle())

        self.logger.debug(f"Connected to {self.uri}")
    except Exception as e:
        self.logger.error(f"Failed to connect to {self.uri}: {e}")
        raise ConnectionUnavailableError(f"Failed to connect: {e}") from e

async def close(self) -> None:
    """Close the WebSocket connection and clean up resources.

    This method cancels the idle monitor task, closes the WebSocket
    connection, and resets all internal state flags.
    """
    async with self._lock:
        if self._closed:
            return

        self._closed = True

        # Cancel monitor task
        if self._monitor_task is not None:
            if not self._monitor_task.done():
                self._monitor_task.cancel()

            try:
                await self._monitor_task
            except aio.CancelledError:
                self.logger.debug(f"Monitor task for {self.uri} cancelled")
            finally:
                self._monitor_task = None

        # Close websocket
        if self.websocket is not None:
            try:
                await self.websocket.close()
            finally:
                self.websocket = None

        self._is_busy = False
        self._is_draining = False
        self.logger.debug(f"Closed connection to {self.uri}")

async def acquire(self) -> None:
    """Acquire the connection for exclusive use.

    This method marks the connection as busy and ensures it is connected.

    Raises:
        ConnectionUnavailableError: If the connection has been closed or
            if connection establishment fails.
        ConnectionBusyError: If the connection is already busy or draining.
    """
    async with self._lock:
        if self._closed:
            raise ConnectionUnavailableError("Connection has been closed")
        if self._is_busy:
            raise ConnectionBusyError("Connection is already busy")
        if self._is_draining:
            raise ConnectionBusyError("Connection is currently draining")

    try:
        await self.connect()
        async with self._lock:
            self._is_busy = True
            self._update_activity()
    except Exception as e:
        self.logger.error(f"Failed to acquire connection to {self.uri}: {e}")
        raise ConnectionUnavailableError(f"Failed to acquire connection: {e}") from e

async def release(self) -> None:
    """Release the connection back to the pool.

    This method marks the connection as no longer busy and updates
    the activity timestamp.
    """
    async with self._lock:
        if not self._is_busy:
            return

        self._is_busy = False
        self._update_activity()

async def ping(self) -> None:
    """Send a ping to the WebSocket server to check connection
    health.

    This method performs a health check by sending a ping frame to the
    server. The connection must be open for this to succeed.

    Raises:
        ConnectionUnavailableError: If the websocket is not connected.
        ConnectionClosedError: If the connection closes during the ping.
    """
    # Check connection state outside lock (quick check)
    if not (
        self.websocket is not None
        and not self._closed
        and self.websocket.state == protocol.OPEN
    ):
        raise ConnectionUnavailableError("Websocket is not connected")

    # Get websocket reference under lock
    async with self._lock:
        if not self.is_connected:
            raise ConnectionUnavailableError("Websocket is not connected")
        websocket = self.websocket

    # Perform ping outside lock to avoid blocking other operations
    connection_closed = False
    connection_closed_error = None
    try:
        await websocket.ping()
        async with self._lock:
            self._update_activity()
    except ConnectionClosed as e:
        self.logger.error(f"Connection closed during ping: {e}")
        connection_closed = True
        connection_closed_error = e

    if connection_closed:
        await self.close()
        raise ConnectionClosedError(
            f"Connection closed during ping: {connection_closed_error}"
        ) from connection_closed_error

@contextlib.asynccontextmanager
async def session(self) -> t.AsyncGenerator[WebsocketConnection, None]:
    """Create a context manager session for the connection.

    The connection is automatically acquired on entry and released on exit.

    Yields:
        The WebsocketConnection instance.

    Example:
        async with connection.session():
            await connection.send("Hello")
            response = await connection.recv()
    """
    await self.acquire()
    try:
        yield self
    finally:
        await self.release()

async def send(self, message: str | bytes) -> None:
    """Send a message through the WebSocket connection.

    Args:
        message: The message to send (string or bytes).

    Raises:
        ConnectionBusyError: If the connection has not been acquired.
        ConnectionUnavailableError: If the websocket is not connected.
        ConnectionClosedError: If the connection closes during sending.
    """
    if not self.is_busy:
        raise ConnectionBusyError("Connection must be acquired before sending messages")

    # Check connection state outside lock
    if not (
        self.websocket is not None
        and not self._closed
        and self.websocket.state == protocol.OPEN
    ):
        raise ConnectionUnavailableError("Websocket is not connected")

    # Get websocket reference under lock
    async with self._lock:
        websocket = self.websocket
        if websocket is None:
            raise ConnectionUnavailableError("Websocket is not connected")

    # Perform send outside lock
    try:
        await websocket.send(message)
        async with self._lock:
            self._update_activity()
    except ConnectionClosed as e:
        await self.close()
        self.logger.error(f"Connection closed while sending message: {e}")
        raise ConnectionClosedError(f"Connection closed while sending message: {e}") from e

async def recv(self) -> str | bytes:
    """Receive a message from the WebSocket connection.

    Returns:
        The received message (string or bytes).

    Raises:
        ConnectionBusyError: If the connection has not been acquired.
        ConnectionUnavailableError: If the websocket is not connected.
        ConnectionClosedError: If the connection closes during receiving.
    """
    if not self._is_busy:
        raise ConnectionBusyError("Connection must be acquired before receiving messages")

    # Check connection state outside lock
    if not (
        self.websocket is not None
        and not self._closed
        and self.websocket.state == protocol.OPEN
    ):
        raise ConnectionUnavailableError("Websocket is not connected")

    # Get websocket reference under lock
    async with self._lock:
        websocket = self.websocket
        if websocket is None:
            raise ConnectionUnavailableError("Websocket is not connected")

    # Perform recv outside lock
    try:
        message = await websocket.recv()
        async with self._lock:
            self._update_activity()
        return message
    except ConnectionClosed as e:
        await self.close()
        self.logger.error(f"Connection closed while receiving message: {e}")
        raise ConnectionClosedError(f"Connection closed while receiving message: {e}") from e