LLVM tools: Added support for 'skip' and executing a custom script

If a tryjob's status is set to 'skip', then that tryjob will not be
considered when bisecting LLVM. And for executing a custom script,
the exit code will be used by a mapping that sets the 'status' to
the equivalent value of the exit code.

BUG=None
TEST=Ran the 'update_tryjob_status.py' with the 'skip' option and
successfully set the tryjob's status to 'skip'. Also, created a testing
script to serve as the custom script. Successfully set the 'status'
value based off of the custom script's exit code.

Change-Id: I0ae3a9b53467b13ff5ce26aeff46b31563cfc7e0
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/1772472
Reviewed-by: George Burgess <gbiv@chromium.org>
Reviewed-by: Manoj Gupta <manojgupta@chromium.org>
Tested-by: Salud Lemus <saludlemus@google.com>

3 files changed, 148 insertions, 45 deletions

diff --git a/llvm_tools/get_llvm_hash.py b/llvm_tools/get_llvm_hash.py
index 20e00eb5..3a0df8fe 100755
--- a/llvm_tools/get_llvm_hash.py
+++ b/llvm_tools/get_llvm_hash.py
@@ -18,6 +18,19 @@ import subprocess
 import tempfile
 
 
+def CheckCommand(cmd):
+  """Executes the command using Popen()."""
+
+  cmd_obj = subprocess.Popen(
+      cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+
+  stdout, _ = cmd_obj.communicate()
+
+  if cmd_obj.returncode:
+    print(stdout)
+    raise subprocess.CalledProcessError(cmd_obj.returncode, cmd)
+
+
 @contextmanager
 def CreateTempLLVMRepo(temp_dir):
   """Adds a LLVM worktree to 'temp_dir'.
@@ -47,13 +60,7 @@ def CreateTempLLVMRepo(temp_dir):
       temp_dir, 'master'
   ]
 
-  add_worktree_cmd_obj = subprocess.Popen(
-      add_worktree_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-  _, stderr = add_worktree_cmd_obj.communicate()
-
-  if add_worktree_cmd_obj.returncode:
-    raise ValueError('Failed to add worktree for %s: %s' %
-                     (abs_path_to_llvm_project_dir, stderr))
+  CheckCommand(add_worktree_cmd)
 
   try:
     yield temp_dir
@@ -106,23 +113,11 @@ def GetAndUpdateLLVMProjectInLLVMTools():
         'git', '-C', abs_path_to_llvm_project_dir, 'checkout', 'master'
     ]
 
-    checkout_cmd_obj = subprocess.Popen(
-        checkout_to_master_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    _, stderr = checkout_cmd_obj.communicate()
-
-    if checkout_cmd_obj.returncode:
-      raise ValueError('Failed to checkout to master for %s: %s' %
-                       (abs_path_to_llvm_project_dir, stderr))
+    CheckCommand(checkout_to_master_cmd)
 
     update_master_cmd = ['git', '-C', abs_path_to_llvm_project_dir, 'pull']
 
-    update_cmd_obj = subprocess.Popen(
-        update_master_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    _, stderr = update_cmd_obj.communicate()
-
-    if update_cmd_obj.returncode:
-      raise ValueError('Failed to fetch from chromium mirror of LLVM for %s: %s'
-                       % (abs_path_to_llvm_project_dir, stderr))
+    CheckCommand(update_master_cmd)
 
   return abs_path_to_llvm_project_dir
 
diff --git a/llvm_tools/llvm_bisection.py b/llvm_tools/llvm_bisection.py
index b4705d49..361b9183 100755
--- a/llvm_tools/llvm_bisection.py
+++ b/llvm_tools/llvm_bisection.py
@@ -154,7 +154,8 @@ def GetStartAndEndRevision(start, end, tryjobs):
     ]
 
   Returns:
-    The new start version and end version for bisection.
+    The new start version and end version for bisection, a set of revisions
+    that are 'pending' and a set of revisions that are to be skipped.
 
   Raises:
     ValueError: The value for 'status' is missing or there is a mismatch
@@ -198,20 +199,32 @@ def GetStartAndEndRevision(start, end, tryjobs):
   # Find all revisions that are 'pending' within 'good_rev' and 'bad_rev'.
   #
   # NOTE: The intent is to not launch tryjobs between 'good_rev' and 'bad_rev'
-  # that have already been launched (this list is used when constructing the
+  # that have already been launched (this set is used when constructing the
   # list of revisions to launch tryjobs for).
-  pending_revisions = [
+  pending_revisions = {
       tryjob['rev']
       for tryjob in tryjobs
       if tryjob['status'] == TryjobStatus.PENDING.value and
       good_rev < tryjob['rev'] < bad_rev
-  ]
+  }
 
-  return good_rev, bad_rev, pending_revisions
+  # Find all revisions that are to be skipped within 'good_rev' and 'bad_rev'.
+  #
+  # NOTE: The intent is to not launch tryjobs between 'good_rev' and 'bad_rev'
+  # that have already been marked as 'skip' (this set is used when constructing
+  # the list of revisions to launch tryjobs for).
+  skip_revisions = {
+      tryjob['rev']
+      for tryjob in tryjobs
+      if tryjob['status'] == TryjobStatus.SKIP.value and
+      good_rev < tryjob['rev'] < bad_rev
+  }
+
+  return good_rev, bad_rev, pending_revisions, skip_revisions
 
 
 def GetRevisionsBetweenBisection(start, end, parallel, src_path,
-                                 pending_revisions):
+                                 pending_revisions, skip_revisions):
   """Gets the revisions between 'start' and 'end'.
 
   Sometimes, the LLVM source tree's revisions do not increment by 1 (there is
@@ -224,8 +237,10 @@ def GetRevisionsBetweenBisection(start, end, parallel, src_path,
     end: The end revision.
     parallel: The number of tryjobs to create between 'start' and 'end'.
     src_path: The absolute path to the LLVM source tree to use.
-    pending_revisions: A list of 'pending' revisions that are between 'start'
-    and 'end'.
+    pending_revisions: A set containing 'pending' revisions that are between
+    'start' and 'end'.
+    skip_revisions: A set containing revisions between 'start' and 'end' that
+    are to be skipped.
 
   Returns:
     A list of revisions between 'start' and 'end'.
@@ -242,7 +257,8 @@ def GetRevisionsBetweenBisection(start, end, parallel, src_path,
   # this.
   for cur_revision in range(start + 1, end):
     try:
-      if cur_revision not in pending_revisions:
+      if cur_revision not in pending_revisions and \
+          cur_revision not in skip_revisions:
         # Verify that the current revision exists by finding its corresponding
         # git hash in the LLVM source tree.
         new_llvm.GetGitHashForVersion(src_path, cur_revision)
@@ -258,16 +274,14 @@ def GetRevisionsBetweenBisection(start, end, parallel, src_path,
   if not index_step:
     index_step = 1
 
-  # Starting at 'index_step' because the first element would be close to
-  # 'start' (similar to ('parallel' + 1) for the last element).
   result = [valid_revisions[index] \
-            for index in range(index_step, len(valid_revisions), index_step)]
+            for index in range(0, len(valid_revisions), index_step)]
 
   return result
 
 
 def GetRevisionsListAndHashList(start, end, parallel, src_path,
-                                pending_revisions):
+                                pending_revisions, skip_revisions):
   """Determines the revisions between start and end."""
 
   new_llvm = LLVMHash()
@@ -278,16 +292,13 @@ def GetRevisionsListAndHashList(start, end, parallel, src_path,
         src_path = new_repo
 
       # Get a list of revisions between start and end.
-      revisions = GetRevisionsBetweenBisection(start, end, parallel, src_path,
-                                               pending_revisions)
+      revisions = GetRevisionsBetweenBisection(
+          start, end, parallel, src_path, pending_revisions, skip_revisions)
 
       git_hashes = [
           new_llvm.GetGitHashForVersion(src_path, rev) for rev in revisions
       ]
 
-  assert revisions, ('No revisions between start %d and end %d to create '
-                     'tryjobs.' % (start, end))
-
   return revisions, git_hashes
 
 
@@ -323,13 +334,25 @@ def main():
   _ValidateStartAndEndAgainstJSONStartAndEnd(
       start, end, bisect_contents['start'], bisect_contents['end'])
 
-  # Pending revisions are between 'start_revision' and 'end_revision'.
-  start_revision, end_revision, pending_revisions = GetStartAndEndRevision(
-      start, end, bisect_contents['jobs'])
+  # Pending and skipped revisions are between 'start_revision' and
+  # 'end_revision'.
+  start_revision, end_revision, pending_revisions, skip_revisions = \
+      GetStartAndEndRevision(start, end, bisect_contents['jobs'])
 
   revisions, git_hashes = GetRevisionsListAndHashList(
       start_revision, end_revision, args_output.parallel, args_output.src_path,
-      pending_revisions)
+      pending_revisions, skip_revisions)
+
+  if not revisions:
+    no_revisions_message = (
+        'No revisions between start %d and end '
+        '%d to create tryjobs' % (start_revision, end_revision))
+
+    if skip_revisions:
+      no_revisions_message += '\nThe following tryjobs were skipped:\n' \
+          + '\n'.join(str(rev) for rev in skip_revisions)
+
+    raise ValueError(no_revisions_message)
 
   # Check if any revisions that are going to be added as a tryjob exist already
   # in the 'jobs' list.
diff --git a/llvm_tools/update_tryjob_status.py b/llvm_tools/update_tryjob_status.py
index 8ef68bfb..8425f2c6 100755
--- a/llvm_tools/update_tryjob_status.py
+++ b/llvm_tools/update_tryjob_status.py
@@ -12,11 +12,13 @@ import argparse
 import enum
 import json
 import os
+import subprocess
 import sys
 
 from assert_not_in_chroot import VerifyOutsideChroot
 from cros_utils import command_executer
 from patch_manager import _ConvertToASCII
+from test_helpers import CreateTemporaryJsonFile
 
 
 class TryjobStatus(enum.Enum):
@@ -25,6 +27,11 @@ class TryjobStatus(enum.Enum):
   GOOD = 'good'
   BAD = 'bad'
   PENDING = 'pending'
+  SKIP = 'skip'
+
+  # Executes the script passed into the command line (this script's exit code
+  # determines the 'status' value of the tryjob).
+  CUSTOM_SCRIPT = 'custom_script'
 
   # Uses the result returned by 'cros buildresult'.
   AUTO = 'auto'
@@ -38,6 +45,28 @@ class BuilderStatus(enum.Enum):
   RUNNING = 'running'
 
 
+class CustomScriptStatus(enum.Enum):
+  """Exit code values of a custom script."""
+
+  # NOTE: Not using 1 for 'bad' because the custom script can raise an
+  # exception which would cause the exit code of the script to be 1, so the
+  # tryjob's 'status' would be updated when there is an exception.
+  #
+  # Exit codes are as follows:
+  #   0: 'good'
+  #   124: 'bad'
+  #   125: 'skip'
+  GOOD = 0
+  BAD = 124
+  SKIP = 125
+
+
+custom_script_exit_value_mapping = {
+    CustomScriptStatus.GOOD.value: TryjobStatus.GOOD.value,
+    CustomScriptStatus.BAD.value: TryjobStatus.BAD.value,
+    CustomScriptStatus.SKIP.value: TryjobStatus.SKIP.value
+}
+
 builder_status_mapping = {
     BuilderStatus.PASS.value: TryjobStatus.GOOD.value,
     BuilderStatus.FAIL.value: TryjobStatus.BAD.value,
@@ -84,6 +113,15 @@ def GetCommandLineArgs():
       default=cros_root,
       help='the path to the chroot (default: %(default)s)')
 
+  # Add argument for the custom script to execute for the 'custom_script'
+  # option in '--set_status'.
+  parser.add_argument(
+      '--custom_script',
+      help='The absolute path to the custom script to execute (its exit code '
+      'should be %d for \'good\', %d for \'bad\', or %d for \'skip\')' %
+      (CustomScriptStatus.GOOD.value, CustomScriptStatus.BAD.value,
+       CustomScriptStatus.SKIP.value))
+
   args_output = parser.parse_args()
 
   if not os.path.isfile(args_output.status_file) or \
@@ -91,6 +129,11 @@ def GetCommandLineArgs():
     raise ValueError('File does not exist or does not ending in \'.json\' '
                      ': %s' % args_output.status_file)
 
+  if args_output.set_status == TryjobStatus.CUSTOM_SCRIPT.value and \
+      not args_output.custom_script:
+    raise ValueError('Please provide the absolute path to the script to '
+                     'execute.')
+
   return args_output
 
 
@@ -145,7 +188,8 @@ def GetStatusFromCrosBuildResult(chroot_path, buildbucket_id):
   return str(tryjob_contents['%d' % buildbucket_id]['status'])
 
 
-def UpdateTryjobStatus(revision, set_status, status_file, chroot_path):
+def UpdateTryjobStatus(revision, set_status, status_file, chroot_path,
+                       custom_script):
   """Updates a tryjob's 'status' field based off of 'set_status'.
 
   Args:
@@ -156,6 +200,8 @@ def UpdateTryjobStatus(revision, set_status, status_file, chroot_path):
       'cros buildresult'.
     status_file: The .JSON file that contains the tryjobs.
     chroot_path: The absolute path to the chroot (used by 'cros buildresult').
+    custom_script: The absolute path to a script that will be executed which
+    will determine the 'status' value of the tryjob.
   """
 
   # Format of 'bisect_contents':
@@ -202,6 +248,44 @@ def UpdateTryjobStatus(revision, set_status, status_file, chroot_path):
 
     bisect_contents['jobs'][tryjob_index]['status'] = builder_status_mapping[
         build_result]
+  elif set_status == TryjobStatus.SKIP:
+    bisect_contents['jobs'][tryjob_index]['status'] = TryjobStatus.SKIP.value
+  elif set_status == TryjobStatus.CUSTOM_SCRIPT:
+    # Create a temporary file to write the contents of the tryjob at index
+    # 'tryjob_index' (the temporary file path will be passed into the custom
+    # script as a command line argument).
+    with CreateTemporaryJsonFile() as temp_json_file:
+      with open(temp_json_file, 'w') as tryjob_file:
+        json.dump(
+            bisect_contents['jobs'][tryjob_index],
+            tryjob_file,
+            indent=4,
+            separators=(',', ': '))
+
+      exec_script_cmd = [custom_script, temp_json_file]
+
+      # Execute the custom script to get the exit code.
+      exec_script_cmd_obj = subprocess.Popen(
+          exec_script_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+      _, stderr = exec_script_cmd_obj.communicate()
+
+      # Invalid exit code by the custom script.
+      if exec_script_cmd_obj.returncode not in custom_script_exit_value_mapping:
+        name_of_json_file = os.path.join(
+            os.path.dirname(status_file), os.path.basename(temp_json_file))
+        os.rename(temp_json_file, name_of_json_file)
+        raise ValueError(
+            'Custom script %s exit code %d did not match '
+            'any of the expected exit codes: %d for \'good\', %d '
+            'for \'bad\', or %d for \'skip\'\nPlease check %s for information '
+            'about the tryjob: %s' %
+            (custom_script, exec_script_cmd_obj.returncode,
+             CustomScriptStatus.GOOD.value, CustomScriptStatus.BAD.value,
+             CustomScriptStatus.SKIP.value, name_of_json_file, stderr))
+
+    bisect_contents['jobs'][tryjob_index][
+        'status'] = custom_script_exit_value_mapping[exec_script_cmd_obj
+                                                     .returncode]
   else:
     raise ValueError('Invalid \'set_status\' option provided: %s' % set_status)
 
@@ -217,7 +301,8 @@ def main():
   args_output = GetCommandLineArgs()
 
   UpdateTryjobStatus(args_output.revision, TryjobStatus(args_output.set_status),
-                     args_output.status_file, args_output.chroot_path)
+                     args_output.status_file, args_output.chroot_path,
+                     args_output.custom_script)
 
 
 if __name__ == '__main__':