update.nix: Run update scripts in parallel

To make updating large attribute sets faster, the update scripts
are now run in parallel.

Please note the following changes in semantics:

- The string passed to updateScript needs to be a path to an executable file.
- The updateScript can also be a list: the tail elements will then be passed
  to the head as command line arguments.

2 files changed, 98 insertions, 34 deletions

diff --git a/maintainers/scripts/update.nix b/maintainers/scripts/update.nix
index 8d1e47c6bc9..120cd5552f4 100755
--- a/maintainers/scripts/update.nix
+++ b/maintainers/scripts/update.nix
@@ -1,6 +1,8 @@
 { package ? null
 , maintainer ? null
 , path ? null
+, max-workers ? null
+, keep-going ? null
 }:
 
 # TODO: add assert statements
@@ -105,27 +107,24 @@ let
         % nix-shell maintainers/scripts/update.nix --argstr path gnome3
 
     to run update script for all package under an attribute path.
-  '';
 
-  runUpdateScript = package: ''
-    echo -ne " - ${package.name}: UPDATING ..."\\r
-    ${package.updateScript} &> ${(builtins.parseDrvName package.name).name}.log
-    CODE=$?
-    if [ "$CODE" != "0" ]; then
-      echo " - ${package.name}: ERROR       "
-      echo ""
-      echo "--- SHOWING ERROR LOG FOR ${package.name} ----------------------"
-      echo ""
-      cat ${(builtins.parseDrvName package.name).name}.log
-      echo ""
-      echo "--- SHOWING ERROR LOG FOR ${package.name} ----------------------"
-      exit $CODE
-    else
-      rm ${(builtins.parseDrvName package.name).name}.log
-    fi
-    echo " - ${package.name}: DONE.       "
+    You can also add
+
+        --argstr max-workers 8
+
+    to increase the number of jobs in parallel, or
+
+        --argstr keep-going true
+
+    to continue running when a single update fails.
   '';
 
+  packageData = package: {
+    name = package.name;
+    pname = (builtins.parseDrvName package.name).name;
+    updateScript = pkgs.lib.toList package.updateScript;
+  };
+
 in pkgs.stdenv.mkDerivation {
   name = "nixpkgs-update-script";
   buildCommand = ''
@@ -139,21 +138,7 @@ in pkgs.stdenv.mkDerivation {
     exit 1
   '';
   shellHook = ''
-    echo ""
-    echo "Going to be running update for following packages:"
-    echo "${builtins.concatStringsSep "\n" (map (x: " - ${x.name}") packages)}"
-    echo ""
-    read -n1 -r -p "Press space to continue..." confirm
-    if [ "$confirm" = "" ]; then
-      echo ""
-      echo "Running update for:"
-      ${builtins.concatStringsSep "\n" (map runUpdateScript packages)}
-      echo ""
-      echo "Packages updated!"
-      exit 0
-    else
-      echo "Aborting!"
-      exit 1
-    fi
+    unset shellHook # do not contaminate nested shells
+    exec ${pkgs.python3.interpreter} ${./update.py} ${pkgs.writeText "packages.json" (builtins.toJSON (map packageData packages))}${pkgs.lib.optionalString (max-workers != null) " --max-workers=${max-workers}"}${pkgs.lib.optionalString (keep-going == "true") " --keep-going"}
   '';
 }
diff --git a/maintainers/scripts/update.py b/maintainers/scripts/update.py
new file mode 100644
index 00000000000..eb7d0ef2647
--- /dev/null
+++ b/maintainers/scripts/update.py
@@ -0,0 +1,79 @@
+import argparse
+import concurrent.futures
+import json
+import os
+import subprocess
+import sys
+
+updates = {}
+
+def eprint(*args, **kwargs):
+    print(*args, file=sys.stderr, **kwargs)
+
+def run_update_script(package):
+    eprint(f" - {package['name']}: UPDATING ...")
+
+    subprocess.run(package['updateScript'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=True)
+
+
+def main(max_workers, keep_going, packages):
+    with open(sys.argv[1]) as f:
+        packages = json.load(f)
+
+    eprint()
+    eprint('Going to be running update for following packages:')
+    for package in packages:
+        eprint(f" - {package['name']}")
+    eprint()
+
+    confirm = input('Press Enter key to continue...')
+    if confirm == '':
+        eprint()
+        eprint('Running update for:')
+
+        with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+            for package in packages:
+                updates[executor.submit(run_update_script, package)] = package
+
+            for future in concurrent.futures.as_completed(updates):
+                package = updates[future]
+
+                try:
+                    future.result()
+                    eprint(f" - {package['name']}: DONE.")
+                except subprocess.CalledProcessError as e:
+                    eprint(f" - {package['name']}: ERROR")
+                    eprint()
+                    eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
+                    eprint()
+                    eprint(e.stdout.decode('utf-8'))
+                    with open(f"{package['pname']}.log", 'wb') as f:
+                        f.write(e.stdout)
+                    eprint()
+                    eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
+
+                    if not keep_going:
+                        sys.exit(1)
+
+        eprint()
+        eprint('Packages updated!')
+        sys.exit()
+    else:
+        eprint('Aborting!')
+        sys.exit(130)
+
+parser = argparse.ArgumentParser(description='Update packages')
+parser.add_argument('--max-workers', '-j', dest='max_workers', type=int, help='Number of updates to run concurrently', nargs='?', default=4)
+parser.add_argument('--keep-going', '-k', dest='keep_going', action='store_true', help='Do not stop after first failure')
+parser.add_argument('packages', help='JSON file containing the list of package names and their update scripts')
+
+if __name__ == '__main__':
+    args = parser.parse_args()
+
+    try:
+        main(args.max_workers, args.keep_going, args.packages)
+    except (KeyboardInterrupt, SystemExit) as e:
+        for update in updates:
+            update.cancel()
+
+        sys.exit(e.code if isinstance(e, SystemExit) else 130)