path: root/packages/meshbay-node/src/meshbay_node/transfers.py

"""
Transfer slots: how many downloads and uploads a node runs at once.

A download is invisible to the node today. `pipelinedDownload` sends eight
independent `file_req` messages and reassembles the answers; nothing tells the
node a transfer started, and nothing tells it one ended. There is nothing to
count and so nothing to cap — which is why this exists before any cap does.

The unit is the **lease**: the node's record that a peer is transferring
something, held for the length of the transfer and released by name. Six
properties are load-bearing, and each one is a decision:

  - **`tr` is drawn by the client**, like `upload_id`. Re-opening after a
    reconnect with the same `tr` is idempotent, so a reconnect cannot charge a
    member twice for one transfer.
  - **A lease is scoped to the connection**, never to the account. It dies with
    the session, which is what makes the primary reclaim deterministic.
  - **A lease covers a job, not a file.** A directory zip is dozens of files and
    one lease.
  - **Nothing is persisted.** A restart drops every session anyway; a lease that
    outlived the process would be a slot nothing can release.
  - **Leases are counted, not bytes.** What a slot protects is concurrency —
    open file handles, disk seeks, the channel buffer each transfer keeps full.
  - **Per-member first, then node-wide.** A member at their own cap queues
    behind their own transfers and never holds a node-wide slot a second member
    has none of. Reversed, whoever arrives first takes everything.

This module is deliberately free of asyncio and of the transport: it decides,
and the caller does the I/O. `sweep()` is called on a clock the caller owns, and
every method returns what changed so the caller can push it. That is what makes
the failure modes in §5 of ~/next/improve-downloads.md testable at all — a
queue that only reveals itself through a DataChannel is a queue nobody can
prove things about.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field

log = logging.getLogger(__name__)

DOWNLOAD = "download"
UPLOAD = "upload"
KINDS = (DOWNLOAD, UPLOAD)

# Node-wide defaults. The operator's own values arrive from node.toml/roster.db
# via `set_capacity` — these apply when they have said nothing.
DEFAULT_MAX_CONCURRENT = 8
# Per account, per group. Absent means this, not "unlimited": a group that
# predates the setting coming back unlimited would leave the node-wide cap as
# the only control, which is the situation this exists to end.
DEFAULT_MAX_PER_MEMBER = 2

# A grant nobody takes up is a slot nobody can use. Long enough for a client to
# send its first chunk request, short enough that a browser that died between
# the grant and that request does not hold a slot until the idle timeout.
GRANT_DEADLINE_SECS = 30.0
# Silence on a granted lease. The session dying is the primary reclaim and is
# immediate; this only catches a peer that vanished without the connection
# noticing, so it can afford to be generous.
IDLE_TIMEOUT_SECS = 120.0
# Per account, per kind. Unbounded queues are how a node runs out of memory
# politely; past this the client keeps the rest in its own list.
MAX_QUEUED_PER_MEMBER = 32
# How many times a lease may be granted and not taken up before it is closed
# rather than queued again. Without a bound the requeue is a permanent cycle,
# and a node logs the same reclaim every 30 s until it restarts.
MAX_MISSED_GRANTS = 3

# Why a lease ended, as it reaches the peer.
REASON_DONE = "done"
REASON_CANCELLED = "cancelled"
REASON_PAUSED = "paused"
REASON_FAILED = "failed"
REASON_SESSION_GONE = "session_gone"
REASON_IDLE = "idle"
REASON_NOT_TAKEN_UP = "not_taken_up"
REASON_ABANDONED = "abandoned"


@dataclass
class Lease:
    tr: str
    kind: str
    session_key: str
    user_id: str
    group_id: str
    bytes: int = 0
    chunks: int = 0
    state: str = "queued"          # "queued" | "granted"
    created_at: float = 0.0
    granted_at: float | None = None
    # Set the first time anything happens under this lease. Distinguishes "the
    # client never came back for its slot" from "the client went quiet": the
    # first is a grant to revoke and pass on, the second a transfer to reclaim.
    used: bool = False
    last_seen: float = 0.0
    # How many grants this lease has been given and not taken up. Bounded
    # because the requeue is otherwise a permanent cycle: revoked, put back,
    # granted again a millisecond later because there is room, revoked 30 s
    # later, for ever. Seen doing exactly that in a node's log, every 30 s,
    # minutes after the transfers involved had finished.
    missed_grants: int = 0

    @property
    def member(self) -> tuple[str, str]:
        return (self.group_id, self.user_id)


@dataclass
class TransferSlots:
    """Every lease on this node, and the queues behind them."""

    caps: dict[str, int] = field(
        default_factory=lambda: {k: DEFAULT_MAX_CONCURRENT for k in KINDS})
    # The node-wide default per member, per kind.
    per_member: dict[str, int] = field(
        default_factory=lambda: {k: DEFAULT_MAX_PER_MEMBER for k in KINDS})
    # Per-group overrides: {group_id: {kind: n}}. The cap is a group's setting
    # (its operator signs it), while the pools are the machine's — so this is
    # the one dimension that is not node-wide, and a lookup rather than a field.
    group_limits: dict[str, dict[str, int]] = field(default_factory=dict)
    leases: dict[str, Lease] = field(default_factory=dict)
    # FIFO of `tr`, per kind. Order is arrival; a member at their own cap is
    # skipped rather than blocking the head, or one member's limit would stall
    # the whole node.
    queues: dict[str, list[str]] = field(
        default_factory=lambda: {k: [] for k in KINDS})

    # ── counting ────────────────────────────────────────────────────────────

    def in_use(self, kind: str) -> int:
        return sum(1 for x in self.leases.values()
                   if x.kind == kind and x.state == "granted")

    def member_in_use(self, kind: str, member: tuple[str, str]) -> int:
        return sum(1 for x in self.leases.values()
                   if x.kind == kind and x.state == "granted"
                   and x.member == member)

    def queued_for(self, kind: str, member: tuple[str, str]) -> int:
        return sum(1 for tr in self.queues[kind]
                   if (x := self.leases.get(tr)) and x.member == member)

    def ahead_of(self, lease: Lease) -> int:
        """How many are in front of this one in its queue."""
        try:
            return self.queues[lease.kind].index(lease.tr)
        except ValueError:
            return 0

    def member_cap(self, kind: str, member: tuple[str, str]) -> int:
        """This member's cap in this group: the group's own, else the default.

        Absent means the default, never "unlimited" — a group that predates the
        setting coming back unlimited would leave the node-wide cap as the only
        control, which is the situation slots exist to end.
        """
        group_id = member[0]
        override = self.group_limits.get(group_id, {}).get(kind)
        if override is not None:
            return int(override)
        return self.per_member.get(kind, DEFAULT_MAX_PER_MEMBER)

    def _has_room(self, kind: str, member: tuple[str, str]) -> bool:
        # Per-member first: see the module docstring.
        if self.member_in_use(kind, member) >= self.member_cap(kind, member):
            return False
        return self.in_use(kind) < self.caps.get(kind, DEFAULT_MAX_CONCURRENT)

    # ── the operations a peer asks for ──────────────────────────────────────

    def open(self, *, tr: str, kind: str, session_key: str, user_id: str,
             group_id: str, bytes: int = 0, chunks: int = 0,
             now: float | None = None) -> tuple[Lease | None, str]:
        """Ask for a slot. Returns (lease, error_code); one of them is falsy.

        Idempotent on `tr`: re-opening a lease this session already holds
        returns it unchanged rather than charging for a second one. That is what
        makes a client's reconnect safe, and it is checked before anything else
        because every other branch below would otherwise double-count.
        """
        now = time.monotonic() if now is None else now
        if kind not in KINDS:
            return None, "bad_kind"
        existing = self.leases.get(tr)
        if existing is not None:
            if existing.session_key != session_key:
                # Someone else's lease id. Refused rather than adopted: a `tr`
                # is drawn at random by its owner, so a collision is either a
                # bug or a peer guessing, and neither should move a slot between
                # connections.
                return None, "not_your_transfer"
            return existing, ""

        member = (group_id, user_id)
        if self.queued_for(kind, member) >= MAX_QUEUED_PER_MEMBER:
            return None, "too_many_queued"

        lease = Lease(tr=tr, kind=kind, session_key=session_key,
                      user_id=user_id, group_id=group_id,
                      bytes=int(bytes or 0), chunks=int(chunks or 0),
                      created_at=now, last_seen=now)
        self.leases[tr] = lease
        if self._has_room(kind, member):
            self._grant(lease, now)
        else:
            self.queues[kind].append(tr)
        return lease, ""

    def _grant(self, lease: Lease, now: float) -> None:
        lease.state = "granted"
        lease.granted_at = now
        lease.last_seen = now
        lease.used = False

    def touch(self, tr: str, now: float | None = None) -> bool:
        """Something happened under this lease. False if it is not granted."""
        lease = self.leases.get(tr)
        if lease is None or lease.state != "granted":
            return False
        lease.used = True
        lease.missed_grants = 0
        lease.last_seen = time.monotonic() if now is None else now
        return True

    def close(self, tr: str, reason: str = REASON_DONE,
              now: float | None = None) -> tuple[Lease | None, list[Lease]]:
        """Give a slot back. Returns (the closed lease, newly granted ones).

        The only place a lease is destroyed, and the only caller of the pump —
        two functions that both released would be this repo's flow-control
        lesson one feature later.
        """
        lease = self.leases.pop(tr, None)
        if lease is None:
            return None, []
        if lease.tr in self.queues[lease.kind]:
            self.queues[lease.kind].remove(lease.tr)
        log.debug("transfer: closed %s (%s, %s)", tr[:8], lease.kind, reason)
        return lease, self._pump(lease.kind, now)

    def release_session(self, session_key: str,
                        now: float | None = None) -> tuple[list[Lease], list[Lease]]:
        """The connection is gone; everything it held goes with it.

        The deterministic reclaim, and the reason a lease is scoped to a
        connection rather than to an account: a tab closed, a browser quit and a
        network that dropped all arrive here, and none of them needs a timer.
        """
        gone = [x for x in self.leases.values() if x.session_key == session_key]
        for lease in gone:
            self.leases.pop(lease.tr, None)
            if lease.tr in self.queues[lease.kind]:
                self.queues[lease.kind].remove(lease.tr)
        granted: list[Lease] = []
        for kind in KINDS:
            if any(x.kind == kind for x in gone):
                granted.extend(self._pump(kind, now))
        return gone, granted

    def sweep(self, now: float | None = None) -> tuple[list[tuple[Lease, str]],
                                                       list[Lease]]:
        """Reclaim what the session teardown cannot see.

        Two different failures, deliberately told apart:
        a grant nobody took up (the client died between asking and starting)
        goes back to the tail of the queue; a granted transfer that has gone
        quiet is closed, and the peer is told, so its widget can offer a resume
        rather than sit on a lie.
        """
        now = time.monotonic() if now is None else now
        ended: list[tuple[Lease, str]] = []
        requeued = False
        for lease in list(self.leases.values()):
            if lease.state != "granted":
                continue
            if not lease.used and lease.granted_at is not None \
                    and now - lease.granted_at > GRANT_DEADLINE_SECS:
                lease.missed_grants += 1
                lease.granted_at = None
                if lease.missed_grants >= MAX_MISSED_GRANTS:
                    # It has had its chances. Closing it is what ends the cycle,
                    # and the peer is told so a client that is somehow still
                    # there can ask again from a clean state rather than hold a
                    # slot it has never once used.
                    self.leases.pop(lease.tr, None)
                    ended.append((lease, REASON_ABANDONED))
                else:
                    lease.state = "queued"
                    self.queues[lease.kind].append(lease.tr)
                    ended.append((lease, REASON_NOT_TAKEN_UP))
                    requeued = True
            elif lease.used and now - lease.last_seen > IDLE_TIMEOUT_SECS:
                self.leases.pop(lease.tr, None)
                ended.append((lease, REASON_IDLE))
        granted: list[Lease] = []
        if ended or requeued:
            for kind in KINDS:
                granted.extend(self._pump(kind, now))
        return ended, granted

    # ── the queue ───────────────────────────────────────────────────────────

    def _pump(self, kind: str, now: float | None = None) -> list[Lease]:
        """Grant to whoever can start, in arrival order, skipping who cannot.

        Called from exactly one place per release. Walking past a member who is
        at their own cap is the whole reason this is a walk and not a `pop(0)`:
        granting strictly in order lets one member's limit stall every other
        member behind them.
        """
        now = time.monotonic() if now is None else now
        granted: list[Lease] = []
        for tr in list(self.queues[kind]):
            lease = self.leases.get(tr)
            if lease is None:                      # closed while queued
                self.queues[kind].remove(tr)
                continue
            if self.in_use(kind) >= self.caps.get(kind, DEFAULT_MAX_CONCURRENT):
                break                              # the node is full; stop
            if not self._has_room(kind, lease.member):
                continue                           # this member is; skip them
            self.queues[kind].remove(tr)
            self._grant(lease, now)
            granted.append(lease)
        return granted

    # ── what the operator sees ──────────────────────────────────────────────

    def set_group_limits(self, group_id: str, limits: dict[str, int],
                         now: float | None = None) -> list[Lease]:
        """One group's per-member caps, as its operator signed them."""
        current = dict(self.group_limits.get(group_id, {}))
        for kind, value in limits.items():
            if kind in KINDS:
                current[kind] = max(1, int(value))
        self.group_limits[group_id] = current
        granted: list[Lease] = []
        for kind in KINDS:
            granted.extend(self._pump(kind, now))
        return granted

    def set_caps(self, *, node: dict[str, int] | None = None,
                 per_member: dict[str, int] | None = None,
                 now: float | None = None) -> list[Lease]:
        """Change a cap live. Raising one may start queued transfers at once.

        Lowering never interrupts a transfer that is running, for the same
        reason lowering the stream cap does not stop a film: the new value
        governs what starts next.
        """
        for kind, value in (node or {}).items():
            if kind in KINDS:
                self.caps[kind] = max(1, int(value))
        for kind, value in (per_member or {}).items():
            if kind in KINDS:
                self.per_member[kind] = max(1, int(value))
        granted: list[Lease] = []
        for kind in KINDS:
            granted.extend(self._pump(kind, now))
        return granted

    def snapshot(self) -> dict:
        """The whole picture, for `GET /api/transfers` and the summary log.

        When somebody reports a transfer stuck at "waiting", this is the only
        thing that will say whether the node ever had them in a queue.
        """
        return {
            "pools": {
                kind: {
                    "in_use": self.in_use(kind),
                    "cap": self.caps.get(kind, DEFAULT_MAX_CONCURRENT),
                    "per_member": self.per_member.get(kind,
                                                      DEFAULT_MAX_PER_MEMBER),
                    "queued": len(self.queues[kind]),
                } for kind in KINDS
            },
            "leases": [
                {
                    "tr": x.tr[:12],
                    "kind": x.kind,
                    "state": x.state,
                    "user_id": x.user_id,
                    "group_id": x.group_id,
                    "bytes": x.bytes,
                    "used": x.used,
                    "ahead": self.ahead_of(x) if x.state == "queued" else 0,
                }
                # Never a filename or a path: a lease carries none, and this is
                # the one place it would be tempting to add one for a prettier
                # log line.
                for x in sorted(self.leases.values(),
                                key=lambda l: (l.kind, l.state, l.created_at))
            ],
        }

    def summary(self) -> str:
        p = self.snapshot()["pools"]
        return " ".join(
            f"{kind[0]}={p[kind]['in_use']}/{p[kind]['cap']}"
            f"(q{p[kind]['queued']})" for kind in KINDS)


# ── Reads that carry no lease ───────────────────────────────────────────────

# How many distinct files one session may be reading at once without a lease.
#
# Browsing a group is never subject to a transfer slot — not the poster grid,
# not the covers, not opening a photo or a PDF to look at it. A member must be
# able to browse a group that is at capacity exactly as they browse an idle one.
# That is a requirement, and §3.4 of ~/next/improve-downloads.md satisfies it
# structurally: a transfer is what the transfers widget shows, and nothing else
# takes a slot.
#
# But "not leased" cannot mean "unbounded", or a client that simply omits `tr`
# transfers outside every cap and the caps are decoration.
#
# **The number comes from what the shipped client legitimately does, and it was
# wrong once already.** §3.4.1 of ~/next/improve-downloads.md argued for two,
# reasoning about *viewers*: a photo viewer shows one photo, a preview modal one
# document, and the second is for prefetching the next photo. That reasoning
# forgot the music player, which warms a read-ahead window — `prefetchDepth()`
# in music-player.js returns 5 on Wi-Fi, 3 otherwise — so playing an album has
# six files in flight and the fourth was refused with `transfer_required`.
# Reported the day MNP 3.0 shipped, as "I try to play a track and it tells me to
# download it instead".
#
# Twelve: six for the music read-ahead at its widest, two for a photo viewer
# and its own prefetch running at the same time in the same session, and the
# rest as headroom for the next app that reads ahead. `test_leaseless_reads.py`
# derives the floor from music-player.js itself, so raising the client's
# prefetch without raising this fails rather than reaching a person.
#
# Generosity is cheap here and refusal is not. This is a fairness control among
# cooperating clients, not a security boundary — a client that lies gets twelve
# files at a time instead of its member cap, which is bounded and audited —
# whereas refusing a legitimate read breaks a stated requirement.
#
# Deliberately a count of files and not a byte budget: a RAW photo out of a
# camera is 60-80 MB and is browsing, a 40 MB archive is a download, and no
# size threshold separates them. What separates them is which function asked.
#
# What it costs, stated plainly: a client that lies — labelling a bulk download
# as a view — gets two files at a time instead of its member cap. That is the
# residual, it is bounded, it is audited, and it is the same kind of statement
# as the cap itself. **This is a fairness control among cooperating clients**,
# not a defence against a member determined to saturate a node's disk. The
# answer to that member is `member revoke`.
MAX_LEASELESS_IN_FLIGHT = 12

# A leaseless read has no "close" message, so it ends when the last chunk goes
# out — or, when a viewer is closed mid-file and simply stops asking, when it
# has been quiet this long.
LEASELESS_IDLE_SECS = 60


class LeaselessReads:
    """
    The files one session is reading without a lease, and the bound on them.

    Per session rather than per member: this is not a resource pool, it is a
    ceiling on what one connection can do while claiming to be browsing. A
    member with three tabs open is browsing in three tabs, which is fine.
    """

    def __init__(self, limit: int = MAX_LEASELESS_IN_FLIGHT,
                 idle: float = LEASELESS_IDLE_SECS) -> None:
        self.limit = limit
        self.idle = idle
        self._seen: dict[str, float] = {}

    def admit(self, file_id: str, now: float | None = None) -> bool:
        """May this session read `file_id` without a lease right now?

        True for a file it is already reading, whatever the count: refusing a
        chunk halfway through a photo because the limit moved would be worse
        than never having admitted it.
        """
        when = time.monotonic() if now is None else now
        self._expire(when)
        if file_id in self._seen:
            self._seen[file_id] = when
            return True
        if len(self._seen) >= self.limit:
            return False
        self._seen[file_id] = when
        return True

    def finish(self, file_id: str) -> None:
        """The last chunk went out; the slot is free at once rather than in a
        minute."""
        self._seen.pop(file_id, None)

    def _expire(self, now: float) -> None:
        # A viewer closed mid-file stops asking and says nothing. Without this
        # the session would carry two dead entries and refuse every later
        # preview, which is the bound turning into a bug.
        for file_id, last in list(self._seen.items()):
            if now - last > self.idle:
                del self._seen[file_id]

    def __len__(self) -> int:
        return len(self._seen)