Merge "Reuse --remote-image-dir" into main am: ebfb9e9c84

Original change: https://android-review.googlesource.com/c/platform/tools/acloud/+/2754430

Change-Id: Iba007b77ef54f56e4009251eac47cd3baadda791
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

6 files changed, 198 insertions, 16 deletions

diff --git a/internal/lib/cvd_utils.py b/internal/lib/cvd_utils.py
index bb6804c7..14f51178 100644
--- a/internal/lib/cvd_utils.py
+++ b/internal/lib/cvd_utils.py
@@ -17,10 +17,12 @@
 import collections
 import fnmatch
 import glob
+import json
 import logging
 import os
 import posixpath as remote_path
 import re
+import shlex
 import subprocess
 import tempfile
 import zipfile
@@ -611,6 +613,46 @@ def ParseRemoteHostAddress(instance_name):
     return None
 
 
+def LoadRemoteImageArgs(ssh_obj, remote_args_path):
+    """Load launch_cvd arguments from a remote path.
+
+    This method assumes that one acloud process accesses the path at a time.
+
+    Args:
+        ssh_obj: An Ssh object.
+        remote_image_dir: The remote path containing the arguments.
+
+    Returns:
+        A list of strings, the launch_cvd arguments.
+        None if the directory has not been initialized.
+    """
+    # If the file doesn't exist, the command returns 0 and outputs nothing.
+    args_str = ssh_obj.Run(shlex.quote(
+        f"test ! -f {remote_args_path} || cat {remote_args_path}"))
+    if not args_str:
+        return None
+    try:
+        return json.loads(args_str)
+    except json.JSONDecodeError as e:
+        logger.error("Unable to load %s: %s", remote_args_path, e)
+        return None
+
+
+def SaveRemoteImageArgs(ssh_obj, remote_args_path, launch_cvd_args):
+    """Save launch_cvd arguments to a remote path.
+
+    This method assumes that one acloud process accesses the path at a time.
+
+    Args:
+        ssh_obj: An Ssh object.
+        remote_args_path: The remote path containing the arguments.
+        launch_cvd_args: A list of strings, the launch_cvd arguments.
+    """
+    # The json string is interpreted twice by SSH client and remote shell.
+    args_str = shlex.quote(json.dumps(launch_cvd_args))
+    ssh_obj.Run(shlex.quote(f"echo {args_str} > {remote_args_path}"))
+
+
 def GetConfigFromRemoteAndroidInfo(ssh_obj, remote_image_dir):
     """Get config from android-info.txt on a remote host or a GCE instance.
 
diff --git a/internal/lib/cvd_utils_test.py b/internal/lib/cvd_utils_test.py
index 34294f8f..edb0ea04 100644
--- a/internal/lib/cvd_utils_test.py
+++ b/internal/lib/cvd_utils_test.py
@@ -309,6 +309,55 @@ class CvdUtilsTest(driver_test_lib.BaseDriverTest):
             "host-goldfish-192.0.2.1-5554-123456-sdk_x86_64-sdk")
         self.assertIsNone(result)
 
+    def testLoadRemoteImageArgs(self):
+        """Test LoadRemoteImageArgs."""
+        self.assertEqual(os.path, cvd_utils.remote_path)
+
+        with tempfile.TemporaryDirectory(prefix="cvd_utils") as temp_dir:
+            env = os.environ.copy()
+            env["HOME"] = temp_dir
+            # Execute the commands locally.
+            mock_ssh = mock.Mock()
+            mock_ssh.Run.side_effect = lambda cmd: subprocess.check_output(
+                "sh -c " + cmd, shell=True, cwd=temp_dir, env=env, text=True)
+            args_path = os.path.join(temp_dir, "args.txt")
+
+            # Test with an uninitialized directory.
+            args = cvd_utils.LoadRemoteImageArgs(mock_ssh, args_path)
+
+            self.assertIsNone(args)
+            mock_ssh.Run.assert_called_once_with(
+                f"'test ! -f {args_path} || cat {args_path}'")
+            self.assertFalse(os.path.exists(args_path))
+
+            # Test with an initialized directory.
+            self.CreateFile(args_path, b'["ok"]')
+
+            args = cvd_utils.LoadRemoteImageArgs(
+                mock_ssh, args_path)
+
+            self.assertEqual(args, ["ok"])
+
+    def testSaveRemoteImageArgs(self):
+        """Test SaveRemoteImageArgs."""
+        self.assertEqual(os.path, cvd_utils.remote_path)
+
+        with tempfile.TemporaryDirectory(prefix="cvd_utils") as temp_dir:
+            env = os.environ.copy()
+            env["HOME"] = temp_dir
+            mock_ssh = mock.Mock()
+            mock_ssh.Run.side_effect = lambda cmd: subprocess.check_call(
+                "sh -c " + cmd, shell=True, cwd=temp_dir, env=env, text=True)
+            args_path = os.path.join(temp_dir, "args.txt")
+
+            cvd_utils.SaveRemoteImageArgs(mock_ssh, args_path, ["ok"])
+
+            mock_ssh.Run.assert_called_with(
+                f"""'echo '"'"'["ok"]'"'"' > {args_path}'""")
+            with open(args_path, "r", encoding="utf-8") as args_file:
+                self.assertEqual(args_file.read().strip(), '["ok"]')
+
+
     def testGetConfigFromRemoteAndroidInfo(self):
         """Test GetConfigFromRemoteAndroidInfo."""
         mock_ssh = mock.Mock()
diff --git a/internal/lib/ssh.py b/internal/lib/ssh.py
index 53916b88..524a297e 100755
--- a/internal/lib/ssh.py
+++ b/internal/lib/ssh.py
@@ -29,7 +29,7 @@ logger = logging.getLogger(__name__)
 _SSH_CMD = ("-i %(rsa_key_file)s -o LogLevel=ERROR -o ControlPath=none "
             "-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no")
 _SSH_IDENTITY = "-l %(login_user)s %(ip_addr)s"
-_SSH_CMD_MAX_RETRY = 5
+SSH_CMD_DEFAULT_RETRY = 5
 _SSH_CMD_RETRY_SLEEP = 3
 _CONNECTION_TIMEOUT = 10
 _MAX_REPORTED_ERROR_LINES = 10
@@ -108,6 +108,9 @@ def _SshLogOutput(cmd, timeout=None, show_output=False, hide_error_msg=False):
         show_output: Boolean, True to show command output in screen.
         hide_error_msg: Boolean, True to hide error message.
 
+    Returns:
+        A string, stdout and stderr.
+
     Raises:
         errors.DeviceConnectionError: Failed to connect to the GCE instance.
         subprocess.CalledProcessError: The process exited with an error on the instance.
@@ -145,6 +148,7 @@ def _SshLogOutput(cmd, timeout=None, show_output=False, hide_error_msg=False):
         if constants.ERROR_MSG_WEBRTC_NOT_SUPPORT in stdout:
             raise errors.LaunchCVDFail(constants.ERROR_MSG_WEBRTC_NOT_SUPPORT)
         raise subprocess.CalledProcessError(process.returncode, cmd)
+    return stdout
 
 
 def _GetErrorMessage(stdout):
@@ -185,7 +189,7 @@ def _FilterUnusedContent(content):
 
 
 def ShellCmdWithRetry(cmd, timeout=None, show_output=False,
-                      retry=_SSH_CMD_MAX_RETRY):
+                      retry=SSH_CMD_DEFAULT_RETRY):
     """Runs a shell command on remote device.
 
     If the network is unstable and causes SSH connect fail, it will retry. When
@@ -199,12 +203,15 @@ def ShellCmdWithRetry(cmd, timeout=None, show_output=False,
         show_output: Boolean, True to show command output in screen.
         retry: Integer, the retry times.
 
+    Returns:
+        A string, stdout and stderr.
+
     Raises:
         errors.DeviceConnectionError: For any non-zero return code of remote_cmd.
         errors.LaunchCVDFail: Happened on launch_cvd with specific pattern of error message.
         subprocess.CalledProcessError: The process exited with an error on the instance.
     """
-    utils.RetryExceptionType(
+    return utils.RetryExceptionType(
         exception_types=(errors.DeviceConnectionError,
                          errors.LaunchCVDFail,
                          subprocess.CalledProcessError),
@@ -257,7 +264,7 @@ class Ssh():
                 extra_args_ssh_tunnel)
 
     def Run(self, target_command, timeout=None, show_output=False,
-            retry=_SSH_CMD_MAX_RETRY):
+            retry=SSH_CMD_DEFAULT_RETRY):
         """Run a shell command over SSH on a remote instance.
 
         Example:
@@ -273,11 +280,15 @@ class Ssh():
             timeout: Integer, the maximum time to wait for the command to respond.
             show_output: Boolean, True to show command output in screen.
             retry: Integer, the retry times.
+
+        Returns:
+            A string, stdout and stderr.
         """
-        ShellCmdWithRetry(self.GetBaseCmd(constants.SSH_BIN) + " " + target_command,
-                          timeout,
-                          show_output,
-                          retry)
+        return ShellCmdWithRetry(
+            self.GetBaseCmd(constants.SSH_BIN) + " " + target_command,
+            timeout,
+            show_output,
+            retry)
 
     def GetBaseCmd(self, execute_bin):
         """Get a base command over SSH on a remote instance.
@@ -346,7 +357,7 @@ class Ssh():
                 "Ssh isn't ready in the remote instance.") from e
 
     @utils.TimeExecute(function_description="Waiting for SSH server")
-    def WaitForSsh(self, timeout=None, max_retry=_SSH_CMD_MAX_RETRY):
+    def WaitForSsh(self, timeout=None, max_retry=SSH_CMD_DEFAULT_RETRY):
         """Wait until the remote instance is ready to accept commands over SSH.
 
         Args:
diff --git a/internal/lib/ssh_test.py b/internal/lib/ssh_test.py
index d35a6086..72ba3970 100644
--- a/internal/lib/ssh_test.py
+++ b/internal/lib/ssh_test.py
@@ -93,8 +93,9 @@ class SshTest(driver_test_lib.BaseDriverTest):
     def testSshRunCmd(self):
         """Test ssh run command."""
         self.Patch(subprocess, "Popen", return_value=self.created_subprocess)
+        self.created_subprocess.communicate.return_value = ("stdout", "")
         ssh_object = ssh.Ssh(self.FAKE_IP, self.FAKE_SSH_USER, self.FAKE_SSH_PRIVATE_KEY_PATH)
-        ssh_object.Run("command")
+        self.assertEqual("stdout", ssh_object.Run("command"))
         expected_cmd = (
             "exec /usr/bin/ssh -i /fake/acloud_rea -o LogLevel=ERROR -o ControlPath=none "
             "-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no "
@@ -109,11 +110,12 @@ class SshTest(driver_test_lib.BaseDriverTest):
     def testSshRunCmdwithExtraArgs(self):
         """test ssh rum command with extra command."""
         self.Patch(subprocess, "Popen", return_value=self.created_subprocess)
+        self.created_subprocess.communicate.return_value = ("stdout", "")
         ssh_object = ssh.Ssh(self.FAKE_IP,
                              self.FAKE_SSH_USER,
                              self.FAKE_SSH_PRIVATE_KEY_PATH,
                              self.FAKE_EXTRA_ARGS_SSH)
-        ssh_object.Run("command")
+        self.assertEqual("stdout", ssh_object.Run("command"))
         expected_cmd = (
             "exec /usr/bin/ssh -i /fake/acloud_rea -o LogLevel=ERROR -o ControlPath=none "
             "-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no "
diff --git a/public/actions/remote_host_cf_device_factory.py b/public/actions/remote_host_cf_device_factory.py
index 403e15d9..bfb43c84 100644
--- a/public/actions/remote_host_cf_device_factory.py
+++ b/public/actions/remote_host_cf_device_factory.py
@@ -41,6 +41,7 @@ logger = logging.getLogger(__name__)
 _ALL_FILES = "*"
 _HOME_FOLDER = os.path.expanduser("~")
 _TEMP_PREFIX = "acloud_remote_host"
+_IMAGE_ARGS_FILE_NAME = "acloud_image_args.txt"
 
 
 class RemoteHostDeviceFactory(base_device_factory.BaseDeviceFactory):
@@ -90,15 +91,18 @@ class RemoteHostDeviceFactory(base_device_factory.BaseDeviceFactory):
             A string, representing instance name.
         """
         start_time = time.time()
+        self._compute_client.SetStage(constants.STAGE_SSH_CONNECT)
         instance = self._InitRemotehost()
         start_time = self._compute_client.RecordTime(
             constants.TIME_GCE, start_time)
 
         process_artifacts_timestart = start_time
+        self._compute_client.SetStage(constants.STAGE_ARTIFACT)
         image_args = self._ProcessRemoteHostArtifacts()
         start_time = self._compute_client.RecordTime(
             constants.TIME_ARTIFACT, start_time)
 
+        self._compute_client.SetStage(constants.STAGE_BOOT_UP)
         error_msg = self._LaunchCvd(image_args, process_artifacts_timestart)
         start_time = self._compute_client.RecordTime(
             constants.TIME_LAUNCH, start_time)
@@ -147,7 +151,6 @@ class RemoteHostDeviceFactory(base_device_factory.BaseDeviceFactory):
         Returns:
             A string, representing instance name.
         """
-        self._compute_client.SetStage(constants.STAGE_SSH_CONNECT)
         # Get product name from the img zip file name or TARGET_PRODUCT.
         image_name = os.path.basename(
             self._local_image_artifact) if self._local_image_artifact else ""
@@ -175,7 +178,31 @@ class RemoteHostDeviceFactory(base_device_factory.BaseDeviceFactory):
         return instance
 
     def _ProcessRemoteHostArtifacts(self):
-        """Process remote host artifacts.
+        """Initialize or reuse the images on the remote host.
+
+        Returns:
+            A list of strings, the launch_cvd arguments.
+        """
+        remote_image_dir = self._avd_spec.remote_image_dir
+        if remote_image_dir:
+            remote_args_path = remote_path.join(remote_image_dir,
+                                                _IMAGE_ARGS_FILE_NAME)
+            launch_cvd_args = cvd_utils.LoadRemoteImageArgs(
+                self._ssh, remote_args_path)
+            if launch_cvd_args is not None:
+                logger.info("Reuse the images in %s", remote_image_dir)
+                return launch_cvd_args
+            logger.info("Create images in %s", remote_image_dir)
+
+        launch_cvd_args = self._InitRemoteImageDir()
+
+        if remote_image_dir:
+            cvd_utils.SaveRemoteImageArgs(self._ssh, remote_args_path,
+                                          launch_cvd_args)
+        return launch_cvd_args
+
+    def _InitRemoteImageDir(self):
+        """Create remote host artifacts.
 
         - If images source is local, tool will upload images from local site to
           remote host.
@@ -186,8 +213,6 @@ class RemoteHostDeviceFactory(base_device_factory.BaseDeviceFactory):
         Returns:
             A list of strings, the launch_cvd arguments.
         """
-        # TODO(b/293966645): Check if --remote-image-dir is initialized.
-        self._compute_client.SetStage(constants.STAGE_ARTIFACT)
         self._ssh.Run(f"mkdir -p {self._GetArtifactPath()}")
 
         launch_cvd_args = []
@@ -421,7 +446,6 @@ class RemoteHostDeviceFactory(base_device_factory.BaseDeviceFactory):
         Returns:
             The error message as a string. An empty string represents success.
         """
-        self._compute_client.SetStage(constants.STAGE_BOOT_UP)
         config = cvd_utils.GetConfigFromRemoteAndroidInfo(
             self._ssh, self._GetArtifactPath())
         cmd = cvd_utils.GetRemoteLaunchCvdCmd(
diff --git a/public/actions/remote_host_cf_device_factory_test.py b/public/actions/remote_host_cf_device_factory_test.py
index 057e0bcf..9ee010de 100644
--- a/public/actions/remote_host_cf_device_factory_test.py
+++ b/public/actions/remote_host_cf_device_factory_test.py
@@ -385,6 +385,60 @@ class RemoteHostDeviceFactoryTest(driver_test_lib.BaseDriverTest):
         self.assertFalse(factory.GetFailures())
         self.assertDictEqual({"inst": [log]}, factory.GetLogs())
 
+    @mock.patch("acloud.public.actions.remote_host_cf_device_factory.ssh")
+    @mock.patch("acloud.public.actions.remote_host_cf_device_factory."
+                "cvd_utils")
+    @mock.patch("acloud.public.actions.remote_host_cf_device_factory."
+                "subprocess.check_call")
+    @mock.patch("acloud.public.actions.remote_host_cf_device_factory.glob")
+    @mock.patch("acloud.public.actions.remote_host_cf_device_factory.pull")
+    def testCreateInstanceWithRemoteImageDir(self, _mock_pull, mock_glob,
+                                             _mock_check_call, mock_cvd_utils,
+                                             mock_ssh):
+        """Test CreateInstance with AvdSpec.remote_image_dir."""
+        mock_avd_spec = self._CreateMockAvdSpec()
+        mock_avd_spec.remote_image_dir = "mock_img_dir"
+
+        mock_ssh_obj = mock.Mock()
+        mock_ssh.Ssh.return_value = mock_ssh_obj
+        mock_ssh_obj.GetBaseCmd.return_value = "/mock/ssh"
+        # Test initializing the remote image dir.
+        mock_ssh_obj.Run.return_value = ""
+        mock_glob.glob.return_value = ["/mock/super.img"]
+        factory = remote_host_cf_device_factory.RemoteHostDeviceFactory(
+            mock_avd_spec)
+
+        mock_cvd_utils.GetRemoteHostBaseDir.return_value = "acloud_cf_1"
+        mock_cvd_utils.FormatRemoteHostInstanceName.return_value = "inst"
+        mock_cvd_utils.LoadRemoteImageArgs.return_value = None
+        mock_cvd_utils.AreTargetFilesRequired.return_value = False
+        mock_cvd_utils.UploadExtraImages.return_value = ["arg1"]
+        mock_cvd_utils.ExecuteRemoteLaunchCvd.return_value = ""
+        mock_cvd_utils.FindRemoteLogs.return_value = []
+
+        self._mock_build_api.GetFetchBuildArgs.return_value = ["-test"]
+
+        self.assertEqual("inst", factory.CreateInstance())
+        mock_cvd_utils.LoadRemoteImageArgs.assert_called_once_with(
+            mock_ssh_obj, "mock_img_dir/acloud_image_args.txt")
+        mock_cvd_utils.SaveRemoteImageArgs.assert_called_once_with(
+            mock_ssh_obj, "mock_img_dir/acloud_image_args.txt", ["arg1"])
+        self._mock_build_api.DownloadFetchcvd.assert_called_once()
+        print(mock_cvd_utils.GetRemoteLaunchCvdCmd.call_args[0][3])
+        self.assertEqual(["arg1"],
+                         mock_cvd_utils.GetRemoteLaunchCvdCmd.call_args[0][3])
+
+        # Test reusing the remote image dir.
+        mock_cvd_utils.LoadRemoteImageArgs.return_value = ["arg2"]
+        mock_cvd_utils.SaveRemoteImageArgs.reset_mock()
+        self._mock_build_api.DownloadFetchcvd.reset_mock()
+
+        self.assertEqual("inst", factory.CreateInstance())
+        mock_cvd_utils.SaveRemoteImageArgs.assert_not_called()
+        self._mock_build_api.DownloadFetchcvd.assert_not_called()
+        self.assertEqual(["arg2"],
+                         mock_cvd_utils.GetRemoteLaunchCvdCmd.call_args[0][3])
+
 
 if __name__ == "__main__":
     unittest.main()