binman: Support multithreading for building images

Some images may take a while to build, e.g. if they are large and use slow
compression. Support compiling sections in parallel to speed things up.

Signed-off-by: Simon Glass <sjg@chromium.org>
(fixed to use a separate test file to fix flakiness)

8 files changed, 136 insertions, 6 deletions

diff --git a/tools/binman/binman.rst b/tools/binman/binman.rst
index bc635aa00a..09e7b57198 100644
--- a/tools/binman/binman.rst
+++ b/tools/binman/binman.rst
@@ -1142,6 +1142,22 @@ adds a -v<level> option to the call to binman::
    make BINMAN_VERBOSE=5
 
 
+Building sections in parallel
+-----------------------------
+
+By default binman uses multiprocessing to speed up compilation of large images.
+This works at a section level, with one thread for each entry in the section.
+This can speed things up if the entries are large and use compression.
+
+This feature can be disabled with the '-T' flag, which defaults to a suitable
+value for your machine. This depends on the Python version, e.g on v3.8 it uses
+12 threads on an 8-core machine. See ConcurrentFutures_ for more details.
+
+The special value -T0 selects single-threaded mode, useful for debugging during
+development, since dealing with exceptions and problems in threads is more
+difficult. This avoids any use of ThreadPoolExecutor.
+
+
 History / Credits
 -----------------
 
@@ -1190,3 +1206,5 @@ Some ideas:
 --
 Simon Glass <sjg@chromium.org>
 7/7/2016
+
+.. _ConcurrentFutures: https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor
diff --git a/tools/binman/cmdline.py b/tools/binman/cmdline.py
index 95f9ba27fb..d6156df408 100644
--- a/tools/binman/cmdline.py
+++ b/tools/binman/cmdline.py
@@ -32,6 +32,10 @@ controlled by a description in the board device tree.'''
         default=False, help='Display the README file')
     parser.add_argument('--toolpath', type=str, action='append',
         help='Add a path to the directories containing tools')
+    parser.add_argument('-T', '--threads', type=int,
+          default=None, help='Number of threads to use (0=single-thread)')
+    parser.add_argument('--test-section-timeout', action='store_true',
+          help='Use a zero timeout for section multi-threading (for testing)')
     parser.add_argument('-v', '--verbosity', default=1,
         type=int, help='Control verbosity: 0=silent, 1=warnings, 2=notices, '
         '3=info, 4=detail, 5=debug')
diff --git a/tools/binman/control.py b/tools/binman/control.py
index f57e34daaa..b2113b6e64 100644
--- a/tools/binman/control.py
+++ b/tools/binman/control.py
@@ -628,9 +628,13 @@ def Binman(args):
             tools.PrepareOutputDir(args.outdir, args.preserve)
             tools.SetToolPaths(args.toolpath)
             state.SetEntryArgs(args.entry_arg)
+            state.SetThreads(args.threads)
 
             images = PrepareImagesAndDtbs(dtb_fname, args.image,
                                           args.update_fdt, use_expanded)
+            if args.test_section_timeout:
+                # Set the first image to timeout, used in testThreadTimeout()
+                images[list(images.keys())[0]].test_section_timeout = True
             missing = False
             for image in images.values():
                 missing |= ProcessImage(image, args.update_fdt, args.map,
diff --git a/tools/binman/etype/section.py b/tools/binman/etype/section.py
index c3bac026c1..92d3f3add4 100644
--- a/tools/binman/etype/section.py
+++ b/tools/binman/etype/section.py
@@ -9,10 +9,12 @@ images to be created.
 """
 
 from collections import OrderedDict
+import concurrent.futures
 import re
 import sys
 
 from binman.entry import Entry
+from binman import state
 from dtoc import fdt_util
 from patman import tools
 from patman import tout
@@ -525,15 +527,43 @@ class Entry_section(Entry):
     def GetEntryContents(self):
         """Call ObtainContents() for each entry in the section
         """
+        def _CheckDone(entry):
+            if not entry.ObtainContents():
+                next_todo.append(entry)
+            return entry
+
         todo = self._entries.values()
         for passnum in range(3):
+            threads = state.GetThreads()
             next_todo = []
-            for entry in todo:
-                if not entry.ObtainContents():
-                    next_todo.append(entry)
+
+            if threads == 0:
+                for entry in todo:
+                    _CheckDone(entry)
+            else:
+                with concurrent.futures.ThreadPoolExecutor(
+                        max_workers=threads) as executor:
+                    future_to_data = {
+                        entry: executor.submit(_CheckDone, entry)
+                        for entry in todo}
+                    timeout = 60
+                    if self.GetImage().test_section_timeout:
+                        timeout = 0
+                    done, not_done = concurrent.futures.wait(
+                        future_to_data.values(), timeout=timeout)
+                    # Make sure we check the result, so any exceptions are
+                    # generated. Check the results in entry order, since tests
+                    # may expect earlier entries to fail first.
+                    for entry in todo:
+                        job = future_to_data[entry]
+                        job.result()
+                    if not_done:
+                        self.Raise('Timed out obtaining contents')
+
             todo = next_todo
             if not todo:
                 break
+
         if todo:
             self.Raise('Internal error: Could not complete processing of contents: remaining %s' %
                        todo)
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index 5383eec489..531ea65771 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -308,7 +308,8 @@ class TestFunctional(unittest.TestCase):
     def _DoTestFile(self, fname, debug=False, map=False, update_dtb=False,
                     entry_args=None, images=None, use_real_dtb=False,
                     use_expanded=False, verbosity=None, allow_missing=False,
-                    extra_indirs=None):
+                    extra_indirs=None, threads=None,
+                    test_section_timeout=False):
         """Run binman with a given test file
 
         Args:
@@ -331,6 +332,8 @@ class TestFunctional(unittest.TestCase):
             allow_missing: Set the '--allow-missing' flag so that missing
                 external binaries just produce a warning instead of an error
             extra_indirs: Extra input directories to add using -I
+            threads: Number of threads to use (None for default, 0 for
+                single-threaded)
         """
         args = []
         if debug:
@@ -342,6 +345,10 @@ class TestFunctional(unittest.TestCase):
         if self.toolpath:
             for path in self.toolpath:
                 args += ['--toolpath', path]
+        if threads is not None:
+            args.append('-T%d' % threads)
+        if test_section_timeout:
+            args.append('--test-section-timeout')
         args += ['build', '-p', '-I', self._indir, '-d', self.TestFile(fname)]
         if map:
             args.append('-m')
@@ -412,7 +419,7 @@ class TestFunctional(unittest.TestCase):
 
     def _DoReadFileDtb(self, fname, use_real_dtb=False, use_expanded=False,
                        map=False, update_dtb=False, entry_args=None,
-                       reset_dtbs=True, extra_indirs=None):
+                       reset_dtbs=True, extra_indirs=None, threads=None):
         """Run binman and return the resulting image
 
         This runs binman with a given test file and then reads the resulting
@@ -439,6 +446,8 @@ class TestFunctional(unittest.TestCase):
                 function. If reset_dtbs is True, then the original test dtb
                 is written back before this function finishes
             extra_indirs: Extra input directories to add using -I
+            threads: Number of threads to use (None for default, 0 for
+                single-threaded)
 
         Returns:
             Tuple:
@@ -463,7 +472,8 @@ class TestFunctional(unittest.TestCase):
         try:
             retcode = self._DoTestFile(fname, map=map, update_dtb=update_dtb,
                     entry_args=entry_args, use_real_dtb=use_real_dtb,
-                    use_expanded=use_expanded, extra_indirs=extra_indirs)
+                    use_expanded=use_expanded, extra_indirs=extra_indirs,
+                    threads=threads)
             self.assertEqual(0, retcode)
             out_dtb_fname = tools.GetOutputFilename('u-boot.dtb.out')
 
@@ -4542,5 +4552,22 @@ class TestFunctional(unittest.TestCase):
         data = self._DoReadFile('201_opensbi.dts')
         self.assertEqual(OPENSBI_DATA, data[:len(OPENSBI_DATA)])
 
+    def testSectionsSingleThread(self):
+        """Test sections without multithreading"""
+        data = self._DoReadFileDtb('055_sections.dts', threads=0)[0]
+        expected = (U_BOOT_DATA + tools.GetBytes(ord('!'), 12) +
+                    U_BOOT_DATA + tools.GetBytes(ord('a'), 12) +
+                    U_BOOT_DATA + tools.GetBytes(ord('&'), 4))
+        self.assertEqual(expected, data)
+
+    def testThreadTimeout(self):
+        """Test handling a thread that takes too long"""
+        with self.assertRaises(ValueError) as e:
+            self._DoTestFile('202_section_timeout.dts',
+                             test_section_timeout=True)
+        self.assertIn("Node '/binman/section@0': Timed out obtaining contents",
+                      str(e.exception))
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tools/binman/image.py b/tools/binman/image.py
index 10778f47fe..cdc58b39a4 100644
--- a/tools/binman/image.py
+++ b/tools/binman/image.py
@@ -36,6 +36,8 @@ class Image(section.Entry_section):
         fdtmap_data: Contents of the fdtmap when loading from a file
         allow_repack: True to add properties to allow the image to be safely
             repacked later
+        test_section_timeout: Use a zero timeout for section multi-threading
+            (for testing)
 
     Args:
         copy_to_orig: Copy offset/size to orig_offset/orig_size after reading
@@ -74,6 +76,7 @@ class Image(section.Entry_section):
         self.allow_repack = False
         self._ignore_missing = ignore_missing
         self.use_expanded = use_expanded
+        self.test_section_timeout = False
         if not test:
             self.ReadNode()
 
diff --git a/tools/binman/state.py b/tools/binman/state.py
index dfb1760455..2f56758738 100644
--- a/tools/binman/state.py
+++ b/tools/binman/state.py
@@ -7,6 +7,7 @@
 
 import hashlib
 import re
+import threading
 
 from dtoc import fdt
 import os
@@ -55,6 +56,9 @@ allow_entry_expansion = True
 # to the new ones, the compressed size increases, etc.
 allow_entry_contraction = False
 
+# Number of threads to use for binman (None means machine-dependent)
+num_threads = None
+
 def GetFdtForEtype(etype):
     """Get the Fdt object for a particular device-tree entry
 
@@ -420,3 +424,22 @@ def AllowEntryContraction():
             raised
     """
     return allow_entry_contraction
+
+def SetThreads(threads):
+    """Set the number of threads to use when building sections
+
+    Args:
+        threads: Number of threads to use (None for default, 0 for
+            single-threaded)
+    """
+    global num_threads
+
+    num_threads = threads
+
+def GetThreads():
+    """Get the number of threads to use when building sections
+
+    Returns:
+        Number of threads to use (None for default, 0 for single-threaded)
+    """
+    return num_threads
diff --git a/tools/binman/test/202_section_timeout.dts b/tools/binman/test/202_section_timeout.dts
new file mode 100644
index 0000000000..1481450367
--- /dev/null
+++ b/tools/binman/test/202_section_timeout.dts
@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: GPL-2.0+
+
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		pad-byte = <0x26>;
+		size = <0x28>;
+		section@0 {
+			read-only;
+			size = <0x10>;
+			pad-byte = <0x21>;
+
+			u-boot {
+			};
+		};
+	};
+};