# SPDX-License-Identifier: Apache-2.0
"""OPT_PLAN H7 (``BEVFUSION_HOST_MALLOC``): keep glibc from returning freed per-frame memory to the kernel.

A frame allocates and frees tens of MB of short-lived arrays (the rulebook maps, the voxelizer's per-point arrays,
the decoded cloud). With glibc's defaults the freed top of the heap is trimmed (and large blocks are unmapped), so
the next frame takes those pages again as minor page faults: 15-17 MB (3.8-4.3 k faults, ~9 ms) per PandaSet
pre-processing call on this host (OPT_REPORT round 7). ``mallopt`` is process-wide, so this is applied once at model
load; it only changes when memory goes back to the kernel (the resident size stays at the per-frame peak).

- ``default``: glibc as is;
- ``trim``: ``M_TRIM_THRESHOLD`` = 512 MiB (no top-of-heap trim below it) and ``M_MMAP_THRESHOLD`` = 32 MiB (glibc's
  maximum). Setting either one freezes glibc's dynamic mmap threshold; frozen at its 128 KiB start (the model loads
  this before its large arrays) every per-frame array would be a fresh mapping, which round 7 measured as slower
  than the default, so the threshold is set explicitly;
- ``arena1``: ``trim`` + ``M_ARENA_MAX`` = 1 (worker threads allocate from the main heap too, so their frames reuse
  the same pages; the threads then share the arena lock).
"""
from __future__ import annotations

import ctypes
import ctypes.util
from typing import Optional

__all__ = ["MODES", "apply"]

MODES = ("default", "trim", "arena1")
M_TRIM_THRESHOLD = -1
M_ARENA_MAX = -8
M_MMAP_THRESHOLD = -3
MMAP_BYTES = 32 << 20
TRIM_BYTES = 512 << 20
_APPLIED: Optional[str] = None


def apply(mode: str) -> Optional[str]:
    """Apply ``mode`` once per process (later calls with another mode are ignored); returns the applied mode or None
    (``default``, or not glibc)."""
    global _APPLIED
    if mode not in MODES:
        raise ValueError(f"HOST_MALLOC must be one of {MODES}")
    if mode == "default" or _APPLIED is not None:
        return _APPLIED
    try:
        libc = ctypes.CDLL(ctypes.util.find_library("c") or "libc.so.6")
        mallopt = libc.mallopt
    except (OSError, AttributeError):    # not glibc (musl, macOS): nothing to tune
        return None
    mallopt.argtypes = [ctypes.c_int, ctypes.c_int]
    mallopt.restype = ctypes.c_int
    ok = mallopt(M_MMAP_THRESHOLD, MMAP_BYTES) == 1 and mallopt(M_TRIM_THRESHOLD, TRIM_BYTES) == 1
    if ok and mode == "arena1":
        ok = mallopt(M_ARENA_MAX, 1) == 1
    _APPLIED = mode if ok else None
    return _APPLIED