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

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

@@ -6,6 +6,7 @@
   python3,
   nix,
   xz,
+  zstd,
 }:
 
 # This function is for creating a flat-file binary cache, i.e. the kind created by
@@ -16,9 +17,16 @@
 
 {
   name ? "binary-cache",
+  compression ? "zstd", # one of ["none" "xz" "zstd"]
   rootPaths,
 }:
 
+assert lib.elem compression [
+  "none"
+  "xz"
+  "zstd"
+];
+
 stdenv.mkDerivation {
   inherit name;
 
@@ -28,18 +36,20 @@ stdenv.mkDerivation {
 
   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

@@ -1,3 +1,4 @@
+import argparse
 from functools import partial
 import json
 from multiprocessing import Pool
@@ -10,13 +11,15 @@ def dropPrefix(path, nixPrefix):
     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,
@@ -24,30 +27,48 @@ def processItem(item, nixPrefix, outDir):
 
     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"))
@@ -61,7 +82,14 @@ def main():
         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)