From 11798449b229633fc75d5d5e8536f3e818fef792 Mon Sep 17 00:00:00 2001
From: Andrei Litvin <andy314@gmail.com>
Date: Tue, 29 Nov 2022 06:45:39 -0500
Subject: [PATCH] Run zap in parallel by default for zap_regen_all.py (#23766)

* Run zap in parallel on regen all. This makes my local zap run in 88 seconds instead of about 222

* Restyle

* Fix arguments for generate.py
---
 scripts/tools/zap/generate.py  |  9 +++++++++
 scripts/tools/zap_regen_all.py | 26 ++++++++++++++++++++++++--
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/scripts/tools/zap/generate.py b/scripts/tools/zap/generate.py
index 5c48a41619d1aa..9a10752da064e1 100755
--- a/scripts/tools/zap/generate.py
+++ b/scripts/tools/zap/generate.py
@@ -32,6 +32,7 @@ class CmdLineArgs:
     templateFile: str
     outputDir: str
     runBootstrap: bool
+    parallel: bool = True
 
 
 CHIP_ROOT_DIR = os.path.realpath(
@@ -103,6 +104,9 @@ def runArgumentsParser() -> CmdLineArgs:
                         help='Output directory for the generated files (default: automatically selected)')
     parser.add_argument('--run-bootstrap', default=None, action='store_true',
                         help='Automatically run ZAP bootstrap. By default the bootstrap is not triggered')
+    parser.add_argument('--parallel', action='store_true')
+    parser.add_argument('--no-parallel', action='store_false', dest='parallel')
+    parser.set_defaults(parallel=True)
     args = parser.parse_args()
 
     # By default, this script assumes that the global CHIP template is used with
@@ -259,6 +263,11 @@ def main():
 
     # The maximum memory usage is over 4GB (#15620)
     os.environ["NODE_OPTIONS"] = "--max-old-space-size=8192"
+
+    if cmdLineArgs.parallel:
+        # Parallel-compatible runs will need separate state
+        os.environ["ZAP_TEMPSTATE"] = "1"
+
     runGeneration(cmdLineArgs.zapFile, cmdLineArgs.zclFile, cmdLineArgs.templateFile, cmdLineArgs.outputDir)
 
     prettifiers = [
diff --git a/scripts/tools/zap_regen_all.py b/scripts/tools/zap_regen_all.py
index 18de729f16b3a8..e9483084c56019 100755
--- a/scripts/tools/zap_regen_all.py
+++ b/scripts/tools/zap_regen_all.py
@@ -21,6 +21,8 @@
 import sys
 import subprocess
 import logging
+import multiprocessing
+
 from dataclasses import dataclass
 
 CHIP_ROOT_DIR = os.path.realpath(
@@ -113,6 +115,11 @@ def setupArgumentsParser():
                         help="Don't do any generation, just log what targets would be generated (default: False)")
     parser.add_argument('--run-bootstrap', default=None, action='store_true',
                         help='Automatically run ZAP bootstrap. By default the bootstrap is not triggered')
+
+    parser.add_argument('--parallel', action='store_true')
+    parser.add_argument('--no-parallel', action='store_false', dest='parallel')
+    parser.set_defaults(parallel=True)
+
     return parser.parse_args()
 
 
@@ -253,6 +260,14 @@ def getTargets(type, test_target):
     return targets
 
 
+def _ParallelGenerateOne(target):
+    """
+    Helper method to be passed to multiprocessing parallel generation of
+    items.
+    """
+    target.generate()
+
+
 def main():
     logging.basicConfig(
         level=logging.INFO,
@@ -269,8 +284,15 @@ def main():
         if args.run_bootstrap:
             subprocess.check_call(os.path.join(CHIP_ROOT_DIR, "scripts/tools/zap/zap_bootstrap.sh"), shell=True)
 
-        for target in targets:
-            target.generate()
+        if args.parallel:
+            # Ensure each zap run is independent
+            os.environ['ZAP_TEMPSTATE'] = '1'
+            with multiprocessing.Pool() as pool:
+                for _ in pool.imap_unordered(_ParallelGenerateOne, targets):
+                    pass
+        else:
+            for target in targets:
+                target.generate()
 
 
 if __name__ == '__main__':