Fix USB multiplexing to avoid lockups

Attempts to run firmware_FAFTSetup autotest in a continuous loop have
shown that after some time the system comes into a state when swinging
the usb_mux_sel1 multiplexer between DUT and host does not cause the
new USB device in servo port J3 to show up on neither side.

Further investigation has shown that the problem was actually with the
USB flash itself - unplugging it and plugging back in restored its
ability to initialize when the mux is connected to either DUT or host.

One of the reasons for such behavior could be the fact that the
switching process does not power down the USB interface when moving it
between DUT and host and back.

This change modifies the switching code such that every time the USB
flash is reconnected, it is powered off first.

The servo API in servo.py:Servo class is also being modified:

- a new method, switch_usbkey(), is being added to control which side
  the USB flash is connected to, instead of direct manipulating of the
  multiplexer state. The new method makes sure that proper power
  management takes place when reconnecting.

- another new method can be used to find out current state of the USB
  flash multiplexer: usb_flash_direction()

- enable_usb_hub() and disable_usb_hub() methods are being removed,
  instead of disabling the USB device it is now connected to the host

BUG=chrome-os-partner:15610
TEST=manual
 . ran firmware_FAFTSetup autotest more than 50 times in a row,
   observed it succeed each run, will leave it running overnight.
 . also ran firmware_FwScreenCloseLid, firmware_FwScreenPressPower and
   firmware_InvalidUSB tests, they all succeeded

Change-Id: I3aa34c14ec3929762520492587973db4c7b604da
Reviewed-on: https://gerrit.chromium.org/gerrit/39983
Tested-by: Vadim Bendebury <vbendeb@chromium.org>
Reviewed-by: Yusuf Mohsinally <mohsinally@google.com>
Reviewed-by: Todd Broch <tbroch@chromium.org>
Commit-Queue: Vadim Bendebury <vbendeb@chromium.org>
diff --git a/server/cros/faftsequence.py b/server/cros/faftsequence.py
index ee468a6..8d1411f 100644
--- a/server/cros/faftsequence.py
+++ b/server/cros/faftsequence.py
@@ -239,10 +239,10 @@
         """
         recovery_reason = 0
         logging.info('Try to retrieve recovery reason...')
-        if self.servo.get('usb_mux_sel1') == 'dut_sees_usbkey':
+        if self.servo.get_usbkey_direction() == 'dut':
             self.wait_fw_screen_and_plug_usb()
         else:
-            self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+            self.servo.switch_usbkey('dut')
 
         try:
             self.wait_for_client(install_deps=True)
@@ -316,7 +316,7 @@
                 return
 
         logging.info('Try boot into USB image...')
-        self.servo.enable_usb_hub(host=True)
+        self.servo.switch_usbkey('host')
         self.enable_rec_mode_and_reboot()
         self.wait_fw_screen_and_plug_usb()
         try:
@@ -383,9 +383,9 @@
         if self.check_setup_done('usb_check'):
             return
         if usb_dev:
-            assert self.servo.get('usb_mux_sel1') == 'servo_sees_usbkey'
+            assert self.servo.get_usbkey_direction() == 'host'
         else:
-            self.servo.enable_usb_hub(host=True)
+            self.servo.switch_usbkey('host')
             usb_dev = self.servo.probe_host_usb_dev()
             if not usb_dev:
                 raise error.TestError(
@@ -436,9 +436,9 @@
             host = True
 
         if host is True:
-            self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+            self.servo.switch_usbkey('host')
         elif host is False:
-            self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+            self.servo.switch_usbkey('dut')
 
 
     def get_server_address(self):
@@ -535,7 +535,7 @@
         # Now turn it on.
         self.servo.power_short_press()
         self.wait_for_client()
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.servo.switch_usbkey('dut')
 
         install_cmd = 'chromeos-install --yes'
         if firmware_update:
@@ -574,7 +574,7 @@
             self.clear_saved_firmware()
 
         # 'Unplug' any USB keys in the servo from the dut.
-        self.servo.enable_usb_hub(host=True)
+        self.servo.switch_usbkey('host')
         # Mark usb_check done so it won't check a test image in USB anymore.
         self.mark_setup_done('usb_check')
         self.mark_setup_done('reimage')
@@ -871,14 +871,14 @@
     def wait_fw_screen_and_unplug_usb(self):
         """Wait for firmware warning screen and then unplug the servo USB."""
         time.sleep(self.delay.load_usb)
-        self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        self.servo.switch_usbkey('host')
         time.sleep(self.delay.between_usb_plug)
 
 
     def wait_fw_screen_and_plug_usb(self):
         """Wait for firmware warning screen and then unplug and plug the USB."""
         self.wait_fw_screen_and_unplug_usb()
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.servo.switch_usbkey('dut')
 
 
     def wait_fw_screen_and_press_power(self):
@@ -1021,7 +1021,7 @@
     def enable_keyboard_dev_mode(self):
         logging.info("Enabling keyboard controlled developer mode")
         # Plug out USB disk for preventing recovery boot without warning
-        self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        self.servo.switch_usbkey('host')
         # Rebooting EC with rec mode on. Should power on AP.
         self.enable_rec_mode_and_reboot()
         self.wait_for_client_offline()
@@ -1200,7 +1200,7 @@
         shim to reset TPM values.
         """
         # Unplug USB first to avoid the complicated USB autoboot cases.
-        self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        self.servo.switch_usbkey('host')
         is_dev = self.checkers.crossystem_checker({'devsw_boot': '1'})
         if not is_dev:
             self.enable_dev_mode_and_reboot()
diff --git a/server/cros/servo.py b/server/cros/servo.py
index 7b1965d..9b8d1e2 100644
--- a/server/cros/servo.py
+++ b/server/cros/servo.py
@@ -62,6 +62,8 @@
 
     # Time between an usb disk plugged-in and detected in the system.
     USB_DETECTION_DELAY = 10
+    # Time to keep USB power off before and after USB mux direction is changed
+    USB_POWEROFF_DELAY = 2
 
     KEY_MATRIX_ALT_0 = {
         'ctrl_refresh':  ['0', '0', '0', '1'],
@@ -145,6 +147,13 @@
         self._is_localhost = (servo_host == 'localhost')
         self._target_host = target_host
 
+        # a string, showing what interface (host or dut) the USB device is
+        # connected to.
+        self._usb_position = None
+        self.set('dut_hub_pwren', 'on')
+        self.set('usb_mux_oe1', 'on')
+        self.switch_usbkey('host')
+
         # Commands on the servo host must be run by the superuser. Our account
         # on Beaglebone is root, but locally we might be running as a
         # different user. If so - `sudo ' will have to be added to the
@@ -361,35 +370,6 @@
         """Disable development mode on device."""
         self.set('dev_mode', 'off')
 
-    def enable_usb_hub(self, host=False):
-        """Enable Servo's USB/ethernet hub.
-
-        This is equivalent to plugging in the USB devices attached to Servo to
-        the host (if |host| is True) or dut (if |host| is False).
-        For host=False, requires that the USB out on the servo board is
-        connected to a USB in port on the target device. Servo's USB ports are
-        labeled DUT_HUB_USB1 and DUT_HUB_USB2. Servo's ethernet port is also
-        connected to this hub. Servo's USB port DUT_HUB_IN is the output of the
-        hub.
-        """
-        self.set('dut_hub_pwren', 'on')
-        if host:
-            self.set('usb_mux_oe1', 'on')
-            self.set('usb_mux_sel1', 'servo_sees_usbkey')
-        else:
-            self.set('dut_hub_sel', 'dut_sees_hub')
-
-        self.set('dut_hub_on', 'yes')
-
-
-    def disable_usb_hub(self):
-        """Disable Servo's USB/ethernet hub.
-
-        This is equivalent to unplugging the USB devices attached to Servo.
-        """
-        self.set('dut_hub_on', 'no')
-
-
     def boot_devmode(self):
         """Boot a dev-mode device that is powered off."""
         self.power_short_press()
@@ -500,22 +480,22 @@
           A string of USB disk path, like '/dev/sdb', or None if not existed.
         """
         cmd = 'ls /dev/sd[a-z]'
-        original_value = self.get('usb_mux_sel1')
+        original_value = self.get_usbkey_direction()
 
         # Make the host unable to see the USB disk.
-        if original_value != 'dut_sees_usbkey':
-            self.set('usb_mux_sel1', 'dut_sees_usbkey')
+        if original_value != 'dut':
+            self.switch_usbkey('dut')
             time.sleep(self.USB_DETECTION_DELAY)
         no_usb_set = set(self.system_output(cmd, ignore_status=True).split())
 
         # Make the host able to see the USB disk.
-        self.set('usb_mux_sel1', 'servo_sees_usbkey')
+        self.switch_usbkey('host')
         time.sleep(self.USB_DETECTION_DELAY)
         has_usb_set = set(self.system_output(cmd, ignore_status=True).split())
 
         # Back to its original value.
-        if original_value != 'servo_sees_usbkey':
-            self.set('usb_mux_sel1', original_value)
+        if original_value != self.get_usbkey_direction():
+            self.switch_usbkey(original_value)
             time.sleep(self.USB_DETECTION_DELAY)
 
         diff_set = has_usb_set - no_usb_set
@@ -547,7 +527,7 @@
 
         # Set up Servo's usb mux.
         self.set('prtctl4_pwren', 'on')
-        self.enable_usb_hub(host=True)
+        self.switch_usbkey('host')
         if image_path:
             logging.info('Searching for usb device and copying image to it. '
                          'Please wait a few minutes...')
@@ -591,7 +571,7 @@
             self.enable_recovery_mode()
             self.power_short_press()
             time.sleep(Servo.RECOVERY_BOOT_DELAY)
-            self.set('usb_mux_sel1', 'dut_sees_usbkey')
+            self.switch_usbkey('dut')
             self.disable_recovery_mode()
 
             if host:
@@ -607,7 +587,7 @@
                     logging.error('Host failed to come back up in the allotted '
                                   'time: %d seconds.', wait_timeout)
                 logging.info('Removing the usb key from the DUT.')
-                self.disable_usb_hub()
+                self.switch_usbkey('host')
         except:
             # In case anything went wrong we want to make sure to do a clean
             # reset.
@@ -695,3 +675,51 @@
         if not self.is_localhost():
             image = self._scp_image(image)
         programmer.program_bootprom(board, self, image)
+
+    def switch_usbkey(self, side):
+        """Connect USB flash stick to either host or DUT.
+
+        This function switches the servo multiplexer to provide electrical
+        connection between the USB port J3 and either host or DUT side.
+
+        Switching is accompanied by powercycling of the USB stick, because it
+        sometimes gets wedged if the mux is switched while the stick power is
+        on.
+
+        @param side: a string, either 'dut' or 'host' - indicates which side
+                   the USB flash device is required to be connected to.
+
+        @raise: error.TestError in case the parameter is neither 'dut' not
+                   'host'
+        """
+
+        if self._usb_position == side:
+            return
+
+        if side == 'host':
+            mux_direction = 'servo_sees_usbkey'
+        elif side == 'dut':
+            mux_direction = 'dut_sees_usbkey'
+        else:
+            raise error.TestError('unknown USB mux setting: %s' % side)
+
+        self.set('prtctl4_pwren', 'off')
+        time.sleep(self.USB_POWEROFF_DELAY)
+        self.set('usb_mux_sel1', mux_direction)
+        time.sleep(self.USB_POWEROFF_DELAY)
+        self.set('prtctl4_pwren', 'on')
+
+        self._usb_position = side
+
+
+    def get_usbkey_direction(self):
+        """Get name of the side the USB device is connected to.
+
+        @return a string, either 'dut' or 'host'
+        """
+        if not self._usb_position:
+            if self.get('usb_mux_sel1').starstwith('dut'):
+                self._usb_position = 'dut'
+            else:
+                self._usb_position = 'host'
+        return self._usb_position
diff --git a/server/site_tests/autoupdate_EndToEndTest/autoupdate_EndToEndTest.py b/server/site_tests/autoupdate_EndToEndTest/autoupdate_EndToEndTest.py
index 0e9c224..acf6cde 100755
--- a/server/site_tests/autoupdate_EndToEndTest/autoupdate_EndToEndTest.py
+++ b/server/site_tests/autoupdate_EndToEndTest/autoupdate_EndToEndTest.py
@@ -395,16 +395,17 @@
 
 
     def _servo_dut_reboot(self, host, is_dev_mode, is_using_test_images,
-                          is_disable_usb_hub=False):
+                          disconnect_usbkey=False):
         """Reboots a DUT.
 
         @param host: a host object
         @param is_dev_mode: whether or not the DUT is in dev mode
         @param is_using_test_images: whether or not a test image should be
                assumed
-        @param is_disable_usb_hub: disabled the servo USB hub in between power
-               off/on cycles; this is useful when (for example) a USB booted
-               device need not see the attached USB key after the reboot.
+        @param disconnect_usbkey: detach USB flash device from the DUT before
+               powering it back up; this is useful when (for example) a USB
+               booted device need not see the attached USB key after the
+               reboot.
 
         @raise error.TestFail if DUT fails to reboot.
 
@@ -412,8 +413,8 @@
         logging.info('rebooting dut')
         host.servo.power_long_press()
         _wait(self._WAIT_AFTER_SHUTDOWN_SECONDS, 'after shutdown')
-        if is_disable_usb_hub:
-            host.servo.disable_usb_hub()
+        if disconnect_usbkey:
+            host.servo.switch_usbkey('host')
         self._servo_dut_power_up(host, is_dev_mode)
         if is_using_test_images:
             if not host.wait_up(timeout=host.BOOT_TIMEOUT):
@@ -444,7 +445,7 @@
 
         # Reboot the DUT after installation.
         self._servo_dut_reboot(host, is_dev_mode, False,
-                               is_disable_usb_hub=True)
+                               disconnect_usbkey=True)
 
 
     def _install_test_image(self, host, lorry_image_url, is_dev_mode):
@@ -480,7 +481,7 @@
 
         # Reboot the DUT after installation.
         self._servo_dut_reboot(host, is_dev_mode, True,
-                               is_disable_usb_hub=True)
+                               disconnect_usbkey=True)
 
 
     def _trigger_test_update(self, host, omaha_netloc):
diff --git a/server/site_tests/firmware_FwScreenCloseLid/firmware_FwScreenCloseLid.py b/server/site_tests/firmware_FwScreenCloseLid/firmware_FwScreenCloseLid.py
index 416c9c7..0772be2 100644
--- a/server/site_tests/firmware_FwScreenCloseLid/firmware_FwScreenCloseLid.py
+++ b/server/site_tests/firmware_FwScreenCloseLid/firmware_FwScreenCloseLid.py
@@ -30,7 +30,7 @@
     def wait_yuck_screen_and_close_lid(self):
         """Wait and trigger yuck screen and clod lid."""
         # Insert a corrupted USB stick. A yuck screen is expected.
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.servo.switch_usbkey('dut')
         time.sleep(self.delay.load_usb)
         self.wait_longer_fw_screen_and_close_lid()
 
@@ -40,7 +40,7 @@
         if self.client_attr.has_lid:
             self.assert_test_image_in_usb_disk()
             self.setup_dev_mode(dev_mode=True)
-            self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+            self.servo.switch_usbkey('host')
             usb_dev = self.servo.probe_host_usb_dev()
             # Corrupt the kernel of USB stick. It is needed for triggering a
             # yuck screen later.
@@ -49,7 +49,7 @@
 
     def cleanup(self):
         if self.client_attr.has_lid:
-            self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+            self.servo.switch_usbkey('host')
             usb_dev = self.servo.probe_host_usb_dev()
             # Restore the kernel of USB stick which is corrupted on setup phase.
             self.restore_usb_kernel(usb_dev)
diff --git a/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py b/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py
index 2fa06b1..514dbb4 100644
--- a/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py
+++ b/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py
@@ -28,7 +28,7 @@
     def wait_yuck_screen_and_press_power(self):
         """Insert corrupted USB for yuck screen and press power button."""
         # This USB stick will be removed in cleanup phase.
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.servo.switch_usbkey('dut')
         time.sleep(self.delay.load_usb)
         self.wait_longer_fw_screen_and_press_power()
 
@@ -37,7 +37,7 @@
         super(firmware_FwScreenPressPower, self).setup()
         self.assert_test_image_in_usb_disk()
         self.setup_dev_mode(dev_mode=True)
-        self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        self.servo.switch_usbkey('host')
         usb_dev = self.servo.probe_host_usb_dev()
         # Corrupt the kernel of USB stick. It is needed for triggering a
         # yuck screen later.
@@ -45,7 +45,7 @@
 
 
     def cleanup(self):
-        self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        self.servo.switch_usbkey('host')
         usb_dev = self.servo.probe_host_usb_dev()
         # Restore the kernel of USB stick which is corrupted on setup phase.
         self.restore_usb_kernel(usb_dev)
diff --git a/server/site_tests/firmware_InvalidUSB/firmware_InvalidUSB.py b/server/site_tests/firmware_InvalidUSB/firmware_InvalidUSB.py
index 4ca62c6..939d3c3 100644
--- a/server/site_tests/firmware_InvalidUSB/firmware_InvalidUSB.py
+++ b/server/site_tests/firmware_InvalidUSB/firmware_InvalidUSB.py
@@ -24,7 +24,7 @@
 
     def restore_usb(self):
         """Restore the USB image. USB plugs/unplugs happen in this method."""
-        self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        self.servo.switch_usbkey('host')
         usb_dev = self.servo.probe_host_usb_dev()
         self.restore_usb_kernel(usb_dev)
 
@@ -42,17 +42,17 @@
 
         self.restore_usb()
         time.sleep(self.delay.sync)
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.servo.switch_usbkey('dut')
 
 
     def setup(self):
         super(firmware_InvalidUSB, self).setup()
-        self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        self.servo.switch_usbkey('host')
         usb_dev = self.servo.probe_host_usb_dev()
         self.assert_test_image_in_usb_disk(usb_dev)
         self.corrupt_usb_kernel(usb_dev)
         self.setup_dev_mode(dev_mode=False)
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.servo.switch_usbkey('dut')
 
 
     def cleanup(self):
diff --git a/server/site_tests/platform_InstallTestImage/platform_InstallTestImage.py b/server/site_tests/platform_InstallTestImage/platform_InstallTestImage.py
index e617d0d..d58f0bd 100755
--- a/server/site_tests/platform_InstallTestImage/platform_InstallTestImage.py
+++ b/server/site_tests/platform_InstallTestImage/platform_InstallTestImage.py
@@ -21,7 +21,7 @@
         host.run('chromeos-install --yes',
                  timeout=self._INSTALL_TIMEOUT)
         host.servo.power_long_press()
-        host.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        host.servo.switch_usbkey('host')
         host.servo.power_short_press()
         if not host.wait_up(timeout=host.BOOT_TIMEOUT):
             raise error.TestFail('DUT failed to reboot installed test image'