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

This commit is contained in:
thomasjm 2025-01-05 00:24:37 -08:00 committed by Tom McLaughlin
parent 2f5bd177a3
commit 00a218abb2
6 changed files with 82 additions and 24 deletions

View file

@ -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)