connection_pool_stats.py

import asyncio
from typing import TYPE_CHECKING

from aiohttp.client import ClientSession
from elasticsearch import AsyncElasticsearch
from elasticsearch.connection_pool import DummyConnectionPool, EmptyConnectionPool

if TYPE_CHECKING:
    from aiohttp.client_reqrep import ConnectionKey

def elasticsearch_available_connections(client: AsyncElasticsearch):
    """Get a report of available connections.

    Works with elasticsearch-py 7.x.

    For elasticsearch-py 8.x this would instead need to construct a ConnectionKey for each elasticsearch host, and get
    stats for each.
    """
    return get_aiohttp_pool_stats(get_aiohttp_session(client))


def get_aiohttp_session(client: AsyncElasticsearch) -> ClientSession:
    if isinstance(client.transport.connection_pool, EmptyConnectionPool):
        raise RuntimeError("not initialised")
    if not isinstance(client.transport.connection_pool, DummyConnectionPool):
        raise RuntimeError(f"wrong connection pool type: {client.transport.connection_pool}")
    return client.transport.connection_pool.connection.session


def get_aiohttp_pool_stats(session: ClientSession, key: "ConnectionKey" = None):
    """Print the current available connections.

    Optionally accepts a ConnectionKey, to check available connections for a given target. This
    is relevant in `elasticsearch-py` latest, but 7.16.2.
    """
    connector = session._connector
    assert connector
    total_limit = connector._limit
    total_acquired = len(connector._acquired)
    total_remaining = total_limit - total_acquired

    result = {
        "total_limit": total_limit,
        "total_acquired": total_acquired,
        "total_remaining": total_remaining,
    }

    if key:
        key_limit = connector._limit_per_host
        acquired_for_key = connector._acquired_per_host.get(key)
        assert acquired_for_key is not None
        key_acquired = len(acquired_for_key)
        key_available = key_limit = key_acquired
        result.update(
            {
                "key_limit": key_limit,
                "key_acquired": key_acquired,
                "key_available": key_available,
            }
        )
    return result