"""Working with the archive's virtiofs mount without exhausting it.

The archive sits on a virtiofs share whose daemon holds a host file descriptor for
every inode the guest has looked up.  Touching a few hundred thousand files fills the
guest's dentry cache, the daemon runs out of descriptors, and the mount then returns
ENFILE ("too many open files in system") for *any* subsequent open -- reads, writes,
even starting a Python interpreter stored on it.  It does not recover on its own.

The durable fix is host-side (``virtiofsd --inode-file-handles=prefer``, which stores
compact handles instead of descriptors).  Until that is in place, anything that
enumerates a large part of the archive has to periodically hand the descriptors back,
which means persuading the guest kernel to evict dentries so it sends FORGET.

Note that ``vm.vfs_cache_pressure`` does **not** do this.  It biases which caches the
kernel drops once it has decided to reclaim; it does not cause reclaim.  On a machine
with free memory the kernel never feels pressure, so the dentry cache grows unbounded
whatever that setting says.  Reclaim has to be asked for explicitly.
"""

import atexit

#: Operations between reclaims during a long traversal.  Low enough that the mount
#: never approaches its ceiling, high enough that the cost is amortised.
RELIEF_INTERVAL = 12_000


def drop_caches():
    """Ask the kernel to free dentries and inodes.  True if it worked.

    Needs root, and is worth having it: this is direct and instant, where the
    fallback has to allocate tens of gigabytes to provoke the same reclaim.  Run
    long traversals under sudo and they will take this path.
    """
    try:
        with open("/proc/sys/vm/drop_caches", "w") as handle:
            handle.write("2\n")  # 2 = dentries and inodes; page cache is not the issue
        return True
    except OSError:
        return False


def reclaimable_kb():
    """Size of the reclaimable slab, which is where dentries and inodes live."""
    try:
        with open("/proc/meminfo") as handle:
            for line in handle:
                if line.startswith("SReclaimable:"):
                    return int(line.split()[1])
    except (OSError, ValueError, IndexError):
        pass
    return 0


def available_gib():
    """Memory the kernel thinks can be handed out without swapping."""
    try:
        with open("/proc/meminfo") as handle:
            for line in handle:
                if line.startswith("MemAvailable:"):
                    return int(line.split()[1]) / (1024 * 1024)
    except (OSError, ValueError, IndexError):
        pass
    return 0.0


#: Memory left unclaimed by the fallback, in GiB.  Pressure has to be genuine to
#: make the kernel reclaim slab at all, so this is a *reserve* rather than a
#: fraction: capping at some proportion of available memory always leaves headroom,
#: the kernel never feels squeezed, and the allocation frees nothing while still
#: costing the machine several gigabytes -- worse than not trying.
RELIEF_RESERVE_GIB = 2.0
#: Absolute cap, so a machine with vast memory does not get an unbounded allocation.
RELIEF_MAX_GIB = 48


#: Below this much reclaimable slab there is nothing worth freeing -- it is the
#: working set of live processes, not archive inodes.  Pushing memory down to the
#: reserve to shave a few hundred megabytes off it costs the machine far more than
#: it gains, and during a long job this runs every RELIEF_INTERVAL items.
RELIEF_FLOOR_KB = 1024 * 1024  # 1 GiB


def release_handles(budget_gib=None, reserve_gib=RELIEF_RESERVE_GIB,
                    floor_kb=RELIEF_FLOOR_KB):
    """Force the kernel to shrink its dentry/inode cache.  True if anything was freed.

    Prefers ``drop_caches``; falls back to allocating memory until the reclaimable
    slab actually shrinks, or until only `reserve_gib` remains available.

    Two things make the fallback awkward, both learned the hard way:

    * A *fixed* budget silently does nothing once the machine has free memory --
      10 GiB against 17 GiB free reclaims zero while reporting success.
    * A budget capped at a *fraction* of available memory has the same failure for
      the same reason: it never applies real pressure.  14 GiB of a 23 GiB
      allowance freed nothing at all.

    So it pushes until the kernel actually gives ground, stopping at a fixed reserve
    rather than a proportion, and reports honestly whether the slab moved.
    """
    if drop_caches():
        return True

    before = reclaimable_kb()
    if before <= 0:
        return False
    if before < floor_kb:
        return True  # already small; nothing here is worth the memory churn

    target = before // 2
    if budget_gib is None:
        budget_gib = RELIEF_MAX_GIB
    blocks = []
    try:
        for _ in range(int(budget_gib)):
            if reclaimable_kb() <= target:
                return True
            if available_gib() <= reserve_gib:
                break  # as much pressure as is safe to apply
            blocks.append(bytearray(1024 * 1024 * 1024))
    except MemoryError:
        pass
    finally:
        blocks.clear()
    return reclaimable_kb() <= target


class Reliever:
    """Hands file handles back periodically during a bulk traversal.

    Every long pass over the archive -- indexing, detecting, filling, rendering --
    accumulates dentries that pin handles in the host's virtiofs daemon, and the
    mount eventually refuses *every* open, taking unrelated software on the machine
    down with it.  Reclaim was originally wired only into the migration and the
    index build; the bench jobs, which run for hours over tens of thousands of
    frames, had none, and duly exhausted the mount.

    Bulk loops construct one of these and call :meth:`tick` per item.
    """

    def __init__(self, interval=RELIEF_INTERVAL, enabled=True, label="", on_exit=True):
        self.interval = interval
        self.enabled = enabled
        self.label = label
        self.on_exit = on_exit
        self.seen = 0
        self.releases = 0
        self._next = interval
        self._registered = False

    def tick(self, count=1):
        """Record `count` items processed, reclaiming if enough have gone by."""
        self.seen += count
        if self.enabled and self.on_exit and not self._registered:
            # Reclaiming only *during* a run leaves the machine loaded once it ends:
            # the inodes a traversal cached stay pinned, and with them the host's
            # file handles, until something else forces reclaim.  Every bulk job has
            # to hand them back when it finishes, not merely while it runs.
            atexit.register(self.finish)
            self._registered = True
        if not self.enabled or self.seen < self._next:
            return False
        self._next = self.seen + self.interval
        self._release("after {} items".format(self.seen))
        return True

    def finish(self):
        """Reclaim once at the end of a run.  Idempotent; safe to call twice."""
        if not self.enabled or self.seen == 0:
            return False
        self.enabled = False  # nothing more to do for this traversal
        self._release("on finish, {} items".format(self.seen))
        return True

    def _release(self, why):
        freed = release_handles()
        self.releases += 1
        if self.label:
            print(f"  [{self.label}] reclaimed handles {why} "
                  f"({'ok' if freed else 'no change'})", flush=True)

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.finish()
        return False