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/default.nix

@@ -11,6 +11,7 @@ in
   chocolate-doom2xx = callPackage ./chocolate-doom2xx { };
   windowtolayer = callPackage ./windowtolayer.nix { };
   swaylock-plugin = callPackage ./swaylock-plugin.nix { };
+  firefox-memory-control = callPackage ./firefox-memory-control { };
 
   update-docs-assignments = pkgs.writeShellScriptBin "update-docs-assignments" ''
     exec ${pkgs.python3}/bin/python3 ${../ci/update-docs-assignments.py} "$@"

pkgs/firefox-memory-control/autoconfig.js

@@ -0,0 +1,251 @@
+// This file runs as privileged Firefox AutoConfig code. Keep it in the Nix store.
+(() => {
+  "use strict";
+
+  const { classes: Cc, interfaces: Ci, utils: Cu } = Components;
+  Cu.importGlobalProperties(["IOUtils"]);
+  const Services = {
+    appinfo: Cc["@mozilla.org/xre/app-info;1"].getService(Ci.nsIXULRuntime),
+    console: Cc["@mozilla.org/consoleservice;1"].getService(Ci.nsIConsoleService),
+    env: Cc["@mozilla.org/process/environment;1"].getService(Ci.nsIEnvironment),
+    prefs: Cc["@mozilla.org/preferences-service;1"].getService(Ci.nsIPrefBranch),
+  };
+  const { TabUnloader } = ChromeUtils.importESModule(
+    "moz-src:///browser/components/tabbrowser/TabUnloader.sys.mjs"
+  );
+
+  const PREFIX = "firefox.memoryControl.";
+  const MiB = 1024 * 1024;
+  const log = message => {
+    const line = `[firefox-memory-control] ${message}`;
+    Services.console.logStringMessage(line);
+    if (typeof dump === "function") {
+      dump(`${line}\n`);
+    }
+  };
+
+  const sleep = milliseconds =>
+    new Promise(resolve => {
+      const timer = Cc["@mozilla.org/timer;1"].createInstance(Ci.nsITimer);
+      timer.initWithCallback(resolve, milliseconds, Ci.nsITimer.TYPE_ONE_SHOT);
+    });
+
+  const prefInt = name => Services.prefs.getIntPref(PREFIX + name);
+  const prefBool = name => Services.prefs.getBoolPref(PREFIX + name);
+
+  // procfs files report a size of zero, so IOUtils reads /proc/meminfo as empty.
+  // nsIScriptableInputStream also rejects reads larger than that reported size;
+  // nsIConverterInputStream reads until EOF without relying on it.
+  function availableMemory() {
+    const file = Cc["@mozilla.org/file/local;1"].createInstance(Ci.nsIFile);
+    file.initWithPath("/proc/meminfo");
+    const fileStream = Cc["@mozilla.org/network/file-input-stream;1"].createInstance(
+      Ci.nsIFileInputStream
+    );
+    fileStream.init(file, 0x01, 0, 0);
+    const input = Cc["@mozilla.org/intl/converter-input-stream;1"].createInstance(
+      Ci.nsIConverterInputStream
+    );
+    input.init(fileStream, "UTF-8", 0, 0);
+    const chunk = {};
+    let meminfo = "";
+    while (input.readString(4096, chunk)) {
+      meminfo += chunk.value;
+    }
+    input.close();
+
+    const match = /^MemAvailable:\s+(\d+)\s+kB$/m.exec(meminfo);
+    if (!match) {
+      throw new Error("MemAvailable is absent from /proc/meminfo");
+    }
+    return Number(match[1]) * 1024;
+  }
+
+  async function unloadOne(minInactiveMs) {
+    const sorted = await TabUnloader.getSortedTabs(minInactiveMs);
+    const candidate = sorted.find(tab => TabUnloader.isDiscardable(tab));
+    if (!candidate) {
+      return null;
+    }
+
+    const estimatedBytes = candidate.memory || 0;
+    const unloaded = await TabUnloader.unloadLeastRecentlyUsedTab(minInactiveMs);
+    return unloaded ? { estimatedBytes } : null;
+  }
+
+  function runtimePaths() {
+    const runtimeDir = Services.env.get("XDG_RUNTIME_DIR");
+    if (!runtimeDir || !runtimeDir.startsWith("/")) {
+      throw new Error("XDG_RUNTIME_DIR is not an absolute path");
+    }
+
+    const root = `${runtimeDir}/firefox-memory-control`;
+    return {
+      root,
+      requests: `${root}/requests`,
+      processing: `${root}/processing`,
+      responses: `${root}/responses`,
+    };
+  }
+
+  async function ensureRuntimeDirectories(paths) {
+    for (const path of Object.values(paths)) {
+      await IOUtils.makeDirectory(path, { ignoreExisting: true, permissions: 0o700 });
+    }
+  }
+
+  async function claimRequest(paths) {
+    const children = await IOUtils.getChildren(paths.requests);
+    for (const requestPath of children.sort()) {
+      if (!requestPath.endsWith(".json")) {
+        continue;
+      }
+
+      const leaf = requestPath.slice(requestPath.lastIndexOf("/") + 1);
+      const claimed = `${paths.processing}/${leaf}.${Services.appinfo.processID}`;
+      try {
+        await IOUtils.move(requestPath, claimed, { noOverwrite: true });
+        return claimed;
+      } catch (error) {
+        // Another Firefox instance can win the atomic move.
+      }
+    }
+    return null;
+  }
+
+  async function writeResponse(paths, id, response) {
+    const finalPath = `${paths.responses}/${id}.json`;
+    const temporaryPath = `${finalPath}.${Services.appinfo.processID}.tmp`;
+    await IOUtils.writeUTF8(temporaryPath, JSON.stringify(response));
+    await IOUtils.move(temporaryPath, finalPath, { noOverwrite: true });
+  }
+
+  async function serviceRequest(paths, requestPath) {
+    let request;
+    try {
+      request = JSON.parse(await IOUtils.readUTF8(requestPath));
+      if (!/^[0-9a-f-]{36}$/.test(request.id)) {
+        throw new Error("invalid request id");
+      }
+      if (!Number.isSafeInteger(request.targetBytes) || request.targetBytes <= 0) {
+        throw new Error("targetBytes must be a positive integer");
+      }
+
+      const minInactiveMs = Number.isSafeInteger(request.minInactiveMs)
+        ? Math.max(0, request.minInactiveMs)
+        : 0;
+      const baseline = await availableMemory();
+      let estimatedBytes = 0;
+      let observedBytes = 0;
+      let unloadedTabs = 0;
+
+      while (
+        estimatedBytes < request.targetBytes &&
+        observedBytes < request.targetBytes &&
+        unloadedTabs < 100
+      ) {
+        const result = await unloadOne(minInactiveMs);
+        if (!result) {
+          break;
+        }
+
+        unloadedTabs += 1;
+        estimatedBytes += result.estimatedBytes;
+        await sleep(400);
+        observedBytes = Math.max(0, (await availableMemory()) - baseline);
+      }
+
+      const reachedTarget =
+        estimatedBytes >= request.targetBytes || observedBytes >= request.targetBytes;
+      await writeResponse(paths, request.id, {
+        id: request.id,
+        reachedTarget,
+        targetBytes: request.targetBytes,
+        unloadedTabs,
+        estimatedBytes,
+        observedBytes,
+      });
+    } catch (error) {
+      log(`request failed: ${error}`);
+      if (request && /^[0-9a-f-]{36}$/.test(request.id)) {
+        await writeResponse(paths, request.id, {
+          id: request.id,
+          reachedTarget: false,
+          error: String(error),
+        });
+      }
+    } finally {
+      await IOUtils.remove(requestPath, { ignoreAbsent: true });
+    }
+  }
+
+  const controller = {
+    busy: false,
+    underPressure: false,
+    timer: null,
+    paths: null,
+
+    async tick() {
+      if (this.busy) {
+        return;
+      }
+      this.busy = true;
+
+      try {
+        const request = await claimRequest(this.paths);
+        if (request) {
+          await serviceRequest(this.paths, request);
+          return;
+        }
+
+        if (!prefBool("enabled")) {
+          this.underPressure = false;
+          return;
+        }
+
+        const available = await availableMemory();
+        const low = prefInt("lowAvailableMiB") * MiB;
+        const high = prefInt("highAvailableMiB") * MiB;
+        if (high <= low) {
+          throw new Error("highAvailableMiB must be greater than lowAvailableMiB");
+        }
+
+        if (!this.underPressure && available <= low) {
+          this.underPressure = true;
+          log(`memory pressure entered at ${Math.round(available / MiB)} MiB available`);
+        } else if (this.underPressure && available >= high) {
+          this.underPressure = false;
+          log(`memory pressure cleared at ${Math.round(available / MiB)} MiB available`);
+        }
+
+        if (this.underPressure) {
+          await unloadOne(prefInt("minInactiveMs"));
+        }
+      } catch (error) {
+        log(`poll failed: ${error}`);
+      } finally {
+        this.busy = false;
+      }
+    },
+
+    async start() {
+      try {
+        this.paths = runtimePaths();
+        await ensureRuntimeDirectories(this.paths);
+        const interval = Math.max(250, prefInt("pollIntervalMs"));
+        this.timer = Cc["@mozilla.org/timer;1"].createInstance(Ci.nsITimer);
+        this.timer.initWithCallback(
+          () => this.tick(),
+          interval,
+          Ci.nsITimer.TYPE_REPEATING_SLACK
+        );
+        log(`started; polling every ${interval} ms`);
+        await this.tick();
+      } catch (error) {
+        log(`startup failed: ${error}`);
+      }
+    },
+  };
+
+  controller.start();
+})();

pkgs/firefox-memory-control/default.nix

@@ -0,0 +1,46 @@
+{
+  lib,
+  firefox-unwrapped,
+  wrapFirefox,
+  writeText,
+  writeScriptBin,
+  symlinkJoin,
+  python3,
+  lowAvailableMiB ? 2048,
+  highAvailableMiB ? 3072,
+  pollIntervalMs ? 1000,
+  minInactiveMs ? 300000,
+}:
+let
+  autoConfig = writeText "firefox-memory-control.js" ''
+    defaultPref("firefox.memoryControl.enabled", true);
+    defaultPref("firefox.memoryControl.lowAvailableMiB", ${toString lowAvailableMiB});
+    defaultPref("firefox.memoryControl.highAvailableMiB", ${toString highAvailableMiB});
+    defaultPref("firefox.memoryControl.pollIntervalMs", ${toString pollIntervalMs});
+    defaultPref("firefox.memoryControl.minInactiveMs", ${toString minInactiveMs});
+
+    ${builtins.readFile ./autoconfig.js}
+  '';
+
+  firefox = wrapFirefox firefox-unwrapped {
+    extraAutoConfig = ''
+      pref("general.config.sandbox_enabled", false);
+    '';
+    extraPrefsFiles = [ autoConfig ];
+  };
+
+  freeMemory = writeScriptBin "firefox-free-memory" ''
+    #!${python3}/bin/python3
+    ${builtins.readFile ./firefox-free-memory.py}
+  '';
+in
+symlinkJoin {
+  name = "firefox-memory-control-${firefox.version}";
+  paths = [ firefox freeMemory ];
+
+  meta = firefox.meta // {
+    description = "Firefox with memory-pressure tab unloading and an on-demand reclaim utility";
+    mainProgram = "firefox";
+    platforms = lib.platforms.linux;
+  };
+}

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)