servo_updater: make more robust on resets

long-term, we need to pull this into hdctools, rather than reimplement
everything twice. Short term, this is a fine solution. It essentially
makes sure that
1. we only keep one pty/stm32uart/stm32usb object around for
communication
2. we always reset it properly when it could need a reset e.g. when the
stm32 is rebooting, or a new firmware was flashed
3. it expands the timeout for the chip to come back to 2s

the tiny_servod will eventually also land in hdctools, though for now it
just ensures that we can reset the pyusb communication without larger
issues and a larger refactor.

BUG=chromium:1152838
BRANCH=None
// Timeout before change, runs after change
TEST=sudo servo_updater --board servo_micro
// Timeout before change, runs after change
TEST=sudo servo_updater --board sweetberry
// Timeout before change, runs after change
TEST=sudo servo_updater --board servo_micro --force
// to show the serialname support. This is a fake serial, and it gets
stuck waiting
TEST=sudo servo_updater -s MICRO-S-2009020022 --board servo_micro
// to show the serialname support. This is a real serial and it proceeds
TEST=sudo servo_updater -s MICRO-S-2009020022 --board servo_micro

Change-Id: I747ca69881c13c1aadd8e90a35badecbf4e6a09e
Signed-off-by: Ruben Rodriguez Buchillon <coconutruben@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/2580854
Reviewed-by: Otabek Kasimov <otabek@google.com>
Reviewed-by: Garry Wang <xianuowang@chromium.org>
(cherry picked from commit c1ff0aed409b5ddf03538448e2e4c773b97986f4)
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/2585855

2 files changed, 88 insertions, 48 deletions

diff --git a/extra/tigertool/ecusb/tiny_servod.py b/extra/tigertool/ecusb/tiny_servod.py
new file mode 100644
index 0000000000..aec87995a4
--- /dev/null
+++ b/extra/tigertool/ecusb/tiny_servod.py
@@ -0,0 +1,50 @@
+# Copyright 2020 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Helper class to facilitate communication to servo ec console."""
+
+from ecusb import pty_driver
+from ecusb import stm32uart
+
+
+class TinyServod(object):
+  """Helper class to wrap a pty_driver with interface."""
+
+  def __init__(self, vid, pid, interface, serialname=None, debug=False):
+    """Build the driver and interface.
+
+    Args:
+      vid: servo device vid
+      pid: servo device pid
+      interface: which usb interface the servo console is on
+      serialname: the servo device serial (if available)
+    """
+    self._vid = vid
+    self._pid = pid
+    self._interface = interface
+    self._serial = serialname
+    self._debug = debug
+    self._init()
+
+  def _init(self):
+    self.suart = stm32uart.Suart(vendor=self._vid,
+                                 product=self._pid,
+                                 interface=self._interface,
+                                 serialname=self._serial,
+                                 debuglog=self._debug)
+    self.suart.run()
+    self.pty = pty_driver.ptyDriver(self.suart, [])
+
+  def reinitialize(self):
+    """Reinitialize the connect after a reset/disconnect/etc."""
+    self.close()
+    self._init()
+
+  def close(self):
+    """Close out the connection and release resources.
+
+    Note: if another TinyServod process or servod itself needs the same device
+          it's necessary to call this to ensure the usb device is available.
+    """
+    self.suart.close()
diff --git a/extra/usb_updater/servo_updater.py b/extra/usb_updater/servo_updater.py
index 80a63636d8..e093f55422 100755
--- a/extra/usb_updater/servo_updater.py
+++ b/extra/usb_updater/servo_updater.py
@@ -19,7 +19,7 @@ import json
 
 import fw_update
 import ecusb.tiny_servo_common as c
-
+from ecusb import tiny_servod
 
 class ServoUpdaterException(Exception):
   """Raised on exceptions generated by servo_updater."""
@@ -82,48 +82,25 @@ def flash2(vidpid, serialno, binfile):
   else:
     raise ServoUpdaterException("%s exit with res = %d" % (cmd, res))
 
-def connect(vidpid, iface, serialno, debuglog=False):
-  """Connect to console.
-
-  Args:
-    vidpid: vidpid of desired device.
-    iface: interface to connect.
-    serialno: serial number, to differentiate multiple devices.
-    debuglog: do chatty log.
-
-  Returns:
-    a connected pty object.
-  """
-  # Make sure device is up.
-  c.wait_for_usb(vidpid, serialname=serialno)
-
-  # make a console.
-  pty = c.setup_tinyservod(vidpid, iface,
-            serialname=serialno, debuglog=debuglog)
-
-  return pty
-
-def select(vidpid, iface, serialno, region, debuglog=False):
+def select(tinys, region):
   """Ensure the servo is in the expected ro/rw partition."""
 
   if region not in ["rw", "ro"]:
     raise Exception("Region must be ro or rw")
 
-  pty = connect(vidpid, iface, serialno)
-
   if region is "ro":
     cmd = "reboot"
   else:
     cmd = "sysjump %s" % region
-  pty._issue_cmd(cmd)
-  time.sleep(1)
-  pty.close()
+  tinys.pty._issue_cmd(cmd)
+  time.sleep(2)
+  tinys.reinitialize()
 
-def do_version(vidpid, iface, serialno):
+def do_version(tinys):
   """Check version via ec console 'pty'.
 
   Args:
-    see connect()
+    tinys: TinyServod object
 
   Returns:
     detected version number
@@ -133,26 +110,23 @@ def do_version(vidpid, iface, serialno):
   # ...
   # Build:   tigertail_v1.1.6749-74d1a312e
   """
-  pty = connect(vidpid, iface, serialno)
-
   cmd = '\r\nversion\r\n'
   regex = 'Build:\s+(\S+)[\r\n]+'
 
-  results = pty._issue_cmd_get_results(cmd, [regex])[0]
-  pty.close()
+  results = tinys.pty._issue_cmd_get_results(cmd, [regex])[0]
 
   return results[1].strip(' \t\r\n\0')
 
-def do_updater_version(vidpid, iface, serialno):
+def do_updater_version(tinys):
   """Check whether this uses python updater or c++ updater
 
   Args:
-    see connect()
+    tinys: TinyServod object
 
   Returns:
     updater version number. 2 or 6.
   """
-  vers = do_version(vidpid, iface, serialno)
+  vers = do_version(tinys)
 
   # Servo versions below 58 are from servo-9040.B. Versions starting with _v2
   # are newer than anything _v1, no need to check the exact number. Updater
@@ -270,17 +244,22 @@ def main():
   brdfile, binfile = findfiles(args.board, args.file)
 
   serialno = args.serialno
-  debuglog = (args.verbose is True)
 
   with open(brdfile) as data_file:
       data = json.load(data_file)
-
-  vidpid = "%04x:%04x" % (int(data['vid'], 0), int(data['pid'], 0))
+  vid, pid = int(data['vid'], 0), int(data['pid'], 0)
+  vidpid = "%04x:%04x" % (vid, pid)
   iface = int(data['console'], 0)
   boardname = data['board']
 
+
+  # Make sure device is up.
+  c.wait_for_usb(vidpid, serialname=serialno)
+  # We need a tiny_servod to query some information. Set it up first.
+  tinys = tiny_servod.TinyServod(vid, pid, iface, serialno, args.verbose)
+
   if not args.force:
-    vers = do_version(vidpid, iface, serialno)
+    vers = do_version(tinys)
     print("Current %s version is   %s" % (boardname, vers))
 
     newvers = find_available_version(boardname, binfile)
@@ -288,16 +267,20 @@ def main():
 
     if newvers == vers:
       print("No version update needed")
-      if args.reboot is True:
-        select(vidpid, iface, serialno, "ro", debuglog=debuglog)
+      if args.reboot:
+        select(tinys, 'ro')
       return
     else:
       print("Updating to recommended version.")
 
 
-  select(vidpid, iface, serialno, "ro", debuglog=debuglog)
+  # Make sure the servo MCU is in RO
+  select(tinys, 'ro')
 
-  vers = do_updater_version(vidpid, iface, serialno)
+  vers = do_updater_version(tinys)
+  # To make sure that the tiny_servod here does not interfere with other
+  # processes, close it out.
+  tinys.close()
   if vers == 2:
     flash(brdfile, serialno, binfile)
   elif vers == 6:
@@ -305,9 +288,15 @@ def main():
   else:
     raise ServoUpdaterException("Can't detect updater version")
 
-  select(vidpid, iface, serialno, "rw", debuglog=debuglog)
+  # Make sure device is up.
+  c.wait_for_usb(vidpid, serialname=serialno)
+  # After we have made sure that it's back/available, reconnect the tiny servod.
+  tinys.reinitialize()
+
+  # Make sure the servo MCU is in RW
+  select(tinys, 'rw')
 
-  vers = do_updater_version(vidpid, iface, serialno)
+  vers = do_updater_version(tinys)
   if vers == 2:
     flash(brdfile, serialno, binfile)
   elif vers == 6:
@@ -315,7 +304,8 @@ def main():
   else:
     raise ServoUpdaterException("Can't detect updater version")
 
-  select(vidpid, iface, serialno, "ro", debuglog=debuglog)
+  # Make sure the servo MCU is in RO
+  select(tinys, 'ro')
 
 if __name__ == "__main__":
   main()