"""Лимит памяти на группу процессов и то, чем кончается его превышение.

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

ПОЧЕМУ cgroup v1. В этом контейнере смонтирована именно первая версия
(`/sys/fs/cgroup/memory`), и файлы называются так, как называются. Во второй
версии те же величины лежат в `memory.max`, `memory.events` и `memory.peak`;
имена другие, механизм тот же.

ПОЧЕМУ ПОДПИСИ ПО-АНГЛИЙСКИ. Урок существует в двух языках и цитирует запись
прогона дословно обеими версиями. Русские подписи означали бы либо два
прогона с расходящимися числами, либо кириллицу в английском материале.

ЧТО ЗДЕСЬ ИЗМЕРЯЕТСЯ. Ничего временного: только коды выхода, счётчики из
`/sys/fs/cgroup` и строки ядра из `dmesg`. Всё это — наблюдения, а не замеры
скорости, поэтому от машины они зависят слабо.

ТРЕБОВАНИЯ: право создавать подкаталог в `/sys/fs/cgroup/memory` и читать
`dmesg`. Без них скрипт честно скажет, чего не смог, и не станет печатать
выдуманное.

ЗАПУСК: python3 bench/cgroups/oom.py
Вывод: runs/oom.txt
"""

import os
import re
import shutil
import subprocess
import sys
import time

CG_ROOT = "/sys/fs/cgroup/memory"
MIB = 1024 * 1024

# Программа-подопытная: вписывает себя в группу и ест память страницами,
# КАСАЯСЬ каждой страницы. Без касания страницы не заселяются, и лимит
# памяти не срабатывает — учитывается заселённое, а не запрошенное.
EATER = """
import os, sys, time
group, want_mib = sys.argv[1], int(sys.argv[2])
hold = float(sys.argv[3]) if len(sys.argv) > 3 else 0.0
with open(group + "/cgroup.procs", "w") as fh:
    fh.write(str(os.getpid()))
blocks = []
for _ in range(want_mib // 4):
    block = bytearray(4 * 1024 * 1024)
    for offset in range(0, len(block), 4096):
        block[offset] = 1
    blocks.append(block)
print("allocated", len(blocks) * 4, "MiB", flush=True)
time.sleep(hold)
"""


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


def row(label: str, value: object) -> None:
    print(f"  {label:<40} {value}")


def read(group: str, name: str) -> str:
    with open(f"{group}/{name}", encoding="utf-8") as fh:
        return fh.read().strip()


def make_group(name: str, limit_mib: int) -> str | None:
    group = f"{CG_ROOT}/{name}"
    try:
        os.makedirs(group, exist_ok=True)
        with open(f"{group}/memory.limit_in_bytes", "w", encoding="utf-8") as fh:
            fh.write(str(limit_mib * MIB))
    except OSError as exc:
        print(f"  cannot create a cgroup here: {exc}")
        return None
    return group


def drop_group(group: str) -> None:
    try:
        os.rmdir(group)
    except OSError:
        pass


def run_eater(group: str, want_mib: int) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        [sys.executable, "-c", EATER, group, str(want_mib)],
        capture_output=True,
        text=True,
    )


def kernel_lines() -> list[str]:
    if shutil.which("dmesg") is None:
        return []
    done = subprocess.run(["dmesg"], capture_output=True, text=True)
    if done.returncode != 0:
        return []
    return done.stdout.splitlines()


# ------------------------------------------------------------------ 1
def block1() -> None:
    show("1. THE KILL: the limit is 64 MiB, the program asks for 1600")

    group = make_group("de_oom_kill", 64)
    if group is None:
        return
    done = run_eater(group, 1600)

    row("limit (memory.limit_in_bytes)", f"{int(read(group, 'memory.limit_in_bytes')) // MIB} MiB")
    row("exit code from waitpid", done.returncode)
    row("as the shell reports it", 128 - done.returncode if done.returncode < 0 else done.returncode)
    row("stdout of the program", repr(done.stdout))
    row("stderr of the program", repr(done.stderr))
    row("peak usage before the kill", f"{int(read(group, 'memory.max_usage_in_bytes')) / MIB:.1f} MiB")
    row("times the limit was hit (failcnt)", read(group, "memory.failcnt"))
    drop_group(group)

    print()
    print("  No MemoryError, no traceback, no last words: the process was not")
    print("  asked to stop, it was removed by SIGKILL.")


# ------------------------------------------------------------------ 2
def block2() -> None:
    show("2. THE KERNEL SAYS IT, THE APPLICATION DOES NOT")

    lines = kernel_lines()
    if not lines:
        row("dmesg", "not readable here")
        return
    oom = [ln for ln in lines if "oom-kill:constraint" in ln]
    killed = [ln for ln in lines if "Memory cgroup out of memory" in ln]
    if not oom:
        row("kernel lines about the kill", "none found")
        return

    constraint = re.search(r"constraint=(\w+)", oom[-1])
    memcg = re.search(r"oom_memcg=(\S+?),", oom[-1])
    row("constraint reported by the kernel", constraint.group(1) if constraint else "?")
    row("cgroup that ran out", memcg.group(1) if memcg else "?")
    if killed:
        name = re.search(r"Killed process \d+ \((\S+?)\)", killed[-1])
        row("victim named by the kernel", name.group(1) if name else "?")
    print()
    print("  This is the whole difference between the two logs: the kernel")
    print("  writes down what happened, the application writes nothing,")
    print("  because it never learned anything happened.")


# ------------------------------------------------------------------ 3
def block3() -> None:
    show("3. WHO GETS KILLED: the biggest, not the last one")

    group = make_group("de_oom_victim", 160)
    if group is None:
        return

    # Двое в одной группе: маленький и большой. Убивают того, у кого больше
    # заселённой памяти, а не того, кто попросил последним.
    small = subprocess.Popen(
        [sys.executable, "-c", EATER, group, "40", "30"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    time.sleep(1.5)
    big = subprocess.Popen(
        [sys.executable, "-c", EATER, group, "600"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    big_code = big.wait()
    time.sleep(0.3)
    small_alive = small.poll() is None

    row("small process (40 MiB, holding it)", "alive" if small_alive else "killed")
    row("big process (asked for 600 MiB)", "killed" if big_code == -9 else f"exit {big_code}")
    row("victim", "big" if big_code == -9 and small_alive else "small or both")

    small.kill()
    small.wait()
    drop_group(group)
    print()
    print("  The kernel picks by how much memory a task actually holds, not by")
    print("  who asked last. A small neighbour survives a big one's mistake.")


# ------------------------------------------------------------------ 4
def block4() -> None:
    show("4. THE SAME PROGRAM AND THE SAME REQUEST, A LIMIT THAT FITS")

    group = make_group("de_oom_room", 2048)
    if group is None:
        return
    done = run_eater(group, 1600)
    row("limit", f"{int(read(group, 'memory.limit_in_bytes')) // MIB} MiB")
    row("exit code", done.returncode)
    row("stdout of the program", done.stdout.strip() or "(empty)")
    row("times the limit was hit (failcnt)", read(group, "memory.failcnt"))
    drop_group(group)
    print()
    print("  Same code, same machine, different limit: nothing about the")
    print("  program changed. The environment decided the outcome.")


# ------------------------------------------------------------------ 5
def block5() -> None:
    show("5. PAGE CACHE COUNTS TOO, BUT IT CAN BE GIVEN BACK")

    # ЗАЧЕМ ЭТОТ БЛОК. Утверждение «в лимит группы попадает и файловый кеш»
    # звучит одинаково с «много читать с диска опасно» — а это разные вещи,
    # и разница проверяется здесь: читаем файл вчетверо больше лимита и
    # смотрим, убьют ли. Не убивают: страницы кеша ядро освобождает.
    path = "/tmp/de_cache_probe.bin"
    size_mib = 256
    if not os.path.exists(path) or os.path.getsize(path) != size_mib * MIB:
        with open(path, "wb") as fh:
            fh.write(os.urandom(MIB) * size_mib)

    # Кеш заряжается той группе, которая ЗАПОЛНИЛА его первой. Если файл уже
    # лежит в кеше после записи, чтение не потратит ни байта лимита — и блок
    # покажет ноль, ничего при этом не доказав. Поэтому кеш сбрасывается.
    subprocess.run(["sync"], check=False)
    dropped = True
    try:
        with open("/proc/sys/vm/drop_caches", "w", encoding="utf-8") as fh:
            fh.write("3")
    except OSError:
        dropped = False

    group = make_group("de_oom_cache", 64)
    if group is None:
        return
    reader = """
import os, sys
group, path = sys.argv[1], sys.argv[2]
with open(group + "/cgroup.procs", "w") as fh:
    fh.write(str(os.getpid()))
total = 0
with open(path, "rb") as fh:
    while True:
        chunk = fh.read(1024 * 1024)
        if not chunk:
            break
        total += len(chunk)
print("read", total // (1024 * 1024), "MiB", flush=True)
"""
    done = subprocess.run(
        [sys.executable, "-c", reader, group, path], capture_output=True, text=True
    )
    row("limit", f"{int(read(group, 'memory.limit_in_bytes')) // MIB} MiB")
    row("page cache dropped before the read", "yes" if dropped else "no (not permitted)")
    row("file read through the group", f"{size_mib} MiB")
    row("exit code", done.returncode)
    row("stdout of the program", done.stdout.strip() or "(empty)")
    row("peak usage", f"{int(read(group, 'memory.max_usage_in_bytes')) / MIB:.1f} MiB")
    row("times the limit was hit (failcnt)", read(group, "memory.failcnt"))
    drop_group(group)
    print()
    print("  Four times the limit went through the group and nobody died:")
    print("  cache pages are charged to it, and cache pages can be dropped.")
    print("  What cannot be dropped is anonymous memory, and that is what")
    print("  block 1 was about.")


def main() -> None:
    print(f"Python {sys.version.split()[0]} · Linux {os.uname().release} · cgroup v1")
    if not os.path.isdir(CG_ROOT):
        print(f"no {CG_ROOT} here: this machine uses another cgroup layout")
        return
    block1()
    block2()
    block3()
    block4()
    block5()


if __name__ == "__main__":
    main()