Reduce cros_sdk calls in crosperf.

Improves crosperf performance by reducing calls to ChrootRunCommand,
Instead use RunCommand.

BUG: chromium:564889
TEST:crosperf unit tests
Change-Id: I0705273d1126b81167d7a789514771122a71b8d3
Reviewed-on: https://chromium-review.googlesource.com/424088
Reviewed-by: Caroline Tice <cmtice@chromium.org>
Commit-Queue: Manoj Gupta <manojgupta@chromium.org>
Tested-by: Manoj Gupta <manojgupta@chromium.org>

2 files changed, 114 insertions, 67 deletions

diff --git a/crosperf/download_images.py b/crosperf/download_images.py
index a2a7800f..8ceaa874 100644
--- a/crosperf/download_images.py
+++ b/crosperf/download_images.py
@@ -12,6 +12,8 @@ import test_flag
 
 from cros_utils import command_executer
 
+GS_UTIL = 'chromium/tools/depot_tools/gsutil.py'
+
 
 class MissingImage(Exception):
   """Raised when the requested image does not exist in gs://"""
@@ -21,6 +23,25 @@ class MissingFile(Exception):
   """Raised when the requested file does not exist in gs://"""
 
 
+class RunCommandExceptionHandler(object):
+  """Handle Exceptions from calls to RunCommand"""
+
+  def __init__(self, logger_to_use, log_level, cmd_exec, command):
+    self.logger = logger_to_use
+    self.log_level = log_level
+    self.ce = cmd_exec
+    self.cleanup_command = command
+
+  def HandleException(self, _, e):
+    # Exception handler, Run specified command
+    if self.log_level != 'verbose' and self.cleanup_command is not None:
+      self.logger.LogOutput('CMD: %s' % self.cleanup_command)
+    if self.cleanup_command is not None:
+      _ = self.ce.RunCommand(self.cleanup_command)
+    # Raise exception again
+    raise e
+
+
 class ImageDownloader(object):
   """Download images from Cloud Storage."""
 
@@ -59,11 +80,12 @@ class ImageDownloader(object):
     # Check to see if the image has already been downloaded.  If not,
     # download the image.
     if not os.path.exists(image_path):
-      command = 'gsutil cp %s /tmp/%s' % (image_name, build_id)
+      gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
+      command = '%s cp %s %s' % (gsutil_cmd, image_name, download_path)
 
       if self.log_level != 'verbose':
         self._logger.LogOutput('CMD: %s' % command)
-      status = self._ce.ChrootRunCommand(chromeos_root, command)
+      status = self._ce.RunCommand(command)
       downloaded_image_name = os.path.join(download_path,
                                            'chromiumos_test_image.tar.xz')
       if status != 0 or not os.path.exists(downloaded_image_name):
@@ -79,30 +101,34 @@ class ImageDownloader(object):
       return
 
     # Uncompress and untar the downloaded image.
-    command = ('cd /tmp/%s ; tar -Jxf chromiumos_test_image.tar.xz ' % build_id)
+    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
+    command = ('cd %s ; tar -Jxf chromiumos_test_image.tar.xz ' % download_path)
+    # Cleanup command for exception handler
+    clean_cmd = ('cd %s ; rm -f chromiumos_test_image.bin ' % download_path)
+    exception_handler = RunCommandExceptionHandler(self._logger, self.log_level,
+                                                   self._ce, clean_cmd)
     if self.log_level != 'verbose':
       self._logger.LogOutput('CMD: %s' % command)
       print('(Uncompressing and un-tarring may take a couple of minutes...'
             'please be patient.)')
-    try:
-      retval = self._ce.ChrootRunCommand(chromeos_root, command)
-      if retval != 0:
-        raise MissingImage('Cannot uncompress image: %s.' % build_id)
-    except Exception:
-      # Exception in uncompressing file, so cleanup
-      command = ('cd /tmp/%s ; rm -f chromiumos_test_image.bin; ' % build_id)
-      # Remove the partially extracted bin file to avoid issues with future runs
-      _ = self._ce.ChrootRunCommand(chromeos_root, command)
-      # Raise exception again
-      raise
-
-    # Remove uncompressed file
-    command = ('cd /tmp/%s ; rm -f chromiumos_test_image.tar.xz; ' % build_id)
+    retval = self._ce.RunCommand(
+        command, except_handler=exception_handler.HandleException)
+    if retval != 0:
+      if self.log_level != 'verbose':
+        self._logger.LogOutput('CMD: %s' % clean_cmd)
+        print('(Removing file chromiumos_test_image.bin.)')
+      # Remove partially uncompressed file
+      _ = self._ce.RunCommand(clean_cmd)
+      # Raise exception for failure to uncompress
+      raise MissingImage('Cannot uncompress image: %s.' % build_id)
+
+    # Remove compressed image
+    command = ('cd %s ; rm -f chromiumos_test_image.tar.xz; ' % download_path)
     if self.log_level != 'verbose':
       self._logger.LogOutput('CMD: %s' % command)
       print('(Removing file chromiumos_test_image.tar.xz.)')
     # try removing file, its ok to have an error, print if encountered
-    retval = self._ce.ChrootRunCommand(chromeos_root, command)
+    retval = self._ce.RunCommand(command)
     if retval != 0:
       print('(Warning: Could not remove file chromiumos_test_image.tar.xz .)')
 
@@ -112,9 +138,10 @@ class ImageDownloader(object):
     status = 0
     gs_package_name = ('gs://chromeos-image-archive/%s/%s' %
                        (build_id, package_file_name))
+    gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
     if not test_flag.GetTestMode():
-      cmd = 'gsutil ls %s' % gs_package_name
-      status = self._ce.ChrootRunCommand(chromeos_root, cmd)
+      cmd = '%s ls %s' % (gsutil_cmd, gs_package_name)
+      status = self._ce.RunCommand(cmd)
     if status != 0:
       raise MissingFile('Cannot find autotest package file: %s.' %
                         package_file_name)
@@ -132,33 +159,34 @@ class ImageDownloader(object):
     # Check to see if the package file has already been downloaded.  If not,
     # download it.
     if not os.path.exists(package_path):
-      command = 'gsutil cp %s /tmp/%s' % (gs_package_name, build_id)
+      command = '%s cp %s %s' % (gsutil_cmd, gs_package_name, download_path)
 
       if self.log_level != 'verbose':
         self._logger.LogOutput('CMD: %s' % command)
-      status = self._ce.ChrootRunCommand(chromeos_root, command)
+      status = self._ce.RunCommand(command)
       if status != 0 or not os.path.exists(package_path):
         raise MissingFile('Cannot download package: %s .' % package_path)
 
   def UncompressSingleAutotestFile(self, chromeos_root, build_id,
                                    package_file_name, uncompress_cmd):
     # Uncompress file
-    command = ('cd /tmp/%s ; %s %s' %
-               (build_id, uncompress_cmd, package_file_name))
+    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
+    command = ('cd %s ; %s %s' %
+               (download_path, uncompress_cmd, package_file_name))
 
     if self.log_level != 'verbose':
       self._logger.LogOutput('CMD: %s' % command)
       print('(Uncompressing autotest file %s .)' % package_file_name)
-    retval = self._ce.ChrootRunCommand(chromeos_root, command)
+    retval = self._ce.RunCommand(command)
     if retval != 0:
       raise MissingFile('Cannot uncompress file: %s.' % package_file_name)
     # Remove uncompressed downloaded file
-    command = ('cd /tmp/%s ; rm -f %s' % (build_id, package_file_name))
+    command = ('cd %s ; rm -f %s' % (download_path, package_file_name))
     if self.log_level != 'verbose':
       self._logger.LogOutput('CMD: %s' % command)
       print('(Removing processed autotest file %s .)' % package_file_name)
     # try removing file, its ok to have an error, print if encountered
-    retval = self._ce.ChrootRunCommand(chromeos_root, command)
+    retval = self._ce.RunCommand(command)
     if retval != 0:
       print('(Warning: Could not remove file %s .)' % package_file_name)
 
@@ -167,11 +195,12 @@ class ImageDownloader(object):
     status = 0
     gs_package_name = ('gs://chromeos-image-archive/%s/%s' %
                        (build_id, package_file))
+    gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
     if not test_flag.GetTestMode():
-      cmd = 'gsutil ls %s' % gs_package_name
+      cmd = '%s ls %s' % (gsutil_cmd, gs_package_name)
       if self.log_level != 'verbose':
         self._logger.LogOutput('CMD: %s' % cmd)
-      status = self._ce.ChrootRunCommand(chromeos_root, cmd)
+      status = self._ce.RunCommand(cmd)
       if status != 0:
         print('(Warning: Could not find file %s )' % gs_package_name)
         return 1
@@ -184,17 +213,7 @@ class ImageDownloader(object):
     autotest_server_package_name = ('autotest_server_package.tar.bz2')
     autotest_control_files_name = ('control_files.tar')
 
-    # Quickly verify if the files are there
-    # If not, just exit with warning
-    status = self.VerifyAutotestFilesExist(chromeos_root, build_id,
-                                      autotest_packages_name)
-    if status != 0:
-      default_autotest_dir = '~/trunk/src/third_party/autotest/files'
-      print('(Warning: Could not download autotest packages .)\n'
-            '(Warning: Defaulting autotest path to %s .' %
-            default_autotest_dir)
-      return default_autotest_dir
-
+    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
     # Autotest directory relative path wrt chroot
     autotest_rel_path = os.path.join('/tmp', build_id, 'autotest_files')
     # Absolute Path to download files
@@ -202,6 +221,18 @@ class ImageDownloader(object):
                                  'autotest_files')
 
     if not os.path.exists(autotest_path):
+      # Quickly verify if the files are present on server
+      # If not, just exit with warning
+      status = self.VerifyAutotestFilesExist(chromeos_root, build_id,
+                                             autotest_packages_name)
+      if status != 0:
+        default_autotest_dir = '~/trunk/src/third_party/autotest/files'
+        print('(Warning: Could not find autotest packages .)\n'
+              '(Warning: Defaulting autotest path to %s .' %
+              default_autotest_dir)
+        return default_autotest_dir
+
+      # Files exist on server, download and uncompress them
       self.DownloadSingleAutotestFile(chromeos_root, build_id,
                                       autotest_packages_name)
       self.DownloadSingleAutotestFile(chromeos_root, build_id,
@@ -218,11 +249,11 @@ class ImageDownloader(object):
                                         autotest_control_files_name,
                                         'tar -xvf ')
       # Rename created autotest directory to autotest_files
-      command = ('cd /tmp/%s ; mv autotest autotest_files' % build_id)
+      command = ('cd %s ; mv autotest autotest_files' % download_path)
       if self.log_level != 'verbose':
         self._logger.LogOutput('CMD: %s' % command)
         print('(Moving downloaded autotest files to autotest_files)')
-      retval = self._ce.ChrootRunCommand(chromeos_root, command)
+      retval = self._ce.RunCommand(command)
       if retval != 0:
         raise MissingFile('Could not create directory autotest_files')
 
@@ -237,8 +268,9 @@ class ImageDownloader(object):
     # download it.
     status = 0
     if not test_flag.GetTestMode():
-      cmd = 'gsutil ls %s' % image_name
-      status = self._ce.ChrootRunCommand(chromeos_root, cmd)
+      gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
+      cmd = '%s ls %s' % (gsutil_cmd, image_name)
+      status = self._ce.RunCommand(cmd)
     if status != 0:
       raise MissingImage('Cannot find official image: %s.' % image_name)
 
diff --git a/crosperf/download_images_unittest.py b/crosperf/download_images_unittest.py
index e2bd7318..24a94a37 100755
--- a/crosperf/download_images_unittest.py
+++ b/crosperf/download_images_unittest.py
@@ -61,13 +61,15 @@ class ImageDownloaderTestcast(unittest.TestCase):
     mock_mkdirs.assert_called_with(
         '/usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0')
 
-    # Verify we called ChrootRunCommand once, with proper arguments.
-    self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 1)
-    mock_cmd_exec.ChrootRunCommand.assert_called_with(
-        '/usr/local/home/chromeos', 'gsutil cp '
-        'gs://chromeos-image-archive/lumpy-release/R36-5814.0.0/'
-        'chromiumos_test_image.tar.xz'
-        ' /tmp/lumpy-release/R36-5814.0.0')
+    # Verify we called RunCommand once, with proper arguments.
+    self.assertEqual(mock_cmd_exec.RunCommand.call_count, 1)
+    expected_args = (
+        '/usr/local/home/chromeos/chromium/tools/depot_tools/gsutil.py '
+        'cp gs://chromeos-image-archive/lumpy-release/R36-5814.0.0/'
+        'chromiumos_test_image.tar.xz '
+        '/usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0')
+
+    mock_cmd_exec.RunCommand.assert_called_with(expected_args)
 
     # Reset the velues in the mocks; set os.path.exists to always return True.
     mock_path_exists.reset_mock()
@@ -112,20 +114,33 @@ class ImageDownloaderTestcast(unittest.TestCase):
         '/usr/local/home/chromeos/chroot/tmp/lumpy-release/'
         'R36-5814.0.0/chromiumos_test_image.bin')
 
-    # Verify ChrootRunCommand was called twice, with correct arguments.
-    self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 2)
-    call_args_0 = mock_cmd_exec.ChrootRunCommand.call_args_list[0][0]
-    call_args_1 = mock_cmd_exec.ChrootRunCommand.call_args_list[1][0]
-    expected_arg_0 = ('/usr/local/home/chromeos',
-                      'cd /tmp/lumpy-release/R36-5814.0.0 ; '
-                      'tar -Jxf chromiumos_test_image.tar.xz ')
-    expected_arg_1 = ('/usr/local/home/chromeos',
-                      'cd /tmp/lumpy-release/R36-5814.0.0 ; '
-                      'rm -f chromiumos_test_image.bin; ')
-    self.assertEqual(call_args_0, expected_arg_0)
-    self.assertEqual(call_args_1, expected_arg_1)
-
-    # Set os.path.exists to always return False and run uncompress.
+    # Verify RunCommand was called twice with correct arguments.
+    self.assertEqual(mock_cmd_exec.RunCommand.call_count, 2)
+    print(mock_cmd_exec.RunCommand.call_args_list)
+    # Call 1, should have 2 arguments
+    self.assertEqual(len(mock_cmd_exec.RunCommand.call_args_list[0]), 2)
+    actual_arg = mock_cmd_exec.RunCommand.call_args_list[0][0]
+    expected_arg = (
+        'cd /usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0 ; '
+        'tar -Jxf chromiumos_test_image.tar.xz ',)
+    self.assertEqual(expected_arg, actual_arg)
+    # 2nd arg must be exception handler
+    except_handler_string = 'RunCommandExceptionHandler.HandleException'
+    self.assertTrue(
+        except_handler_string in
+        repr(mock_cmd_exec.RunCommand.call_args_list[0][1]))
+
+    # Call 2, should have 2 arguments
+    self.assertEqual(len(mock_cmd_exec.RunCommand.call_args_list[1]), 2)
+    actual_arg = mock_cmd_exec.RunCommand.call_args_list[1][0]
+    expected_arg = (
+        'cd /usr/local/home/chromeos/chroot/tmp/lumpy-release/R36-5814.0.0 ; '
+        'rm -f chromiumos_test_image.bin ',)
+    self.assertEqual(expected_arg, actual_arg)
+    # 2nd arg must be empty
+    self.assertTrue('{}' in repr(mock_cmd_exec.RunCommand.call_args_list[1][1]))
+
+    # Set os.path.exists to always return True and run uncompress.
     mock_path_exists.reset_mock()
     mock_cmd_exec.reset_mock()
     mock_path_exists.return_value = True
@@ -137,8 +152,8 @@ class ImageDownloaderTestcast(unittest.TestCase):
         '/usr/local/home/chromeos/chroot/tmp/lumpy-release/'
         'R36-5814.0.0/chromiumos_test_image.bin')
 
-    # Verify ChrootRunCommand was not called.
-    self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 0)
+    # Verify RunCommand was not called.
+    self.assertEqual(mock_cmd_exec.RunCommand.call_count, 0)
 
   def test_run(self):