AFDO-Bisect: Write test confirming assumptions re state saving

This CL adds a test which confirms that the state saving used correctly
produces the same exact profiles as if the run went straight through.

BUG=None
TEST=All tests, new and old, pass.

Change-Id: I250bbf051165deeb673dedc61b3c6e6904cf4875
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/1713235
Reviewed-by: George Burgess <gbiv@chromium.org>
Tested-by: Emma Vukelj <emmavukelj@google.com>

4 files changed, 181 insertions, 6 deletions

diff --git a/afdo_tools/bisection/afdo_prof_analysis.py b/afdo_tools/bisection/afdo_prof_analysis.py
index d974a7a8..f50e0b88 100755
--- a/afdo_tools/bisection/afdo_prof_analysis.py
+++ b/afdo_tools/bisection/afdo_prof_analysis.py
@@ -120,18 +120,19 @@ class DeciderState(object):
   def save_state(self):
     state = {'seed': self.seed, 'accumulated_results': self.accumulated_results}
     fd, tmp_file = mkstemp()
+    os.close(fd)
     with open(tmp_file, 'w') as f:
       json.dump(state, f, indent=2)
-    os.close(fd)
     os.rename(tmp_file, self.state_file)
     logging.info('Logged state to %s...', self.state_file)
 
   def run(self, prof, save_run=True):
     """Run the external deciding script on the given profile."""
     if self.saved_results and save_run:
-      result = StatusEnum(self.saved_results.pop(0))
+      result = self.saved_results.pop(0)
       self.accumulated_results.append(result)
-      return result
+      self.save_state()
+      return StatusEnum(result)
 
     filename = prof_to_tmp(prof)
 
@@ -147,7 +148,6 @@ class DeciderState(object):
         raise RuntimeError('Provided decider script returned PROBLEM_STATUS '
                            'when run on profile stored at %s. AFDO Profile '
                            'analysis aborting' % prof_file)
-
       if save_run:
         self.accumulated_results.append(status.value)
       logging.info('Run %d of external script %s returned %s',
diff --git a/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py b/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py
index f95d46e7..24f9e4d0 100755
--- a/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py
+++ b/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py
@@ -126,6 +126,91 @@ class AfdoProfAnalysisE2ETest(unittest.TestCase):
           state_file=state_file,
           extern_decider='problemstatus_external.sh')
 
+  def test_state_assumption(self):
+
+    def compare_runs(tmp_dir, first_ctr, second_ctr):
+      """Compares given prof versions between first and second run in test."""
+      first_prof = '%s/.first_run_%d' % (tmp_dir, first_ctr)
+      second_prof = '%s/.second_run_%d' % (tmp_dir, second_ctr)
+      with open(first_prof) as f:
+        first_prof_text = f.read()
+      with open(second_prof) as f:
+        second_prof_text = f.read()
+      self.assertEqual(first_prof_text, second_prof_text)
+
+    good_prof = {'func_a': ':1\n3: 3\n5: 7\n'}
+    bad_prof = {'func_a': ':2\n4: 4\n6: 8\n'}
+    # add some noise to the profiles; 15 is an arbitrary choice
+    for x in range(15):
+      func = 'func_%d' % x
+      good_prof[func] = ':%d\n' % (x)
+      bad_prof[func] = ':%d\n' % (x + 1)
+    expected = {
+        'bisect_results': {
+            'ranges': [],
+            'individuals': ['func_a']
+        },
+        'good_only_functions': False,
+        'bad_only_functions': False
+    }
+
+    # using a static temp dir rather than a dynamic one because these files are
+    # shared between the bash scripts and this Python test, and the arguments
+    # to the bash scripts are fixed by afdo_prof_analysis.py so it would be
+    # difficult to communicate dynamically generated directory to bash scripts
+    scripts_tmp_dir = '%s/afdo_test_tmp' % os.getcwd()
+    os.mkdir(scripts_tmp_dir)
+    self.addCleanup(shutil.rmtree, scripts_tmp_dir, ignore_errors=True)
+
+    # files used in the bash scripts used as external deciders below
+    # - count_file tracks the current number of calls to the script in total
+    # - local_count_file tracks the number of calls to the script without
+    # interruption
+    count_file = '%s/.count' % scripts_tmp_dir
+    local_count_file = '%s/.local_count' % scripts_tmp_dir
+
+    # runs through whole thing at once
+    initial_seed = self.run_check(
+        good_prof,
+        bad_prof,
+        expected,
+        extern_decider='state_assumption_external.sh')
+    with open(count_file) as f:
+      num_calls = int(f.read())
+    os.remove(count_file)  # reset counts for second run
+    finished_state_file = 'afdo_analysis_state.json.completed.%s' % str(
+        date.today())
+    self.addCleanup(os.remove, finished_state_file)
+
+    # runs the same analysis but interrupted each iteration
+    for i in range(2 * num_calls + 1):
+      no_resume_run = (i == 0)
+      seed = initial_seed if no_resume_run else None
+      try:
+        self.run_check(
+            good_prof,
+            bad_prof,
+            expected,
+            no_resume=no_resume_run,
+            extern_decider='state_assumption_interrupt.sh',
+            seed=seed)
+        break
+      except RuntimeError:
+        # script was interrupted, so we restart local count
+        os.remove(local_count_file)
+    else:
+      raise RuntimeError('Test failed -- took too many iterations')
+
+    for initial_ctr in range(3):  # initial runs unaffected by interruption
+      compare_runs(scripts_tmp_dir, initial_ctr, initial_ctr)
+
+    start = 3
+    for ctr in range(start, num_calls):
+      # second run counter incremented by 4 for each one first run is because
+      # +2 for performing initial checks on good and bad profs each time
+      # +1 for PROBLEM_STATUS run which causes error and restart
+      compare_runs(scripts_tmp_dir, ctr, 6 + (ctr - start) * 4)
+
   def run_check(self,
                 good_prof,
                 bad_prof,
@@ -133,7 +218,9 @@ class AfdoProfAnalysisE2ETest(unittest.TestCase):
                 state_file=None,
                 no_resume=True,
                 out_file=None,
-                extern_decider=None):
+                extern_decider=None,
+                seed=None):
+
     temp_dir = tempfile.mkdtemp()
     self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
 
@@ -149,7 +236,16 @@ class AfdoProfAnalysisE2ETest(unittest.TestCase):
     analysis.FLAGS.good_prof = good_prof_file
     analysis.FLAGS.bad_prof = bad_prof_file
     if state_file:
+      actual_state_file = analysis.FLAGS.state_file
+
+      def cleanup():
+        analysis.FLAGS.state_file = actual_state_file
+
+      self.addCleanup(cleanup)
+
       analysis.FLAGS.state_file = state_file
+
+    analysis.FLAGS.seed = seed
     analysis.FLAGS.no_resume = no_resume
     analysis.FLAGS.analysis_output_file = out_file or '/dev/null'
 
@@ -158,8 +254,9 @@ class AfdoProfAnalysisE2ETest(unittest.TestCase):
     analysis.FLAGS.external_decider = external_script
 
     actual = analysis.main(None)
-    actual.pop('seed')  # nothing to check
+    actual_seed = actual.pop('seed')  # nothing to check
     self.assertEqual(actual, expected)
+    return actual_seed
 
 
 if __name__ == '__main__':
diff --git a/afdo_tools/bisection/state_assumption_external.sh b/afdo_tools/bisection/state_assumption_external.sh
new file mode 100755
index 00000000..1ad78ee2
--- /dev/null
+++ b/afdo_tools/bisection/state_assumption_external.sh
@@ -0,0 +1,40 @@
+#!/bin/bash -eu
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+# This script returns BAD_STATUS if '2' is in the top line of 'func_a's profile
+# and good otherwise
+
+GOOD_STATUS=0
+BAD_STATUS=1
+SKIP_STATUS=125
+PROBLEM_STATUS=127
+
+tmp_dir=$(pwd)/afdo_test_tmp
+count_file=${tmp_dir}/.count
+
+# keep count for purpose of filenames
+if [ -f "${count_file}" ]; then
+  num_call=$(cat "${count_file}")
+else
+  num_call=0
+fi
+
+echo -n $(( ${num_call}+1 )) > "${count_file}"
+
+tmp_file=$(mktemp)
+trap "rm -f '${tmp_file}'" EXIT
+grep -v '^ ' "$1" > "${tmp_file}"
+
+# copy prof to specific file for later test
+if [[ $# -eq 2 ]]; then
+  cp "$1" "${tmp_dir}/.second_run_${num_call}"
+else
+  cp "$1" "${tmp_dir}/.first_run_${num_call}"
+fi
+
+if grep -q 'func_a.*2' "${tmp_file}"; then
+  exit "${BAD_STATUS}"
+fi
+exit "${GOOD_STATUS}"
diff --git a/afdo_tools/bisection/state_assumption_interrupt.sh b/afdo_tools/bisection/state_assumption_interrupt.sh
new file mode 100755
index 00000000..eba3a4b4
--- /dev/null
+++ b/afdo_tools/bisection/state_assumption_interrupt.sh
@@ -0,0 +1,38 @@
+#!/bin/bash -eu
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+# This script returns the result of state_assumption_external.sh on every even
+# iteration, and PROBLEM_STATUS on every odd_iteration
+
+PROBLEM_STATUS=127
+
+tmp_dir=$(pwd)/afdo_test_tmp
+
+count_file="${tmp_dir}/.count"
+if [[ -f "${count_file}" ]]; then
+  num_call=$(cat "${count_file}")
+else
+  num_call=0
+fi
+
+local_count_file=${tmp_dir}/.local_count
+if [[ -f "${local_count_file}" ]]; then
+  local_count=$(cat "${local_count_file}")
+else
+  local_count=0
+fi
+
+echo -n $(( ${local_count}+1 )) > "${local_count_file}"
+
+# Don't want to fail on performance checks hence local_count >= 2
+# but following that, want to fail every other check
+if [[ ${local_count} -ge 2 ]] && [[ $(( ${num_call}%2 )) -ne 0 ]]; then
+  echo -n $(( ${num_call}+1 )) > "${count_file}"
+  exit "${PROBLEM_STATUS}"
+fi
+
+# script just needs any second argument to write profs to .second_run_*
+$(pwd)/state_assumption_external.sh "$1" 'second_run'
+exit $?