mkBinaryCache: support different compression methods: xz (default), zstd, none

doc/build-helpers/images/binarycache.section.md

···
 `rootPaths` must be a list of derivations.
 The transitive closure of these derivations' outputs will be copied into the cache.
 
+## Optional arguments {#sec-pkgs-binary-cache-arguments}
+
+`compression` (`"none"` or `"xz"` or `"zstd"`; _optional_)
+
+: The compression algorithm to use.
+
+  _Default value:_ `zstd`.
+
 ::: {.note}
 This function is meant for advanced use cases.
 The more idiomatic way to work with flat-file binary caches is via the [nix-copy-closure](https://nixos.org/manual/nix/stable/command-ref/nix-copy-closure.html) command.

doc/redirects.json

···
   "sec-pkgs-binary-cache": [
     "index.html#sec-pkgs-binary-cache"
   ],
+  "sec-pkgs-binary-cache-arguments": [
+    "index.html#sec-pkgs-binary-cache-arguments"
+  ],
   "sec-pkgs-binary-cache-example": [
     "index.html#sec-pkgs-binary-cache-example"
   ],

nixos/tests/all-tests.nix

···
   beanstalkd = handleTest ./beanstalkd.nix {};
   bees = handleTest ./bees.nix {};
   benchexec = handleTest ./benchexec.nix {};
-  binary-cache = handleTest ./binary-cache.nix {};
+  binary-cache = handleTest ./binary-cache.nix { compression = "zstd"; };
+  binary-cache-no-compression = handleTest ./binary-cache.nix { compression = "none"; };
+  binary-cache-xz = handleTest ./binary-cache.nix { compression = "xz"; };
   bind = handleTest ./bind.nix {};
   bird = handleTest ./bird.nix {};
   birdwatcher = handleTest ./birdwatcher.nix {};

nixos/tests/binary-cache.nix

···
+{ compression, ... }@args:
+
 import ./make-test-python.nix (
   { lib, pkgs, ... }:
 
   {
-    name = "binary-cache";
+    name = "binary-cache-" + compression;
     meta.maintainers = with lib.maintainers; [ thomasjm ];
 
     nodes.machine =
···
               nativeBuildInputs = [ openssl ];
             }
             ''
-              tar -czf tmp.tar.gz -C "${mkBinaryCache { rootPaths = [ hello ]; }}" .
+              tar -czf tmp.tar.gz -C "${
+                mkBinaryCache {
+                  rootPaths = [ hello ];
+                  inherit compression;
+                }
+              }" .
               openssl enc -aes-256-cbc -salt -in tmp.tar.gz -out $out -k mysecretpassword
             '';
 
···
       machine.succeed("[ -d %s ] || exit 1" % storePath)
     '';
   }
-)
+) args

pkgs/build-support/binary-cache/default.nix

···
   python3,
   nix,
   xz,
+  zstd,
 }:
 
 # This function is for creating a flat-file binary cache, i.e. the kind created by
···
 
 {
   name ? "binary-cache",
+  compression ? "zstd", # one of ["none" "xz" "zstd"]
   rootPaths,
 }:
 
+assert lib.elem compression [
+  "none"
+  "xz"
+  "zstd"
+];
+
 stdenv.mkDerivation {
   inherit name;
 
···
 
   preferLocalBuild = true;
 
-  nativeBuildInputs = [
-    coreutils
-    jq
-    python3
-    nix
-    xz
-  ];
+  nativeBuildInputs =
+    [
+      coreutils
+      jq
+      python3
+      nix
+    ]
+    ++ lib.optional (compression == "xz") xz
+    ++ lib.optional (compression == "zstd") zstd;
 
   buildCommand = ''
     mkdir -p $out/nar
 
-    python ${./make-binary-cache.py}
+    python ${./make-binary-cache.py} --compression "${compression}"
 
     # These directories must exist, or Nix might try to create them in LocalBinaryCacheStore::init(),
     # which fails if mounted read-only

pkgs/build-support/binary-cache/make-binary-cache.py

···
+import argparse
 from functools import partial
 import json
 from multiprocessing import Pool
···
     return path[len(nixPrefix + "/") :]
 
 
-def processItem(item, nixPrefix, outDir):
+def processItem(
+    item, nixPrefix, outDir, compression, compressionCommand, compressionExtension
+):
     narInfoHash = dropPrefix(item["path"], nixPrefix).split("-")[0]
 
-    xzFile = outDir / "nar" / f"{narInfoHash}.nar.xz"
-    with open(xzFile, "wb") as f:
+    narFile = outDir / "nar" / f"{narInfoHash}{compressionExtension}"
+    with open(narFile, "wb") as f:
         subprocess.run(
-            f"nix-store --dump {item['path']} | xz -c",
+            f"nix-store --dump {item['path']} {compressionCommand}",
             stdout=f,
             shell=True,
             check=True,
···
 
     fileHash = (
         subprocess.run(
-            ["nix-hash", "--base32", "--type", "sha256", "--flat", xzFile],
+            ["nix-hash", "--base32", "--type", "sha256", "--flat", narFile],
             capture_output=True,
             check=True,
         )
         .stdout.decode()
         .strip()
     )
-    fileSize = os.path.getsize(xzFile)
+    fileSize = os.path.getsize(narFile)
 
-    finalXzFileName = Path("nar") / f"{fileHash}.nar.xz"
-    os.rename(xzFile, outDir / finalXzFileName)
+    finalNarFileName = Path("nar") / f"{fileHash}{compressionExtension}"
+    os.rename(narFile, outDir / finalNarFileName)
 
     with open(outDir / f"{narInfoHash}.narinfo", "wt") as f:
         f.write(f"StorePath: {item['path']}\n")
-        f.write(f"URL: {finalXzFileName}\n")
-        f.write("Compression: xz\n")
+        f.write(f"URL: {finalNarFileName}\n")
+        f.write(f"Compression: {compression}\n")
         f.write(f"FileHash: sha256:{fileHash}\n")
         f.write(f"FileSize: {fileSize}\n")
         f.write(f"NarHash: {item['narHash']}\n")
         f.write(f"NarSize: {item['narSize']}\n")
-        f.write(f"References: {' '.join(dropPrefix(ref, nixPrefix) for ref in item['references'])}\n")
+        f.write(
+            f"References: {' '.join(dropPrefix(ref, nixPrefix) for ref in item['references'])}\n"
+        )
 
 
 def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--compression", choices=["none", "xz", "zstd"])
+    args = parser.parse_args()
+
+    compressionCommand = {
+        "none": "",
+        "xz": "| xz -c",
+        "zstd": "| zstd",
+    }[args.compression]
+
+    compressionExtension = {
+        "none": "",
+        "xz": ".xz",
+        "zstd": ".zst",
+    }[args.compression]
+
     outDir = Path(os.environ["out"])
     nixPrefix = os.environ["NIX_STORE"]
     numWorkers = int(os.environ.get("NIX_BUILD_CORES", "4"))
···
         f.write(f"StoreDir: {nixPrefix}\n")
 
     with Pool(processes=numWorkers) as pool:
-        worker = partial(processItem, nixPrefix=nixPrefix, outDir=outDir)
+        worker = partial(
+            processItem,
+            nixPrefix=nixPrefix,
+            outDir=outDir,
+            compression=args.compression,
+            compressionCommand=compressionCommand,
+            compressionExtension=compressionExtension,
+        )
         pool.map(worker, closures)