"""Канал между датацентрами: ретранслятор с задержкой и рубильником.

ПОЧЕМУ РЕТРАНСЛЯТОР, А НЕ `tc netem`. Настоящий эмулятор сети ставится
дисциплиной очереди `netem` на сетевом интерфейсе. Она требует прав и, что
важнее здесь, поддержки в ядре: `CONFIG_NET_SCH_NETEM`. В песочнице, где
снимались эти числа, ядро собрано без неё — `tc qdisc add ... netem` отвечает
`Error: Specified qdisc kind is unknown`, и никакими правами это не лечится.
Ретранслятор в пользовательском коде даёт наблюдаемую величину того же рода
(известная задержка между двумя портами) и не трогает ничего за пределами
процесса.

ЧТО ЭТО ЗНАЧИТ ДЛЯ ВЫВОДОВ. Ретранслятор — приближение задержки, а не
эмулятор на уровне пакетов, и разница между ними измерима. Она измерена:
`relaycheck.py` считает, во что обходится сам стенд при нулевой задержке,
сколько раз задержка применяется к одной операции и что меняет замена
сериализующего ретранслятора на конвейерный. Читать выводы, снятые через этот
канал, стоит вместе с блоками 11-14.

ДВА РЕТРАНСЛЯТОРА, И РАЗНИЦА МЕЖДУ НИМИ СОДЕРЖАТЕЛЬНА.

`SerializingLink` — первая версия стенда. Она спит между `recv` и `sendall` в
одном потоке:

    data = src.recv(65536); time.sleep(delay); dst.sendall(data)

Пока поток спит, он не читает следующую порцию. Поэтому k порций, пришедших
подряд, стоят k задержек ПОДРЯД, а не одну: канал не только удлиняется, но и
перестаёт пропускать больше одной порции за время полёта. Для протокола
«запрос — ответ» разницы почти нет; для потока репликации и для `WAIT`,
где по каналу за одну операцию проходит несколько порций, разница есть — и
блок 13 её печатает.

`DelayedLink` — то, чем стенд стал после разбора. Чтение и отправка разнесены:
читающий поток кладёт порцию в очередь со сроком доставки `сейчас + delay` и
сразу читает дальше, отправляющий поток ждёт срок и отправляет. Порции,
пришедшие одновременно, уходят одновременно — это и есть постоянная задержка
распространения. Границ сообщений TCP по-прежнему не хранит, так что и это
приближение; но приближение, у которого задержка не зависит от того, на
сколько порций разбился поток.

ЧТО ЗДЕСЬ ЕЩЁ ЕСТЬ — рубильник `cut()`. Разрыв между датацентрами это не
«стало медленно», это «перестало ходить вовсе», и изображать его увеличенной
задержкой нельзя: узел, ждущий ответа, и узел, получивший сброс соединения,
ведут себя по-разному. `cut()` закрывает все установленные соединения и
перестаёт принимать новые — с точки зрения обеих сторон канал оборван.

СЧЁТЧИКИ. Оба ретранслятора считают порции и байты в каждую сторону
(`stats()`). Это не отладка, а измерительный прибор: число порций и есть то,
сколько раз была применена задержка.
"""

from __future__ import annotations

import queue
import socket
import threading
import time
from dataclasses import dataclass, field


@dataclass
class LinkStats:
    """Сколько раз задержка была применена и к скольким байтам."""

    chunks_up: int = 0
    chunks_down: int = 0
    bytes_up: int = 0
    bytes_down: int = 0
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    def add(self, upstream: bool, size: int) -> None:
        with self._lock:
            if upstream:
                self.chunks_up += 1
                self.bytes_up += size
            else:
                self.chunks_down += 1
                self.bytes_down += size

    @property
    def chunks(self) -> int:
        return self.chunks_up + self.chunks_down


class _BaseLink:
    """Общая часть: слушающий порт, приём соединений, рубильник."""

    def __init__(self, upstream_port: int, delay_ms: float, nodelay: bool = True) -> None:
        self.upstream_port = upstream_port
        self.delay_ms = delay_ms
        # nodelay=False воспроизводит прежний стенд. Оставлено не для
        # совместимости, а затем, чтобы блок 11 мог включить и выключить ровно
        # одну вещь и показать, чего она стоила.
        self.nodelay = nodelay
        self.alive = True
        self.stats = LinkStats()
        self._conns: list[socket.socket] = []
        self._lock = threading.Lock()
        self.sock = socket.socket()
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(("127.0.0.1", 0))
        self.sock.listen(128)
        self.port = self.sock.getsockname()[1]
        threading.Thread(target=self._accept_loop, daemon=True).start()

    def reset_stats(self) -> None:
        self.stats = LinkStats()

    def cut(self) -> None:
        """Оборвать канал: закрыть установленное и не принимать новое."""
        self.alive = False
        with self._lock:
            conns, self._conns = self._conns, []
        for c in conns:
            try:
                c.close()
            except OSError:
                pass

    def _accept_loop(self) -> None:
        while True:
            try:
                client, _ = self.sock.accept()
            except OSError:
                return
            if not self.alive:
                client.close()
                continue
            threading.Thread(target=self._serve, args=(client,), daemon=True).start()

    def _serve(self, client: socket.socket) -> None:
        try:
            server = socket.create_connection(("127.0.0.1", self.upstream_port))
        except OSError:
            client.close()
            return
        # TCP_NODELAY НА ОБА СОКЕТА, И ЭТО НЕ МИКРООПТИМИЗАЦИЯ.
        #
        # Redis ставит этот флаг на своих соединениях сам. Ретранслятор создаёт
        # ДВА НОВЫХ соединения, и на них флага не было. Отсюда пара «алгоритм
        # Нейгла на отправителе — отложенное подтверждение на получателе»:
        # маленькая порция ждёт подтверждения предыдущей, а подтверждение
        # придерживается до 40 мс. Ровно эти 40 мс и оседали в каждом замере,
        # где по каналу шло несколько мелких сообщений подряд, — и именно они
        # стояли в блоке 10 как «слагаемое, от расстояния не зависящее, чем
        # вызвано — замер не устанавливает». Установил блок 11.
        if self.nodelay:
            for s in (client, server):
                s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        with self._lock:
            self._conns += [client, server]
        # upstream=True — порция идёт от клиента к серверу за каналом.
        self._start_pipe(client, server, upstream=True)
        self._start_pipe(server, client, upstream=False)

    def _start_pipe(self, src: socket.socket, dst: socket.socket, upstream: bool) -> None:
        raise NotImplementedError

    @staticmethod
    def _shutdown(*socks: socket.socket) -> None:
        for s in socks:
            try:
                s.close()
            except OSError:
                pass


class SerializingLink(_BaseLink):
    """Первая версия стенда: сон между чтением и отправкой, в одном потоке.

    Оставлена НЕ для совместимости, а как предмет измерения: блок 13 сравнивает
    её с конвейерной версией на одних и тех же операциях, и без неё сравнивать
    было бы не с чем.
    """

    def _start_pipe(self, src: socket.socket, dst: socket.socket, upstream: bool) -> None:
        threading.Thread(target=self._pump, args=(src, dst, upstream), daemon=True).start()

    def _pump(self, src: socket.socket, dst: socket.socket, upstream: bool) -> None:
        try:
            while True:
                data = src.recv(65536)
                if not data:
                    break
                self.stats.add(upstream, len(data))
                # Задержка ДО отправки и В ТОМ ЖЕ потоке: пока спим, следующая
                # порция не читается. Отсюда и вся разница с конвейерной версией.
                time.sleep(self.delay_ms / 1000)
                dst.sendall(data)
        except OSError:
            pass
        finally:
            self._shutdown(src, dst)


class DelayedLink(_BaseLink):
    """Конвейерная версия: чтение не останавливается на время полёта.

    Читающий поток проставляет порции срок доставки и сразу читает дальше;
    отправляющий поток ждёт срок и отправляет. Порции, пришедшие в один момент,
    уходят в один момент, позже на `delay_ms`, — это и есть постоянная задержка
    распространения, а не задержка на порцию.
    """

    _EOF = object()

    def _start_pipe(self, src: socket.socket, dst: socket.socket, upstream: bool) -> None:
        q: queue.Queue = queue.Queue()
        threading.Thread(target=self._reader, args=(src, q, upstream), daemon=True).start()
        threading.Thread(target=self._sender, args=(q, src, dst), daemon=True).start()

    def _reader(self, src: socket.socket, q: queue.Queue, upstream: bool) -> None:
        delay = self.delay_ms / 1000
        try:
            while True:
                data = src.recv(65536)
                if not data:
                    break
                self.stats.add(upstream, len(data))
                q.put((time.monotonic() + delay, data))
        except OSError:
            pass
        finally:
            q.put((time.monotonic() + delay, self._EOF))

    def _sender(self, q: queue.Queue, src: socket.socket, dst: socket.socket) -> None:
        try:
            while True:
                due, data = q.get()
                left = due - time.monotonic()
                if left > 0:
                    time.sleep(left)
                if data is self._EOF:
                    break
                dst.sendall(data)
        except OSError:
            pass
        finally:
            self._shutdown(src, dst)