From 5439b4cf29fb45979f721212965752a994d14929 Mon Sep 17 00:00:00 2001
From: Aditya Deshpande <aditya.deshpande@arm.com>
Date: Thu, 13 Apr 2023 16:32:21 +0100
Subject: [PATCH] Clean up the worktree and exit  gracefully if any pre-build
 or make commands fail. Without this, the script will fail to create a
 worktree when re-run due to the existence of the worktree from the previous
 run.

Signed-off-by: Aditya Deshpande <aditya.deshpande@arm.com>
---
 scripts/code_size_compare.py | 73 +++++++++++++++++++++++++-----------
 1 file changed, 52 insertions(+), 21 deletions(-)

diff --git a/scripts/code_size_compare.py b/scripts/code_size_compare.py
index e8ade7578346..66733a2a9e2e 100755
--- a/scripts/code_size_compare.py
+++ b/scripts/code_size_compare.py
@@ -163,11 +163,15 @@ def _create_git_worktree(self, revision):
         else:
             print("Creating git worktree for", revision)
             git_worktree_path = os.path.join(self.repo_path, "temp-" + revision)
-            subprocess.check_output(
-                [self.git_command, "worktree", "add", "--detach",
-                 git_worktree_path, revision], cwd=self.repo_path,
-                stderr=subprocess.STDOUT
-            )
+            try:
+                subprocess.check_output(
+                    [self.git_command, "worktree", "add", "--detach",
+                    git_worktree_path, revision], cwd=self.repo_path,
+                    stderr=subprocess.STDOUT
+                )
+            except subprocess.CalledProcessError as e:
+                self._handle_CalledProcessError(e,git_worktree_path)
+
         return git_worktree_path
 
     def _build_libraries(self, git_worktree_path):
@@ -175,15 +179,20 @@ def _build_libraries(self, git_worktree_path):
 
         my_environment = os.environ.copy()
         if self.pre_build_commands != '':
+            try:
+                subprocess.check_output(
+                    self.pre_build_commands, env=my_environment, shell=True,
+                    cwd=git_worktree_path, stderr=subprocess.STDOUT,
+                )
+            except subprocess.CalledProcessError as e:
+                self._handle_CalledProcessError(e,git_worktree_path)
+        try:
             subprocess.check_output(
-                self.pre_build_commands, env=my_environment, shell=True,
+                self.make_command, env=my_environment, shell=True,
                 cwd=git_worktree_path, stderr=subprocess.STDOUT,
-        )
-
-        subprocess.check_output(
-            self.make_command, env=my_environment, shell=True,
-            cwd=git_worktree_path, stderr=subprocess.STDOUT,
-        )
+            )
+        except subprocess.CalledProcessError as e:
+            self._handle_CalledProcessError(e,git_worktree_path)
 
     def _gen_code_size_report(self, revision, git_worktree_path):
         """Generate a code size report for each executable and store them
@@ -194,19 +203,28 @@ def _gen_code_size_report(self, revision, git_worktree_path):
             print("Measuring code size for", revision)
 
         # Size for libmbedcrypto.a
-        result = subprocess.check_output(
-            ["size -t library/libmbedcrypto.a"], cwd=git_worktree_path, shell=True
-        )
+        try:
+            result = subprocess.check_output(
+                ["size -t library/libmbedcrypto.a"], cwd=git_worktree_path, shell=True
+            )
+        except subprocess.CalledProcessError as e:
+            self._handle_CalledProcessError(e,git_worktree_path)
         crypto_text = result.decode()
         # Size for libmbedx509.a
-        result = subprocess.check_output(
-            ["size -t library/libmbedx509.a"], cwd=git_worktree_path, shell=True
-        )
+        try:
+            result = subprocess.check_output(
+                ["size -t library/libmbedx509.a"], cwd=git_worktree_path, shell=True
+            )
+        except subprocess.CalledProcessError as e:
+            self._handle_CalledProcessError(e,git_worktree_path)
         x509_text = result.decode()
         # Size for libmbedtls.a
-        result = subprocess.check_output(
-            ["size -t library/libmbedtls.a"], cwd=git_worktree_path, shell=True
-        )
+        try:
+            result = subprocess.check_output(
+                ["size -t library/libmbedtls.a"], cwd=git_worktree_path, shell=True
+            )
+        except subprocess.CalledProcessError as e:
+            self._handle_CalledProcessError(e,git_worktree_path)
         tls_text = result.decode()
 
         def size_text_to_dict(txt):
@@ -320,6 +338,19 @@ def get_comparision_results(self):
         self._get_code_size_for_rev(self.new_rev)
         return self.compare_code_size()
 
+    def _handle_CalledProcessError(self, e: subprocess.CalledProcessError, git_worktree_path):
+        """Handle a CalledProcessError and quit the program gracefully. Remove any
+        extra worktrees so that the script may be called again."""
+
+        # Tell the user what went wrong
+        print("The following command: {} failed and exited with code {}"\
+            .format(e.cmd, e.returncode))
+        print("Process output:\n {}".format(e.output))
+
+        # Quit gracefully by removing the existing worktree
+        self._remove_worktree(git_worktree_path)
+        sys.exit(-1)
+
 def main():
     parser = argparse.ArgumentParser(
         description=(