WebSocket Connection Pool (Async)¶

async def mark_completed(self) -> bool:
    """Atomically mark this pending connection as completed.

    Returns:
        True if this call marked it as completed, False if already completed.
    """
    async with self._lock:
        if self._completed:
            return False
        self._completed = True
        return True

async def start(self) -> None:
    """Start the connection pool.

    Initializes the pool, starts the cleanup task, and optionally
    creates warmup connections if configured.
    """
    if self._started or self._closed:
        return

    self._started = True
    if self._cleanup_task is None or self._cleanup_task.done():
        self._cleanup_task = aio.create_task(self._cleanup_loop())

    # Warmup connections if specified
    if self._warmup_connections > 0:
        await self._warmup_pool()

    self.logger.info(f"Started WebSocket connection pool for {self.uri}")

async def acquire(self) -> WebsocketConnection:
    """Acquire a connection from the pool.

    Attempts to reuse an existing available connection, or creates a new
    one if none are available and the pool limit hasn't been reached.
    Includes automatic retry logic for transient failures.

    Returns:
        An acquired WebsocketConnection ready for use.

    Raises:
        ConnectionPoolUnavailableError: If the pool is closed.
        ConnectionPoolExhaustedError: If maximum connections limit is
            reached.
        ConnectionUnavailableError: If connection acquisition fails
            after retries.
    """
    if not self._started:
        await self.start()

    retry = tenacity.retry(
        stop=tenacity.stop_after_attempt(self._max_retries),
        wait=tenacity.wait_exponential(multiplier=1, max=10),
        retry=tenacity.retry_if_exception_type((
            ConnectionUnavailableError,
            ConnectionClosedError,
            OSError,
            aio.TimeoutError,
        )),
    )
    return await retry(self._acquire)()

async def acquire_new(self) -> WebsocketConnection:
    """Force acquire a new connection, avoiding reuse of existing
    ones.

    This method always creates a fresh connection and returns it in
    acquired state, similar to acquire() but without attempting to
    reuse existing connections.

    Returns:
        A newly created and acquired WebsocketConnection.

    Raises:
        ConnectionPoolUnavailableError: If the pool is closed.
        ConnectionPoolExhaustedError: If maximum connections limit is
            reached.
        ConnectionUnavailableError: If connection creation fails after
            retries.
    """
    if not self._started:
        await self.start()

    retry = tenacity.retry(
        stop=tenacity.stop_after_attempt(self._max_retries),
        wait=tenacity.wait_exponential(multiplier=1, max=10),
        retry=tenacity.retry_if_exception_type((
            ConnectionUnavailableError,
            ConnectionClosedError,
            OSError,
            aio.TimeoutError,
        )),
    )
    return await retry(self._acquire_new)()

async def release(self, connection: WebsocketConnection, *, force_remove: bool = False) -> None:
    """Release a connection back to the pool.

    The connection will be either:
    1. Returned to the available pool immediately if data draining is disabled
    2. Put into pending state for background draining if data draining is enabled
    3. Removed from the pool if force_remove is True or the connection is unhealthy

    Args:
        connection: The WebsocketConnection to release.
        force_remove: If True, the connection will be removed from the pool
            regardless of its state. Defaults to False.
    """
    async with self._lock:
        if connection not in self._busy:
            self.logger.warning(
                f"Attempting to release connection not in busy set: {connection}"
            )
            return

        self._busy.remove(connection)

    # Release connection outside lock
    try:
        await connection.release()
    except Exception as e:
        self.logger.warning(f"Error releasing connection: {e}")
        await self._remove_connection(connection)
        return

    async with self._lock:
        if force_remove or not connection.is_connected or self._closed:
            # Connection should be removed
            if connection.is_connected:
                self._total -= 1
            await connection.close()
            self.logger.debug(f"Removed connection from pool: {connection}")
            return

        # If we don't need to check for server data, directly return to available
        if not self._check_server_data_on_release:
            self._available.append(connection)
            self.logger.debug(f"Released connection directly to available: {connection}")
            return

        # Put connection into pending state for background draining
        pending_conn = PendingConnection(
            connection=connection,
            drain_timeout=self._drain_timeout,
            drain_condition=self._drain_condition,
        )

        self._pending[pending_conn.pending_id] = pending_conn

    # Start background drain task outside lock
    pending_conn.drain_task = aio.create_task(self._drain_connection(pending_conn))

    self.logger.debug(f"Put connection into pending state for draining: {connection}")

@contextlib.asynccontextmanager
async def get_connection(self) -> t.AsyncGenerator[WebsocketConnection, None]:
    """Get a connection from the pool as an async context manager.

    The connection will be automatically released back to the pool when
    the context manager exits.

    Yields:
        A WebsocketConnection from the pool.

    Raises:
        ConnectionPoolUnavailableError: If the pool is closed.
        ConnectionPoolExhaustedError: If maximum connections limit is reached.
        ConnectionUnavailableError: If connection acquisition fails.

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

@contextlib.asynccontextmanager
async def get_new_connection(self) -> t.AsyncGenerator[WebsocketConnection, None]:
    """Get a new connection from the pool as an async context
    manager.

    Forces creation of a new connection without reusing existing ones.
    The connection will be automatically released back to the pool when
    the context manager exits.

    Yields:
        A newly created WebsocketConnection.

    Raises:
        ConnectionPoolUnavailableError: If the pool is closed.
        ConnectionPoolExhaustedError: If maximum connections limit is
            reached.
        ConnectionUnavailableError: If connection creation fails.

    Example:
        async with pool.get_new_connection() as conn:
            await conn.send("Hello")
            response = await conn.recv()
    """
    connection = await self.acquire_new()
    try:
        yield connection
    finally:
        await self.release(connection)

async def close_all(self) -> None:
    """Close all connections and shut down the pool.

    Stops the cleanup task and closes all connections in the pool.
    After calling this method, the pool cannot be used again.
    """
    async with self._lock:
        if self._closed:
            return

        self._closed = True

        # Cancel cleanup task
        if self._cleanup_task and not self._cleanup_task.done():
            self._cleanup_task.cancel()

        # Collect all connections to close
        connections_to_close = []
        connections_to_close.extend(list(self._available))
        connections_to_close.extend(list(self._busy))

        drain_tasks_to_cancel = []
        for pending_conn in self._pending.values():
            connections_to_close.append(pending_conn.connection)
            if pending_conn.drain_task and not pending_conn.drain_task.done():
                drain_tasks_to_cancel.append(pending_conn.drain_task)

        # Clear all tracking structures
        self._available.clear()
        self._busy.clear()
        self._pending.clear()
        self._total = 0

    # Wait for cleanup task to finish
    if self._cleanup_task:
        with contextlib.suppress(aio.CancelledError):
            await self._cleanup_task
        self._cleanup_task = None

    # Cancel all drain tasks
    for task in drain_tasks_to_cancel:
        task.cancel()

    # Wait for all drain tasks to finish
    if drain_tasks_to_cancel:
        await aio.gather(*drain_tasks_to_cancel, return_exceptions=True)

    # Close all connections
    for conn in connections_to_close:
        try:
            await conn.close()
        except Exception as e:
            self.logger.warning(f"Error closing connection during pool shutdown: {e}")

    self.logger.info(f"Closed WebSocket connection pool for {self.uri}")