"""Потолок пула: одна и та же работа при разных размерах пула, три модели.

ЧТО ПРОВЕРЯЕТСЯ. Утверждение «перевод на async ускорит работу с базой».
Нагрузка одна и та же — N запросов, каждый ждёт на сервере фиксированное время
(`pg_sleep`), процессор при этом почти не занят. Меняется только модель
исполнения и размер пула соединений:

  * потоки + синхронный psycopg (пул соединений, поток на соединение);
  * корутины + асинхронный psycopg;
  * корутины + asyncpg.

ПОЧЕМУ `pg_sleep`, А НЕ НАСТОЯЩИЙ ЗАПРОС. Настоящий запрос грузит процессор
сервера, а сервер здесь живёт на тех же двух ядрах, что и клиент: измерялась бы
не модель исполнения, а конкуренция за процессор. `pg_sleep` — это чистое
ожидание, то есть ровно тот случай, ради которого асинхронность и заводят.

ЧТО СЧИТАЕТСЯ. Время всей партии и «эффективная одновременность» — во сколько
запросов сервер был занят в среднем: N * sleep / время. Если модель исполнения
что-то решает, эти числа разойдутся.
"""

import asyncio
import sys
import time
from concurrent.futures import ThreadPoolExecutor

import asyncpg
import psycopg
from psycopg_pool import AsyncConnectionPool, ConnectionPool

sys.path.insert(0, __file__.rsplit("/", 1)[0])
import env  # noqa: E402

QUERIES = 600
SLEEP = 0.005  # секунд на сервере, на запрос
POOL_SIZES = [1, 2, 4, 8, 16, 32]
REPEATS = 3  # берётся лучшее: шум машины только добавляет время, не убавляет


def run_threads(pool_size: int) -> float:
    with ConnectionPool(env.DSN, min_size=pool_size, max_size=pool_size) as pool:
        pool.wait()

        def one(_: int) -> None:
            with pool.connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("select pg_sleep(%s)", (SLEEP,))
                    cur.fetchone()

        # Потоков ровно столько же, сколько соединений: лишний поток всё равно
        # встанет в очередь за соединением, и его единственный вклад — ещё один
        # стек и ещё одно переключение.
        with ThreadPoolExecutor(max_workers=pool_size) as ex:
            started = time.perf_counter()
            list(ex.map(one, range(QUERIES)))
            return time.perf_counter() - started


async def _psycopg_async(pool_size: int) -> float:
    pool = AsyncConnectionPool(env.DSN, min_size=pool_size, max_size=pool_size, open=False)
    await pool.open(wait=True)
    try:

        async def one() -> None:
            async with pool.connection() as conn:
                async with conn.cursor() as cur:
                    await cur.execute("select pg_sleep(%s)", (SLEEP,))
                    await cur.fetchone()

        # ВСЕ задачи запускаются разом — именно так выглядит «переписали на
        # async»: одновременных операций столько, сколько их пришло, а не
        # столько, сколько система может выполнить.
        started = time.perf_counter()
        await asyncio.gather(*(one() for _ in range(QUERIES)))
        return time.perf_counter() - started
    finally:
        await pool.close()


async def _asyncpg(pool_size: int) -> float:
    pool = await asyncpg.create_pool(env.DSN, min_size=pool_size, max_size=pool_size)
    try:

        async def one() -> None:
            async with pool.acquire() as conn:
                await conn.fetchval("select pg_sleep($1)", SLEEP)

        started = time.perf_counter()
        await asyncio.gather(*(one() for _ in range(QUERIES)))
        return time.perf_counter() - started
    finally:
        await pool.close()


def best(fn, size: int) -> float:
    return min(fn(size) for _ in range(REPEATS))


def main() -> None:
    env.describe()
    print(f"{QUERIES} запросов, каждый ждёт на сервере {SLEEP * 1000:.0f} мс, лучшее из {REPEATS}")
    print("«идеал» = QUERIES * SLEEP / размер пула — время, если бы ожидания сложились полностью")
    print("«одновр.» = QUERIES * SLEEP / время — во скольких запросов сервер занят в среднем\n")
    print(
        f"{'пул':>4} | {'идеал':>7} | {'потоки':>8} {'одновр.':>8} | "
        f"{'async pg':>9} {'одновр.':>8} | {'asyncpg':>8} {'одновр.':>8}"
    )
    print("-" * 76)

    work = QUERIES * SLEEP
    for size in POOL_SIZES:
        t_threads = best(run_threads, size)
        t_async = best(lambda s: asyncio.run(_psycopg_async(s)), size)
        t_apg = best(lambda s: asyncio.run(_asyncpg(s)), size)
        print(
            f"{size:>4} | {work / size:>6.2f}с | "
            f"{t_threads:>7.2f}с {work / t_threads:>8.2f} | "
            f"{t_async:>8.2f}с {work / t_async:>8.2f} | "
            f"{t_apg:>7.2f}с {work / t_apg:>8.2f}"
        )


if __name__ == "__main__":
    main()