From 4d12b83ae0437ad26bd0158178eadfac1c608c2d Mon Sep 17 00:00:00 2001 From: DavHau Date: Thu, 5 Jun 2025 23:53:29 +0700 Subject: [PATCH 1/2] parallelRun, parallelMap: init These bash helpers make it trivial to write shell based setup hooks that utilize all cores. This also makes it simpler to optimize existing hooks which are not yet utilizing all cores. Existing hooks which already use `xargs -P` to parallelize their work, can be optimized further by replacing the xargs call with one of the functions added here , eg. `parallelRun` or `parallelMap`. The new shell based functions `parallelRun` and `parallelMap` are superior to `xargs -P`, because: - They perform better as they launch $NIX_BUILD_CORE workers, each handling many jobs, vs `xargss -P` usually launches a new process for each job (anything else is difficult to implement nicely with xargs). - workers can be defined as shell functions, which allows using all declared shell variables and functions inside the worker (e.g isElf or isScript, etc.), vs. `xargs -P` forces the user to create a new shell process, which cannot re-use declared variables and functions. --- pkgs/build-support/setup-hooks/parallel.sh | 89 +++++++++++ .../setup-hooks/tests/default.nix | 20 +++ .../setup-hooks/tests/test-parallel.sh | 146 ++++++++++++++++++ pkgs/stdenv/generic/default.nix | 1 + pkgs/test/default.nix | 2 + 5 files changed, 258 insertions(+) create mode 100644 pkgs/build-support/setup-hooks/parallel.sh create mode 100644 pkgs/build-support/setup-hooks/tests/default.nix create mode 100644 pkgs/build-support/setup-hooks/tests/test-parallel.sh diff --git a/pkgs/build-support/setup-hooks/parallel.sh b/pkgs/build-support/setup-hooks/parallel.sh new file mode 100644 index 000000000000..a7c8330a0de1 --- /dev/null +++ b/pkgs/build-support/setup-hooks/parallel.sh @@ -0,0 +1,89 @@ +# Parallel execution utilities +# These functions provide a framework for parallel processing of jobs from stdin + +# parallelRun - Execute a command in parallel across multiple cores +# +# Reads null-delimited jobs from stdin and distributes them across NIX_BUILD_CORES +# worker processes. Each worker executes the provided command, receiving jobs +# via stdin in null-delimited format. +# +# Usage: some_producer | parallelRun command [args...] +# +# The command receives jobs one at a time via stdin (null-delimited). +# +# Example: +# find . -name '*.log' -print0 | parallelRun sh -c ' +# while read -r -d "" file; do gzip "$file"; done +# ' +parallelRun() { + local pids + local lock + pids=() + lock=$(mktemp -u) + mkfifo "$lock" + for ((i=0; i"$lock" # fd-4 = write side of lock (push token back) + local job + + while :; do + # Acquire the lock: blocks until a token can be read + read -r -n1 >/dev/null <&3 + + # read one job from stdin + # This is guarded by the lock above in order to prevent + # multiple workers from reading from stdin simultaneously. + if ! IFS= read -r -d '' job; then + # If stdin is closed, release lock and exit + printf 'x' >&4 + break + fi + + # Release the lock: write a token back to the lock FIFO + printf 'y' >&4 + + # Forward job to the worker process' stdin + printf '%s\0' "$job" + + done \ + | "$@" # launch the worker process + } & + pids[$i]=$! + done + # launch the workers by writing a token to the lock FIFO + printf 'a' >"$lock" & + # Wait for all workers to finish + for pid in "${pids[@]}"; do + if ! wait "$pid"; then + echo "A parallel job failed with exit code $? (check for errors above)" >&2 + echo -e "Failing Command:\n $@" >&2 + exit 1 + fi + done + rm "$lock" +} + +# parallelMap - Apply a shell function to each job in parallel +# +# A higher-level wrapper around parallelRun that applies a shell function to each +# null-delimited job from stdin. The shell function receives each job as its first +# argument. +# +# Usage: some_producer | parallelMap shell_function [additional_args...] +# +# The shell function is called as: shell_function job [additional_args...] +# for each job read from stdin. +# +# Example: +# compress() { gzip "$1" } +# find . -name '*.log' -print0 | parallelMap compress +parallelMap() { + _wrapper() { + while IFS= read -r -d '' job; do + "$@" "$job" + done + } + parallelRun _wrapper "$@" + unset -f _wrapper +} diff --git a/pkgs/build-support/setup-hooks/tests/default.nix b/pkgs/build-support/setup-hooks/tests/default.nix new file mode 100644 index 000000000000..ad4e284a9279 --- /dev/null +++ b/pkgs/build-support/setup-hooks/tests/default.nix @@ -0,0 +1,20 @@ +{ + stdenv, +}: +{ + # test based on bootstrap tools to prevent rebuilding stdenv on each change + parallel = derivation { + name = "test-parallel-hook"; + system = stdenv.system; + builder = "${stdenv.bootstrapTools}/bin/bash"; + PATH = "${stdenv.bootstrapTools}/bin"; + args = [ + "-c" + '' + . ${../parallel.sh} + . ${./test-parallel.sh} + '' + ]; + meta = { }; + }; +} diff --git a/pkgs/build-support/setup-hooks/tests/test-parallel.sh b/pkgs/build-support/setup-hooks/tests/test-parallel.sh new file mode 100644 index 000000000000..94e0fc6e8b73 --- /dev/null +++ b/pkgs/build-support/setup-hooks/tests/test-parallel.sh @@ -0,0 +1,146 @@ +export NIX_BUILD_CORES=4 + +echo "Testing worker distribution..." + +# Generate 100 jobs to ensure all workers get some +for i in {1..100}; do + printf "job%d\0" $i +done | parallelRun sh -c ' + while IFS= read -r -d "" job; do + sleep 0.05 # Simulate some work + echo "Worker $$ processed $job" >> /tmp/worker-output + done +' + +# Check that all 4 workers were actually utilized +worker_count=$(sort /tmp/worker-output | cut -d" " -f2 | sort -u | wc -l) +if [ "$worker_count" -ne 4 ]; then + echo "ERROR: Expected exactly 4 workers, got $worker_count" + cat /tmp/worker-output + exit 1 +fi +echo "SUCCESS: All 4 workers participated" +rm -f /tmp/worker-output + +echo "Testing error propagation..." + +# Test that errors from workers are propagated +if printf "job1\0job2\0job3\0" | parallelRun sh -c ' + while IFS= read -r -d "" job; do + if [ "$job" = "job2" ]; then + echo "Worker failing on $job" >&2 + exit 1 + fi + echo "Worker processed $job" + done +' 2>/dev/null; then + echo "ERROR: Expected command to fail but it succeeded" + exit 1 +else + echo "SUCCESS: Error was properly propagated" +fi + +echo "Testing error message..." + +error_output=$(printf "job1\0job2\0job3\0" | parallelRun sh -c ' + while IFS= read -r -d "" job; do + if [ "$job" = "job2" ]; then + echo "Worker failing on $job" >&2 + exit 1 + fi + echo "Worker processed $job" + done +' 2>&1 || true) + +if [[ "$error_output" != *"job failed"* ]]; then + echo "ERROR: Expected 'job failed' in error message, got: $error_output" + exit 1 +fi +echo "SUCCESS: Error message was displayed" + +echo "Testing Verify all jobs are processed when no errors occur..." + +# Generate jobs and count processed ones +for i in {1..10}; do + printf "job%d\0" $i +done | parallelRun sh -c ' + while IFS= read -r -d "" job; do + echo "$job" >> /tmp/processed-jobs + done +' + +processed_count=$(wc -l < /tmp/processed-jobs) +if [ "$processed_count" -ne 10 ]; then + echo "ERROR: Expected 10 jobs processed, got $processed_count" + exit 1 +fi +echo "SUCCESS: All 10 jobs were processed" +rm -f /tmp/processed-jobs + +echo "All parallelRun tests passed!" + +# --------------------------------------------------------------------- + +echo "Testing parallelMap basic functionality..." + +# Define a test function +testFunc() { + echo "Processing: $1" >> /tmp/map-output +} + +# Test that parallelMap calls the function with each job +for i in {1..5}; do + printf "item%d\0" $i +done | parallelMap testFunc + +# Check all jobs were processed +processed_map_count=$(wc -l < /tmp/map-output) +if [ "$processed_map_count" -ne 5 ]; then + echo "ERROR: Expected 5 items processed by parallelMap, got $processed_map_count" + exit 1 +fi +echo "SUCCESS: parallelMap processed all 5 items" +rm -f /tmp/map-output + +echo "Testing parallelMap error propagation..." + +# Define a function that fails on specific input +failFunc() { + if [ "$1" = "item2" ]; then + echo "Function failing on $1" >&2 + exit 1 + fi + echo "Function processed $1" +} + +# Test that errors are propagated +if printf "item1\0item2\0item3\0" | parallelMap failFunc 2>/dev/null; then + echo "ERROR: Expected parallelMap to fail but it succeeded" + exit 1 +else + echo "SUCCESS: parallelMap error was properly propagated" +fi + +echo "Testing parallelMap with additional arguments..." + +# Define a function that uses additional arguments +argFunc() { + echo "$1: $2" >> /tmp/map-args-output +} + +# Test with additional arguments +for i in {1..3}; do + printf "value%d\0" $i +done | parallelMap argFunc "PREFIX" + +# Check output contains the prefix +if ! grep -q "PREFIX: value1" /tmp/map-args-output; then + echo "ERROR: parallelMap did not pass additional arguments correctly" + cat /tmp/map-args-output + exit 1 +fi +echo "SUCCESS: parallelMap passed additional arguments correctly" +rm -f /tmp/map-args-output + +echo "All parallelRun and parallelMap tests passed!" +touch $out diff --git a/pkgs/stdenv/generic/default.nix b/pkgs/stdenv/generic/default.nix index 4d9bb81fbbce..a8aba74321b0 100644 --- a/pkgs/stdenv/generic/default.nix +++ b/pkgs/stdenv/generic/default.nix @@ -78,6 +78,7 @@ let ../../build-support/setup-hooks/move-sbin.sh ../../build-support/setup-hooks/move-systemd-user-units.sh ../../build-support/setup-hooks/multiple-outputs.sh + ../../build-support/setup-hooks/parallel.sh ../../build-support/setup-hooks/patch-shebangs.sh ../../build-support/setup-hooks/prune-libtool-files.sh ../../build-support/setup-hooks/reproducible-builds.sh diff --git a/pkgs/test/default.nix b/pkgs/test/default.nix index 7804ea23b06e..b36a82e44a7a 100644 --- a/pkgs/test/default.nix +++ b/pkgs/test/default.nix @@ -212,4 +212,6 @@ with pkgs; build-environment-info = callPackage ./build-environment-info { }; rust-hooks = recurseIntoAttrs (callPackages ../build-support/rust/hooks/test { }); + + setup-hooks = recurseIntoAttrs (callPackages ../build-support/setup-hooks/tests { }); } From e230668356e9f5fe6932f32e847a2328af326b98 Mon Sep 17 00:00:00 2001 From: DavHau Date: Thu, 5 Jun 2025 23:58:42 +0700 Subject: [PATCH 2/2] audit-tmpdir.sh: optimize - make use of parallelMap --- .../build-support/setup-hooks/audit-tmpdir.sh | 65 ++++++------------- .../setup-hooks/tests/default.nix | 31 +++++---- 2 files changed, 36 insertions(+), 60 deletions(-) diff --git a/pkgs/build-support/setup-hooks/audit-tmpdir.sh b/pkgs/build-support/setup-hooks/audit-tmpdir.sh index 25fe203f813b..780f38bc1624 100644 --- a/pkgs/build-support/setup-hooks/audit-tmpdir.sh +++ b/pkgs/build-support/setup-hooks/audit-tmpdir.sh @@ -15,54 +15,27 @@ auditTmpdir() { echo "checking for references to $TMPDIR/ in $dir..." - local tmpdir elf_fifo script_fifo - tmpdir="$(mktemp -d)" - elf_fifo="$tmpdir/elf" - script_fifo="$tmpdir/script" - mkfifo "$elf_fifo" "$script_fifo" - - # Classifier: identify ELF and script files - ( - find "$dir" -type f -not -path '*/.build-id/*' -print0 \ - | while IFS= read -r -d $'\0' file; do - if isELF "$file"; then - printf '%s\0' "$file" >&3 - elif isScript "$file"; then - filename=${file##*/} - dir=${file%/*} - if [ -e "$dir/.$filename-wrapped" ]; then - printf '%s\0' "$file" >&4 + _processFile() { + local file="$1" + if isELF "$file"; then + if { printf :; patchelf --print-rpath "$file"; } | grep -q -F ":$TMPDIR/"; then + echo "RPATH of binary $file contains a forbidden reference to $TMPDIR/" + exit 1 + fi + elif isScript "$file"; then + filename=${i##*/} + dir=${i%/*} + if [ -e "$dir/.$filename-wrapped" ]; then + if grep -q -F "$TMPDIR/" "$file"; then + echo "wrapper script $file contains a forbidden reference to $TMPDIR/" + exit 1 fi fi - done - exec 3>&- 4>&- - ) 3> "$elf_fifo" 4> "$script_fifo" & + fi + } - # Handler: check RPATHs concurrently - ( - xargs -0 -r -P "$NIX_BUILD_CORES" -n 1 sh -c ' - if { printf :; patchelf --print-rpath "$1"; } | grep -q -F ":$TMPDIR/"; then - echo "RPATH of binary $1 contains a forbidden reference to $TMPDIR/" - exit 1 - fi - ' _ < "$elf_fifo" - ) & - local pid_elf=$! + find "$dir" -type f -not -path '*/.build-id/*' -print0 \ + | parallelMap _processFile - # Handler: check wrapper scripts concurrently - local pid_script - ( - xargs -0 -r -P "$NIX_BUILD_CORES" -n 1 sh -c ' - if grep -q -F "$TMPDIR/" "$1"; then - echo "wrapper script $1 contains a forbidden reference to $TMPDIR/" - exit 1 - fi - ' _ < "$script_fifo" - ) & - local pid_script=$! - - wait "$pid_elf" || { echo "Some binaries contain forbidden references to $TMPDIR/. Check the error above!"; exit 1; } - wait "$pid_script" || { echo "Some scripts contain forbidden references to $TMPDIR/. Check the error above!"; exit 1; } - - rm -r "$tmpdir" + unset -f _processFile } diff --git a/pkgs/build-support/setup-hooks/tests/default.nix b/pkgs/build-support/setup-hooks/tests/default.nix index ad4e284a9279..36cd73a1a9ee 100644 --- a/pkgs/build-support/setup-hooks/tests/default.nix +++ b/pkgs/build-support/setup-hooks/tests/default.nix @@ -3,18 +3,21 @@ }: { # test based on bootstrap tools to prevent rebuilding stdenv on each change - parallel = derivation { - name = "test-parallel-hook"; - system = stdenv.system; - builder = "${stdenv.bootstrapTools}/bin/bash"; - PATH = "${stdenv.bootstrapTools}/bin"; - args = [ - "-c" - '' - . ${../parallel.sh} - . ${./test-parallel.sh} - '' - ]; - meta = { }; - }; + parallel = + (derivation { + name = "test-parallel-hook"; + system = stdenv.system; + builder = "${stdenv.bootstrapTools}/bin/bash"; + PATH = "${stdenv.bootstrapTools}/bin"; + args = [ + "-c" + '' + . ${../parallel.sh} + . ${./test-parallel.sh} + '' + ]; + }) + // { + meta = { }; + }; }