"""Куда уходит запрос к шардированной таблице и сколько он стоит.

ЗАЧЕМ ЭТОТ СКРИПТ. Про запрос без ключа шардирования говорят «он идёт во все
шарды». Это верно и недоговорено: во сколько раз это дороже, зависит не от
числа шардов, а от того, ходит ли координатор к ним по очереди или
одновременно, — и в PostgreSQL это одна настройка, выключенная по умолчанию.

ЧТО ИМЕННО МЕРЯЕТСЯ.

  1. План: запрос с ключом шардирования и без него. Это не замер, а то, что
     печатает сам PostgreSQL, — отсечение секций видно в плане буквально.
  2. Сколько порций уходит к шардам за один запрос. Порция — единица, к
     которой ретранслятор применяет задержку, то есть это число обменов.
  2. Что координатор отправляет шарду: ретранслятор разбирает протокол и
     печатает сообщения. Отсюда число обменов на шард.
  3. Время запроса на канале с задержкой: с ключом, без ключа по очереди,
     без ключа одновременно (async_capable), и то же с параллельной
     фиксацией (parallel_commit). Для 4 и 8 шардов.
  4. Сколько кругов оплачивается строго друг за другом. Каждая конфигурация
     меряется при двух задержках; прирост времени, делённый на прирост
     круга, — это число последовательных кругов. Так видно, какие из пяти
     обменов async_capable и parallel_commit действительно совмещают.

ПРОТОКОЛ ЗАМЕРА. Один запрос — отдельная транзакция, как в приложении без
явных BEGIN. Соединения к шардам прогреты: первый запрос в каждой конфигурации
выбрасывается, иначе в замер попала бы установка соединения. Девять кругов по
30 запросов, ключи из фиксированного зерна; печатается медиана круга, лучший
и худший круг и разброс между ними. Порядок конфигураций внутри круга
чередуется, чтобы дрейф машины не ложился на одну из них.

ЗАДЕРЖКА — ПАРАМЕТР, А НЕ СВОЙСТВО МИРА. Основная — 1 мс в каждую сторону,
круг 2 мс: порядок задержки между стойками или зонами доступности одного
региона. Вторая, 3 мс, нужна только для того, чтобы посчитать число
последовательных кругов. Выводы сформулированы в кругах, а не в
миллисекундах, и переносятся на любую задержку умножением.

ОШИБКА ЭТОГО СКРИПТА, ОСТАВЛЕННАЯ В ИСТОРИИ. Первая редакция заканчивалась
фразой, написанной до прогона: при одновременном походе в шарды цена
запроса без ключа «растёт с числом шардов гораздо медленнее». Прогон её
опроверг: при переходе с 4 на 8 шардов она росла почти так же, как у похода
по очереди. Отсюда трассировка протокола (блок 2) и счёт
последовательных кругов (блок 6): они показывают, что async_capable
совмещает только выборку, а parallel_commit — только фиксацию.

ЗАПУСК:
    DE_BENCH_DSN="postgres://postgres@localhost:5433/postgres" \\
      python3 bench/sharding/routing.py
Вывод: runs/routing.txt
"""

from __future__ import annotations

import random
import statistics
import sys
import time

from stand import Stand, version

ROUNDS = 9
PER_ROUND = 30
SEED = 20260923


def show(title: str) -> None:
    print()
    print(title)
    print("-" * len(title))


KEY_Q = "select order_id, amount from orders where user_id = %s"
NONKEY_Q = "select order_id, amount from orders where email = %s"


def plan(conn, sql: str) -> list[str]:
    return [r[0] for r in conn.execute("explain (costs off) " + sql).fetchall()]


def time_query(conn, sql: str, arg) -> float:
    t = time.perf_counter()
    conn.execute(sql, (arg,)).fetchall()
    return (time.perf_counter() - t) * 1000


CONFIGS = [
    ("with the shard key", KEY_Q, {"async_capable": "false", "parallel_commit": "false"}),
    ("no key, shards in turn", NONKEY_Q, {"async_capable": "false", "parallel_commit": "false"}),
    ("no key, async_capable", NONKEY_Q, {"async_capable": "true", "parallel_commit": "false"}),
    ("no key, async + parallel_commit", NONKEY_Q, {"async_capable": "true", "parallel_commit": "true"}),
]


def measure(stand: Stand, users: int) -> dict[str, list[float]]:
    rng = random.Random(SEED)
    medians: dict[str, list[float]] = {label: [] for label, _, _ in CONFIGS}
    for rnd in range(ROUNDS):
        order = CONFIGS if rnd % 2 == 0 else list(reversed(CONFIGS))
        for label, sql, opts in order:
            for k, v in opts.items():
                stand.set_option(k, v)
            with stand.connect() as conn:
                arg0 = 1 if sql is KEY_Q else "user1@example.com"
                time_query(conn, sql, arg0)  # прогрев соединений к шардам
                times = []
                for _ in range(PER_ROUND):
                    u = rng.randint(1, users)
                    arg = u if sql is KEY_Q else f"user{u}@example.com"
                    times.append(time_query(conn, sql, arg))
            medians[label].append(statistics.median(times))
    return medians


def main() -> None:
    print(f"{version()} | routing a query to shards through postgres_fdw")
    print("  shards are databases on one server; each sits behind its own delaying relay")
    print(f"  {ROUNDS} rounds x {PER_ROUND} queries, median of a round; seed {SEED}")

    users = 20_000
    stand = Stand(4, 0.0, users=users)
    stand.build()

    show("1. WHERE THE QUERY GOES: THE PLAN WITH AND WITHOUT THE SHARD KEY")
    with stand.connect() as conn:
        print("  where user_id = 42  (the table is partitioned by hash of user_id)")
        for line in plan(conn, "select * from orders where user_id = 42"):
            print(f"    {line}")
        print("  where email = 'user42@example.com'")
        for line in plan(conn, "select * from orders where email = 'user42@example.com'"):
            print(f"    {line}")
        stand.set_option("async_capable", "true")
        print("  the same, with async_capable 'true' on the foreign servers")
        for line in plan(conn, "select * from orders where email = 'user42@example.com'"):
            print(f"    {line}")
        stand.set_option("async_capable", "false")

    show("2. WHAT THE COORDINATOR SENDS TO A SHARD FOR ONE QUERY")
    print("  decoded from the wire by the relay; connections already open")
    with stand.connect() as conn:
        conn.execute(NONKEY_Q, ("user41@example.com",)).fetchall()  # прогрев
        for link in stand.links:
            link.trace = []
        conn.execute(NONKEY_Q, ("user42@example.com",)).fetchall()
        traces = [link.trace for link in stand.links]
        for link in stand.links:
            link.trace = None
        print("  query without the shard key, shard 0:")
        for n, message in enumerate(traces[0], 1):
            first = message.splitlines()[0]
            print(f"    {n}. {first}")
        per_shard = [len(t) for t in traces]
        print(f"  messages per shard, all shards: {per_shard}")
        conn.execute(KEY_Q, (41,)).fetchall()
        for link in stand.links:
            link.trace = []
        conn.execute(KEY_Q, (42,)).fetchall()
        per_shard_key = [len(link.trace) for link in stand.links]
        for link in stand.links:
            link.trace = None
        print(f"  query with the shard key, messages per shard: {per_shard_key}")
    print()
    print("  Five exchanges per shard touched: open a remote transaction, declare a")
    print("  cursor, fetch, close it, commit. A query without the key pays all five")
    print("  on every shard.")
    stand.drop()

    delays = (1.0, 3.0)
    results: dict[tuple[int, float], dict[str, list[float]]] = {}
    for n in (4, 8):
        stand = Stand(n, 0.0, users=users)
        stand.build()
        for d in delays:
            stand.attach_links(d)
            results[(n, d)] = measure(stand, users)
        stand.drop()

    block = 3
    for n in (4, 8):
        d = delays[0]
        show(f"{block}. TIME PER QUERY, {n} SHARDS, {d:g} MS EACH WAY")
        block += 1
        r = results[(n, d)]
        print(f"  {'configuration':<34} {'median':>8} {'best round':>11} {'worst round':>12} {'spread':>7}")
        for label, _, _ in CONFIGS:
            m = r[label]
            med = statistics.median(m)
            print(
                f"  {label:<34} {med:>6.1f} ms {min(m):>8.1f} ms {max(m):>9.1f} ms"
                f" {(max(m) - min(m)) / med:>6.0%}"
            )
        key = statistics.median(r["with the shard key"])
        seq = statistics.median(r["no key, shards in turn"])
        both = statistics.median(r["no key, async + parallel_commit"])
        wins = sum(
            1 for a, b in zip(r["no key, async + parallel_commit"], r["no key, shards in turn"]) if a < b
        )
        print()
        print(f"  no key in turn / with the key:              {seq / key:.1f}x")
        print(f"  no key async + parallel_commit / with key:  {both / key:.1f}x")
        print(f"  async + parallel_commit faster than in turn in {wins} of {ROUNDS} rounds")

    show(f"{block}. HOW MANY ROUND TRIPS ARE PAID ONE AFTER ANOTHER")
    block += 1
    lo, hi = delays
    print(f"  each configuration timed at {lo:g} ms and at {hi:g} ms each way;")
    print(f"  (time at {hi:g} - time at {lo:g}) / ({2 * hi:g} - {2 * lo:g} ms) = round trips in sequence")
    # Сколько кругов ДОЛЖНО получиться, если из пяти обменов блока 2
    # совмещается только то, что обещает документация: выборка при
    # async_capable и фиксация при parallel_commit. Совмещённый обмен стоит
    # один круг на все шарды вместо одного на каждый, то есть экономит N - 1.
    expected = {
        "with the shard key": lambda n: 5,
        "no key, shards in turn": lambda n: 5 * n,
        "no key, async_capable": lambda n: 5 * n - (n - 1),
        "no key, async + parallel_commit": lambda n: 5 * n - 2 * (n - 1),
    }
    print(f"  {'configuration':<34} {'4 shards':>9} {'expected':>9} {'8 shards':>9} {'expected':>9}")
    for label, _, _ in CONFIGS:
        cells = []
        for n in (4, 8):
            t_lo = statistics.median(results[(n, lo)][label])
            t_hi = statistics.median(results[(n, hi)][label])
            cells.append((t_hi - t_lo) / (2 * (hi - lo)))
        e4, e8 = expected[label](4), expected[label](8)
        print(f"  {label:<34} {cells[0]:>9.1f} {e4:>9} {cells[1]:>9.1f} {e8:>9}")
    print()
    print("  expected: of the five exchanges, fetch overlaps under async_capable and")
    print("  commit under parallel_commit; an overlapped exchange costs one round trip")
    print("  for all shards instead of one per shard, saving N - 1")
    print()
    print("  With the key the count is five - the five exchanges of block 2 - and it")
    print("  does not depend on the number of shards. Without the key, in turn, it is")
    print("  five per shard. async_capable overlaps the fetches and parallel_commit")
    print("  the commits; opening the remote transaction, declaring the cursor and")
    print("  closing it still happen shard after shard, so the count keeps growing")
    print("  with the number of shards.")

    show(f"{block}. WHAT GROWS WITH THE NUMBER OF SHARDS, {lo:g} MS EACH WAY")
    print(f"  {'configuration':<34} {'4 shards':>9} {'8 shards':>9} {'8 / 4':>7}")
    for label, _, _ in CONFIGS:
        a = statistics.median(results[(4, lo)][label])
        b = statistics.median(results[(8, lo)][label])
        print(f"  {label:<34} {a:>6.1f} ms {b:>6.1f} ms {b / a:>6.2f}x")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(1)