home-manager/firefox: Add memory control

Wrap Firefox with a privileged AutoConfig controller that unloads tabs under configurable MemAvailable hysteresis and exposes an on-demand reclaim command. Enable it by default for Linux GUI homes.

pkgs/firefox-memory-control/firefox-free-memory.py

@@ -0,0 +1,139 @@
+import argparse
+import json
+import os
+import re
+import sys
+import time
+import uuid
+from pathlib import Path
+
+
+SIZE_RE = re.compile(r'^([0-9]+(?:\.[0-9]+)?)\s*([kmgt]?i?b?)?$', re.I)
+DURATION_RE = re.compile(r'^([0-9]+(?:\.[0-9]+)?)\s*(ms|s|m|h)?$', re.I)
+
+
+def parse_size(value):
+    match = SIZE_RE.match(value)
+    if not match:
+        raise argparse.ArgumentTypeError(f'invalid size: {value!r}')
+
+    number = float(match.group(1))
+    suffix = (match.group(2) or 'b').lower().removesuffix('b').removesuffix('i')
+    powers = {'': 0, 'k': 1, 'm': 2, 'g': 3, 't': 4}
+    size = round(number * 1024 ** powers[suffix])
+    if size <= 0:
+        raise argparse.ArgumentTypeError('size must be greater than zero')
+    return size
+
+
+def parse_duration(value):
+    match = DURATION_RE.match(value)
+    if not match:
+        raise argparse.ArgumentTypeError(f'invalid duration: {value!r}')
+
+    number = float(match.group(1))
+    suffix = (match.group(2) or 's').lower()
+    factors = {'ms': 1, 's': 1000, 'm': 60_000, 'h': 3_600_000}
+    return round(number * factors[suffix])
+
+
+def format_size(size):
+    for suffix in ('TiB', 'GiB', 'MiB', 'KiB'):
+        unit = 1024 ** {'KiB': 1, 'MiB': 2, 'GiB': 3, 'TiB': 4}[suffix]
+        if size >= unit:
+            return f'{size / unit:.2f} {suffix}'
+    return f'{size} B'
+
+
+def runtime_root():
+    runtime_dir = os.environ.get('XDG_RUNTIME_DIR')
+    if not runtime_dir or not os.path.isabs(runtime_dir):
+        raise RuntimeError('XDG_RUNTIME_DIR is not set to an absolute path')
+    return Path(runtime_dir) / 'firefox-memory-control'
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Ask a running memory-controlled Firefox to unload tabs'
+    )
+    parser.add_argument('size', type=parse_size, help='desired reclaim amount, for example 2G')
+    parser.add_argument(
+        '--min-inactive',
+        type=parse_duration,
+        default=0,
+        metavar='DURATION',
+        help='only unload tabs inactive for this long (default: 0s)',
+    )
+    parser.add_argument(
+        '--timeout',
+        type=float,
+        default=60,
+        metavar='SECONDS',
+        help='maximum time to wait for Firefox (default: 60)',
+    )
+    args = parser.parse_args()
+
+    root = runtime_root()
+    requests = root / 'requests'
+    responses = root / 'responses'
+    requests.mkdir(mode=0o700, parents=True, exist_ok=True)
+    responses.mkdir(mode=0o700, parents=True, exist_ok=True)
+
+    request_id = str(uuid.uuid4())
+    request_path = requests / f'{request_id}.json'
+    temporary_path = requests / f'.{request_id}.{os.getpid()}.tmp'
+    response_path = responses / f'{request_id}.json'
+    request = {
+        'id': request_id,
+        'targetBytes': args.size,
+        'minInactiveMs': args.min_inactive,
+    }
+
+    temporary_path.write_text(json.dumps(request), encoding='utf-8')
+    os.chmod(temporary_path, 0o600)
+    temporary_path.replace(request_path)
+
+    deadline = time.monotonic() + args.timeout
+    try:
+        while time.monotonic() < deadline:
+            try:
+                response = json.loads(response_path.read_text(encoding='utf-8'))
+                break
+            except FileNotFoundError:
+                time.sleep(0.1)
+        else:
+            raise RuntimeError(
+                'timed out waiting for Firefox; start Firefox from the '
+                'firefox-memory-control package'
+            )
+    finally:
+        request_path.unlink(missing_ok=True)
+
+    response_path.unlink(missing_ok=True)
+    if 'error' in response:
+        raise RuntimeError(response['error'])
+
+    unloaded_tabs = response['unloadedTabs']
+    estimated_bytes = response['estimatedBytes']
+    observed_bytes = response['observedBytes']
+    target_bytes = response['targetBytes']
+    print(
+        f'unloaded {unloaded_tabs} tab(s); '
+        f'Firefox estimated {format_size(estimated_bytes)} reclaimable; '
+        f'MemAvailable increased by {format_size(observed_bytes)}'
+    )
+    if not response['reachedTarget']:
+        print(
+            f'could not reach the requested {format_size(target_bytes)}',
+            file=sys.stderr,
+        )
+        return 2
+    return 0
+
+
+if __name__ == '__main__':
+    try:
+        sys.exit(main())
+    except (OSError, RuntimeError) as error:
+        print(f'firefox-free-memory: {error}', file=sys.stderr)
+        sys.exit(1)