cros_flash: Allow stateful payload not to be copied to device

Some devices have very limited size stateful partition. Currently, cros
flash is failing for such devices because cros flash copies the payloads
into the DUT first before updating the device. Hence, the lack of space
problem. This CL adds a new flag --no-copy-payloads-to-device that if
passed the stateful update payload is not copied to the device first,
rather it is piped to the 'tar' command that untars the stateful update.

This approach in theory should make cros flash faster because it
eliminates one disk write and one disk read. However, experiments shows
that this flag makes cros flash slower but about 10-15 seconds (for a
electro device). Probably, the reason could be how inefficient the piped
file is read by the tar command. Hence, we decided not to enable this
option by default.

BUG=b:170359436
TEST=cros flash --no-ping <dut-ip> latest/reef/test
TEST=cros flash --no-ping --no-copy-payloads-to-device <dut-ip> latest/reef/test

Change-Id: Iecfbf3460bb4c48918edf63045d8372d7170cfc0
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2555531
Tested-by: Amin Hassani <ahassani@chromium.org>
Commit-Queue: Amin Hassani <ahassani@chromium.org>
Commit-Queue: Mike Frysinger <vapier@chromium.org>
Auto-Submit: Amin Hassani <ahassani@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
diff --git a/cli/cros/cros_flash.py b/cli/cros/cros_flash.py
index f62f5b6..22f8703 100644
--- a/cli/cros/cros_flash.py
+++ b/cli/cros/cros_flash.py
@@ -153,7 +153,12 @@
     update.add_argument(
         '--send-payload-in-parallel', default=False, action='store_true',
         help=('To speed up transfer payload files for long haul, chop '
-              'payload in chunks and transfer them in parallel'))
+              'payload in chunks and transfer them in parallel.'))
+    update.add_argument(
+        '--no-copy-payloads-to-device', dest='copy_payloads_to_device',
+        action='store_false', default=True,
+        help=('Do not copy the update payloads to the device. For now this '
+              'only works for the stateful payload.'))
     usb = parser.add_argument_group('USB specific options')
     usb.add_argument(
         '--install', default=False, action='store_true',
@@ -202,7 +207,8 @@
           yes=self.options.yes,
           force=self.options.force,
           debug=self.options.debug,
-          send_payload_in_parallel=self.options.send_payload_in_parallel)
+          send_payload_in_parallel=self.options.send_payload_in_parallel,
+          copy_payloads_to_device=self.options.copy_payloads_to_device)
       logging.notice('cros flash completed successfully.')
     except dev_server_wrapper.ImagePathError:
       logging.error('To get the latest remote image, please run:\n'
diff --git a/cli/flash.py b/cli/flash.py
index 2a6071d..16c137c 100644
--- a/cli/flash.py
+++ b/cli/flash.py
@@ -363,7 +363,8 @@
                board=None, src_image_to_delta=None, wipe=True, debug=False,
                yes=False, force=False, ssh_private_key=None, ping=True,
                disable_verification=False, send_payload_in_parallel=False,
-               clear_tpm_owner=False, version=None):
+               clear_tpm_owner=False, version=None,
+               copy_payloads_to_device=True):
     """Initializes RemoteDeviceUpdater"""
     if not stateful_update and not rootfs_update:
       raise ValueError('No update operation to perform; either stateful or'
@@ -389,6 +390,7 @@
     self.force = force
     self.send_payload_in_parallel = send_payload_in_parallel
     self.version = version
+    self.copy_payloads_to_device = copy_payloads_to_device
 
   def Cleanup(self):
     """Cleans up the temporary directory."""
@@ -526,7 +528,8 @@
               yes=self.yes,
               send_payload_in_parallel=self.send_payload_in_parallel,
               resolve_app_id_mismatch=True,
-              transfer_class=auto_updater_transfer.LocalTransfer)
+              transfer_class=auto_updater_transfer.LocalTransfer,
+              copy_payloads_to_device=self.copy_payloads_to_device)
           chromeos_AU.RunUpdate()
 
         except Exception:
@@ -552,7 +555,7 @@
           reboot=True, wipe=True, ssh_private_key=None, ping=True,
           disable_rootfs_verification=False, clear_cache=False, yes=False,
           force=False, debug=False, send_payload_in_parallel=False,
-          clear_tpm_owner=False, version=None):
+          clear_tpm_owner=False, version=None, copy_payloads_to_device=True):
   """Flashes a device, USB drive, or file with an image.
 
   This provides functionality common to `cros flash` and `brillo flash`
@@ -584,6 +587,9 @@
     send_payload_in_parallel: Transfer payloads in chunks in parallel to speed
         up transmissions for long haul between endpoints.
     version: Default version.
+    copy_payloads_to_device: If True, update payloads are copied to the
+        Chromium OS device first. Otherwise, they are piped through SSH.
+        Currently, this only applies to the stateful payloads.
 
   Raises:
     FlashError: An unrecoverable error occured.
@@ -630,7 +636,8 @@
         ping=ping,
         disable_verification=disable_rootfs_verification,
         send_payload_in_parallel=send_payload_in_parallel,
-        version=version)
+        version=version,
+        copy_payloads_to_device=copy_payloads_to_device)
     updater.Run()
   elif device.scheme == commandline.DEVICE_SCHEME_USB:
     path = osutils.ExpandPath(device.path) if device.path else ''
diff --git a/lib/auto_updater.py b/lib/auto_updater.py
index 64a4476..0571294 100644
--- a/lib/auto_updater.py
+++ b/lib/auto_updater.py
@@ -148,7 +148,8 @@
                reboot=True, disable_verification=False,
                send_payload_in_parallel=False, payload_filename=None,
                staging_server=None, clear_tpm_owner=False,
-               resolve_app_id_mismatch=False, ignore_appid=False):
+               resolve_app_id_mismatch=False, ignore_appid=False,
+               copy_payloads_to_device=True):
     """Initialize a ChromiumOSUpdater for auto-update a chromium OS device.
 
     Args:
@@ -190,6 +191,9 @@
           App ID. This allows mismatching the source and target version boards.
           One specific use case is updating between <board> and
           <board>-kernelnext images which have different App IDs.
+      copy_payloads_to_device: If True, update payloads are copied to the
+          Chromium OS device first. Otherwise, they are piped through SSH.
+          Currently, this only applies to the stateful payloads.
     """
     super(ChromiumOSUpdater, self).__init__(device, payload_dir)
 
@@ -234,6 +238,7 @@
     self._clear_tpm_owner = clear_tpm_owner
     self._resolve_app_id_mismatch = resolve_app_id_mismatch
     self._ignore_appid = ignore_appid
+    self._copy_payloads_to_device = copy_payloads_to_device
 
   @property
   def is_au_endtoendtest(self):
@@ -525,18 +530,25 @@
   def UpdateStateful(self):
     """Update the stateful partition of the device."""
     try:
-      stateful_update_payload = os.path.join(
-          self.device.work_dir, auto_updater_transfer.STATEFUL_FILENAME)
+      if self._copy_payloads_to_device:
+        self._transfer_obj.TransferStatefulUpdate()
+        stateful_update_payload = os.path.join(
+            self.device.work_dir, auto_updater_transfer.STATEFUL_FILENAME)
+      else:
+        stateful_update_payload = os.path.join(
+            self.payload_dir, auto_updater_transfer.STATEFUL_FILENAME)
 
       updater = stateful_updater.StatefulUpdater(self.device)
       updater.Update(
           stateful_update_payload,
+          is_payload_on_device=self._copy_payloads_to_device,
           update_type=(stateful_updater.StatefulUpdater.UPDATE_TYPE_CLOBBER if
                        self._clobber_stateful else None))
 
-      # Delete the stateful update file on success so it doesn't occupy extra
-      # disk space. On failure it will get cleaned up.
-      self.device.DeletePath(stateful_update_payload)
+      if self._copy_payloads_to_device:
+        # Delete the stateful update file on success so it doesn't occupy extra
+        # disk space. On failure it will get cleaned up.
+        self.device.DeletePath(stateful_update_payload)
     except stateful_updater.Error as e:
       error_msg = 'Stateful update failed with error: %s' % str(e)
       logging.exception(error_msg)
@@ -571,15 +583,6 @@
     # remainder of the update process.
     self.device.DeletePath(self.device_payload_dir, recursive=True)
 
-  def RunUpdateStateful(self):
-    """Run all processes needed by updating stateful.
-
-    1. Copy files to remote device needed by stateful update.
-    2. Do stateful update.
-    """
-    self._transfer_obj.TransferStatefulUpdate()
-    self.UpdateStateful()
-
   def RebootAndVerify(self):
     """Reboot and verify the remote device.
 
@@ -649,19 +652,20 @@
     """Update the device with image of specific version."""
     self._transfer_obj.CheckPayloads()
 
-    self._transfer_obj.TransferUpdateUtilsPackage()
-
-    restore_stateful = self.CheckRestoreStateful()
-    if restore_stateful:
-      self.RestoreStateful()
-
-    # Perform device updates.
+    restore_stateful = False
     if self._do_rootfs_update:
+      self._transfer_obj.TransferUpdateUtilsPackage()
+
+      restore_stateful = self.CheckRestoreStateful()
+      if restore_stateful:
+        self.RestoreStateful()
+
+      # Perform device updates.
       self.RunUpdateRootfs()
       logging.info('Rootfs update completed.')
 
     if self._do_stateful_update and not restore_stateful:
-      self.RunUpdateStateful()
+      self.UpdateStateful()
       logging.info('Stateful update completed.')
 
     if self._clear_tpm_owner:
@@ -897,7 +901,6 @@
     """Restore stateful partition for device."""
     logging.warning('Restoring the stateful partition.')
     self.PreSetupStatefulUpdate()
-    self._transfer_obj.TransferStatefulUpdate()
     self.ResetStatefulPartition()
     self.UpdateStateful()
     self.PostCheckStatefulUpdate()
diff --git a/lib/auto_updater_unittest.py b/lib/auto_updater_unittest.py
index 97790a5..be9fe64 100644
--- a/lib/auto_updater_unittest.py
+++ b/lib/auto_updater_unittest.py
@@ -34,7 +34,7 @@
 class ChromiumOSBaseUpdaterMock(partial_mock.PartialCmdMock):
   """Mock out all update and verify functions in ChromiumOSUpdater."""
   TARGET = 'chromite.lib.auto_updater.ChromiumOSUpdater'
-  ATTRS = ('RestoreStateful', 'UpdateStateful', 'UpdateRootfs',
+  ATTRS = ('RestoreStateful', 'UpdateRootfs',
            'SetupRootfsUpdate', 'RebootAndVerify')
 
   def __init__(self):
@@ -43,9 +43,6 @@
   def RestoreStateful(self, _inst, *_args, **_kwargs):
     """Mock out RestoreStateful."""
 
-  def UpdateStateful(self, _inst, *_args, **_kwargs):
-    """Mock out UpdateStateful."""
-
   def UpdateRootfs(self, _inst, *_args, **_kwargs):
     """Mock out UpdateRootfs."""
 
@@ -144,26 +141,6 @@
       self.assertFalse(
           self._transfer_mock.patched['TransferStatefulUpdate'].called)
 
-  def testTransferForStateful(self):
-    """Test Transfer functions' code path for stateful update.
-
-    When stateful update is enabled, update-utils and stateful update payload
-    are transferred. Rootfs update payload is not.
-    """
-    with remote_access.ChromiumOSDeviceHandler(
-        remote_access.TEST_IP) as device:
-      CrOS_AU = auto_updater.ChromiumOSUpdater(
-          device, None, self._payload_dir, do_rootfs_update=False,
-          transfer_class=auto_updater_transfer.LocalTransfer)
-      CrOS_AU.RunUpdate()
-      self.assertTrue(
-          self._transfer_mock.patched['CheckPayloads'].called)
-      self.assertTrue(
-          self._transfer_mock.patched['TransferUpdateUtilsPackage'].called)
-      self.assertFalse(
-          self._transfer_mock.patched['TransferRootfsUpdate'].called)
-      self.assertTrue(
-          self._transfer_mock.patched['TransferStatefulUpdate'].called)
 
 class ChromiumOSUpdatePreCheckTest(ChromiumOSUpdaterBaseTest):
   """Test precheck function."""
@@ -230,27 +207,47 @@
       CrOS_AU = auto_updater.ChromiumOSUpdater(
           device, None, self._payload_dir, do_stateful_update=False,
           transfer_class=auto_updater_transfer.LocalTransfer)
+      stateful_update_mock = self.PatchObject(auto_updater.ChromiumOSUpdater,
+                                              'UpdateStateful')
+
       CrOS_AU.RunUpdate()
       self.assertTrue(
           self._base_updater_mock.patched['SetupRootfsUpdate'].called)
       self.assertTrue(self._base_updater_mock.patched['UpdateRootfs'].called)
-      self.assertFalse(self._base_updater_mock.patched['UpdateStateful'].called)
+      stateful_update_mock.assert_not_called()
 
-  def testRunStateful(self):
-    """Test the update functions are called correctly.
-
-    Only UpdateStateful is called for stateful update.
-    """
+  def testUpdateStatefulCopyStateful(self):
+    """Tests stateful update when stateful payload is copied to device."""
     with remote_access.ChromiumOSDeviceHandler(
         remote_access.TEST_IP) as device:
       CrOS_AU = auto_updater.ChromiumOSUpdater(
           device, None, self._payload_dir, do_rootfs_update=False,
           transfer_class=auto_updater_transfer.LocalTransfer)
+      update_mock = self.PatchObject(stateful_updater.StatefulUpdater, 'Update')
+
       CrOS_AU.RunUpdate()
-      self.assertFalse(
-          self._base_updater_mock.patched['SetupRootfsUpdate'].called)
-      self.assertFalse(self._base_updater_mock.patched['UpdateRootfs'].called)
-      self.assertTrue(self._base_updater_mock.patched['UpdateStateful'].called)
+      update_mock.assert_called_with(
+          auto_updater_transfer.STATEFUL_FILENAME, is_payload_on_device=True,
+          update_type=None)
+      self.assertTrue(
+          self._transfer_mock.patched['TransferStatefulUpdate'].called)
+
+  def testUpdateStatefulNoCopyStateful(self):
+    """Tests stateful update when stateful payload is not copied to device."""
+    with remote_access.ChromiumOSDeviceHandler(
+        remote_access.TEST_IP) as device:
+      CrOS_AU = auto_updater.ChromiumOSUpdater(
+          device, None, self._payload_dir, do_rootfs_update=False,
+          transfer_class=auto_updater_transfer.LocalTransfer,
+          copy_payloads_to_device=False)
+      update_mock = self.PatchObject(stateful_updater.StatefulUpdater, 'Update')
+
+      CrOS_AU.RunUpdate()
+      update_mock.assert_called_with(
+          os.path.join(self._payload_dir,
+                       auto_updater_transfer.STATEFUL_FILENAME),
+          is_payload_on_device=False,
+          update_type=None)
 
   def testMismatchedAppId(self):
     """Tests if App ID is set to empty when there's a mismatch."""
diff --git a/lib/stateful_updater.py b/lib/stateful_updater.py
index 0429ffb..732ae0e 100644
--- a/lib/stateful_updater.py
+++ b/lib/stateful_updater.py
@@ -51,22 +51,29 @@
     self._update_type_file = os.path.join(self._stateful_dir,
                                           self._UPDATE_TYPE_FILE)
 
-  def Update(self, payload_path_on_device, update_type=None):
+  def Update(self, payload_path, is_payload_on_device=True, update_type=None):
     """Updates the stateful partition given the update file.
 
     Args:
-      payload_path_on_device: The path to the stateful update (stateful.tgz)
-        on the DUT.
+      payload_path: The path to the stateful update (stateful.tgz).
+      is_payload_on_device: True if the payload is on the device. False if it
+        is on the workstation.
       update_type: The type of the stateful update to be marked. Accepted
         values: 'standard' (default) and 'clobber'.
     """
-    if not self._device.IfPathExists(payload_path_on_device):
-      raise Error('Missing the file: %s' % payload_path_on_device)
-
     try:
       cmd = ['tar', '--ignore-command-error', '--overwrite',
-             '--directory', self._stateful_dir, '-xzf', payload_path_on_device]
-      self._device.run(cmd)
+             '--directory', self._stateful_dir, '-xzf']
+      if is_payload_on_device:
+        if not self._device.IfPathExists(payload_path):
+          raise Error('Missing the file: %s' % payload_path)
+
+        cmd += [payload_path]
+        self._device.run(cmd)
+      else:
+        with open(payload_path, 'rb') as f:
+          cmd += ['-']
+          self._device.run(cmd, input=f)
     except cros_build_lib.RunCommandError as e:
       raise Error('Failed to untar the stateful update with error %s' % e)