build-unit-test-docker: build docker stages in parallel

Signed-off-by: Patrick Williams <patrick@stwcx.xyz>
Change-Id: I1899b2cacdf8f668c3f1226959d1b5a044b5fc6c
diff --git a/scripts/build-unit-test-docker b/scripts/build-unit-test-docker
index 392865f..8f6c73f 100755
--- a/scripts/build-unit-test-docker
+++ b/scripts/build-unit-test-docker
@@ -19,6 +19,7 @@
 
 import os
 import sys
+import threading
 from datetime import date
 from hashlib import sha256
 from sh import docker, git, nproc, uname
@@ -45,7 +46,10 @@
 elif arch == "x86_64":
     docker_base = ""
 else:
-    print(f"Unsupported system architecture({arch}) found for docker image")
+    print(
+        f"Unsupported system architecture({arch}) found for docker image",
+        file=sys.stderr,
+    )
     sys.exit(1)
 
 # Packages to include in image.
@@ -372,30 +376,58 @@
     return result
 
 
+pkg_lock = threading.Lock()
+
+
 def pkg_generate(pkg):
-    if "__pkg_built" in packages[pkg]:
-        return None
+    class pkg_thread(threading.Thread):
+        def run(self):
+            pkg_lock.acquire()
+            deps = [
+                packages[deppkg]["__thread"]
+                for deppkg in sorted(packages[pkg].get("depends", []))
+            ]
+            pkg_lock.release()
+            for deppkg in deps:
+                deppkg.join()
 
-    for deppkg in sorted(packages[pkg].get("depends", [])):
-        pkg_generate(deppkg)
-
-    dockerfile = f"""
+            dockerfile = f"""
 FROM {docker_base_img_name}
 {pkg_copycmds(pkg)}
 {pkg_build(pkg)}
 """
 
-    tag = docker_img_tagname(pkg_stagename(pkg), dockerfile)
-    packages[pkg]["__tag"] = tag
-    docker_img_build(tag, dockerfile)
+            pkg_lock.acquire()
+            tag = docker_img_tagname(pkg_stagename(pkg), dockerfile)
+            packages[pkg]["__tag"] = tag
+            pkg_lock.release()
 
-    packages[pkg]["__pkg_built"] = True
+            try:
+                self.exception = None
+                docker_img_build(pkg, tag, dockerfile)
+            except Exception as e:
+                self.package = pkg
+                self.exception = e
+
+    packages[pkg]["__thread"] = pkg_thread()
 
 
 def pkg_generate_packages():
-    for pkg in sorted(packages.keys()):
+    for pkg in packages.keys():
         pkg_generate(pkg)
 
+    pkg_lock.acquire()
+    pkg_threads = [packages[p]["__thread"] for p in packages.keys()]
+    for t in pkg_threads:
+        t.start()
+    pkg_lock.release()
+
+    for t in pkg_threads:
+        t.join()
+        if t.exception:
+            print(f"Package {t.package} failed!", file=sys.stderr)
+            raise t.exception
+
 
 def docker_img_tagname(pkgname, dockerfile):
     result = docker_image_name
@@ -406,8 +438,8 @@
     return result
 
 
-def docker_img_build(tag, dockerfile):
-    for line in docker.build(
+def docker_img_build(pkg, tag, dockerfile):
+    docker.build(
         proxy_args,
         "--network=host",
         "--force-rm",
@@ -415,9 +447,10 @@
         tag,
         "-",
         _in=dockerfile,
-        _iter=True,
-    ):
-        print(line, end="", flush=True)
+        _out=(
+            lambda line: print(pkg + ":", line, end="", file=sys.stderr, flush=True)
+        ),
+    )
 
 
 # Look up the HEAD for missing a static rev.
@@ -601,7 +634,7 @@
 
 # Build the stage docker images.
 docker_base_img_name = docker_img_tagname("base", dockerfile_base)
-docker_img_build(docker_base_img_name, dockerfile_base)
+docker_img_build("base", docker_base_img_name, dockerfile_base)
 pkg_generate_packages()
 
 dockerfile = f"""
@@ -627,4 +660,4 @@
 """
 
 # Do the final docker build
-docker_img_build(docker_image_name, dockerfile)
+docker_img_build("final", docker_image_name, dockerfile)