devplacepy/devplacepy/services/containers/workspace/quota.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

from dataclasses import dataclass

from devplacepy.database import get_int_setting, get_table

RULES_TABLE = "workspace_quota_rules"

DEFAULTS: dict[str, int] = {
    "max_workspaces": 2,
    "max_tunnels": 5,
    "disk_quota_mb": 2048,
    "egress_quota_mb": 10240,
    "idle_warn_minutes": 45,
    "idle_stop_minutes": 60,
    "retention_days": 14,
    "purge_after_days": 7,
    "disk_warn_percent": 80,
    "cpu_millicores": 2000,
    "memory_mb": 2048,
}

SETTING_KEYS: dict[str, str] = {
    "max_workspaces": "workspace_max_per_user",
    "max_tunnels": "workspace_max_tunnels",
    "disk_quota_mb": "workspace_disk_quota_mb",
    "egress_quota_mb": "workspace_egress_quota_mb",
    "idle_warn_minutes": "workspace_idle_warn_minutes",
    "idle_stop_minutes": "workspace_idle_stop_minutes",
    "retention_days": "workspace_retention_days",
    "purge_after_days": "workspace_purge_after_days",
    "disk_warn_percent": "workspace_disk_warn_percent",
    "cpu_millicores": "workspace_cpu_millicores",
    "memory_mb": "workspace_memory_mb",
}

RULE_COLUMNS = (
    "max_workspaces",
    "max_tunnels",
    "disk_quota_mb",
    "egress_quota_mb",
    "idle_stop_minutes",
    "retention_days",
    "cpu_millicores",
    "memory_mb",
)

INSTANCE_OVERRIDE_COLUMNS = (
    "cpu_millicores",
    "memory_mb",
    "disk_quota_mb",
)


def format_cpu(millicores: int) -> str:
    if millicores <= 0:
        return ""
    return f"{millicores / 1000:.3f}".rstrip("0").rstrip(".")


def format_memory(megabytes: int) -> str:
    if megabytes <= 0:
        return ""
    return f"{megabytes}m"


@dataclass(frozen=True)
class Limits:
    max_workspaces: int
    max_tunnels: int
    disk_quota_mb: int
    egress_quota_mb: int
    idle_warn_minutes: int
    idle_stop_minutes: int
    retention_days: int
    purge_after_days: int
    disk_warn_percent: int
    cpu_millicores: int
    memory_mb: int

    def disk_quota_bytes(self) -> int:
        return self.disk_quota_mb * 1024 * 1024

    def egress_quota_bytes(self) -> int:
        return self.egress_quota_mb * 1024 * 1024

    def cpu_limit(self) -> str:
        return format_cpu(self.cpu_millicores)

    def mem_limit(self) -> str:
        return format_memory(self.memory_mb)


def _global_value(key: str) -> int:
    return get_int_setting(SETTING_KEYS[key], DEFAULTS[key])


def _rule_for(owner_kind: str, owner_id: str) -> dict | None:
    if not owner_id:
        return None
    table = get_table(RULES_TABLE)
    return table.find_one(owner_kind=owner_kind, owner_id=owner_id, deleted_at=None)


def resolve(user_uid: str = "", instance: dict | None = None) -> Limits:
    rule = _rule_for("user", user_uid) if user_uid else None
    values: dict[str, int] = {}
    for key in DEFAULTS:
        value = _global_value(key)
        if rule and key in RULE_COLUMNS:
            override = rule.get(key)
            if override:
                value = int(override)
        if instance and key in INSTANCE_OVERRIDE_COLUMNS:
            instance_override = instance.get(f"workspace_{key}")
            if instance_override:
                value = int(instance_override)
        values[key] = max(0, value)
    if values["idle_warn_minutes"] >= values["idle_stop_minutes"]:
        values["idle_warn_minutes"] = max(1, values["idle_stop_minutes"] - 1)
    return Limits(**values)


def percent_used(used: int, quota: int) -> int:
    if quota <= 0:
        return 0
    return min(100, int(used * 100 / quota))