VMAuWorker: save VM memory in case of VMTest failure

Saving the memory will enable inspection of the run-time state
of the VM, at the time VMTest ended.

While there: capture VM state in a separate try block from
log files. We should try to capture VM state even if we have
some trouble with log files.

BUG=chromium:321855
TEST=Manual (see below)
CQ-DEPEND=CL:177476
CQ-DEPEND=CL:178625

Manual test
-----------
- Download and untar
  https://storage.cloud.google.com/chromeos-image-archive/trybot-amd64-generic-paladin/R33-5064.0.0-b684/attempt1_failed_SimpleTestVerify_1_autotest_tests_vm_state.tar
- cd vm_state
- cros_start_vm --ssh_port=9230 --image_path=$PWD/vm_disk.bin --mem_path=$PWD/vm_memory.bin
  you should see the login screen without first seeing the boot logo

Change-Id: Id3e9f7031b703ffee3ed1673454d17a2b92fc200
Reviewed-on: https://chromium-review.googlesource.com/178636
Reviewed-by: Chris Sosa <sosa@chromium.org>
Commit-Queue: mukesh agrawal <quiche@chromium.org>
Tested-by: mukesh agrawal <quiche@chromium.org>
diff --git a/au_test_harness/vm_au_worker.py b/au_test_harness/vm_au_worker.py
index 8d95c63..cde5f3c 100644
--- a/au_test_harness/vm_au_worker.py
+++ b/au_test_harness/vm_au_worker.py
@@ -27,12 +27,17 @@
       cros_build_lib.Die('Need board to convert base image to vm.')
     self.whitelist_chrome_crashes = options.whitelist_chrome_crashes
 
-  def _KillExistingVM(self, pid_file):
+  def _KillExistingVM(self, pid_file, save_mem_path=None):
     """Kills an existing VM specified by the pid_file."""
-    if os.path.exists(pid_file):
-      cmd = ['./bin/cros_stop_vm', '--kvm_pid=%s' % pid_file]
-      cros_build_lib.RunCommand(cmd, print_cmd=False, error_code_ok=True,
-                                cwd=constants.CROSUTILS_DIR)
+    if not os.path.exists(pid_file):
+      return
+
+    cmd = ['./bin/cros_stop_vm', '--kvm_pid=%s' % pid_file]
+    if save_mem_path is not None:
+      cmd.append('--mem_path=%s' % save_mem_path)
+
+    cros_build_lib.RunCommand(cmd, print_cmd=False, error_code_ok=True,
+                              cwd=constants.CROSUTILS_DIR)
 
   def CleanUp(self):
     """Stop the vm after a test."""
@@ -45,7 +50,7 @@
     # well as the archive stage of cbuildbot. Make a private copy of
     # the VM image, to avoid any conflict.
     _, private_image_path = tempfile.mkstemp(
-        prefix="%s." % buildbot_constants.VM_IMAGE_PREFIX)
+        prefix="%s." % buildbot_constants.VM_DISK_PREFIX)
     shutil.copy(self.vm_image_path, private_image_path)
     self.TestInfo('Copied shared disk image %s to %s.' %
                   (self.vm_image_path, private_image_path))
@@ -61,15 +66,24 @@
     if not os.path.isdir(parent_dir):
       os.makedirs(parent_dir)
 
+    # Copy logs. Must be done before moving image, as this creates
+    # |fail_directory|.
     try:
-      # Copy logs. Must be done before moving image, as this creates
-      # |fail_directory|.
       shutil.copytree(log_directory, fail_directory)
-      self._KillExistingVM(self._kvm_pid_file)
+    except shutil.Error as e:
+      cros_build_lib.Warning(
+          'Ignoring errors while copying logs: %s', e)
+
+    # Copy VM state. This includes the disk image, and the memory
+    # image.
+    try:
+      _, mem_image_path = tempfile.mkstemp(
+        dir=fail_directory, prefix="%s." % buildbot_constants.VM_MEM_PREFIX)
+      self._KillExistingVM(self._kvm_pid_file, save_mem_path=mem_image_path)
       shutil.move(self.vm_image_path, fail_directory)
     except shutil.Error as e:
       cros_build_lib.Warning(
-          'Ignoring errors while copying logs or VM disk image: %s', e)
+          'Ignoring errors while copying VM files: %s', e)
 
   def UpdateImage(self, image_path, src_image_path='', stateful_change='old',
                   proxy_port='', private_key_path=None):