udisks2: Be more tolerant of device-startup time

When switching an SDWire to 'host' mode, it may take time for the
partition to appear.

During that time the device (e.g. '/dev/sdo1') may not exist, so trying
to cat e.g. '/sys/call/block/sdo1' fails. Handle this by returning a
zero size in that case, as is done when the returned size is empty.

Even when the block device is present, udisks2 may take a short time to
make the device available. Handle this by retrying a few times, until
things settle.

Add the call to _wait_for_medium() in write_files() while we are here,
to match what is done in write_image(). Also fix the parameter type for
write_files() and mention the exception it may raise.

Fixes: labgrid-project#1357

Signed-off-by: Simon Glass <sjg@chromium.org>

labgrid/driver/usbstoragedriver.py

@@ -45,6 +45,7 @@ class USBStorageDriver(Driver):
     )
     WAIT_FOR_MEDIUM_TIMEOUT = 10.0 # s
     WAIT_FOR_MEDIUM_SLEEP = 0.5 # s
+    MOUNT_RETRIES = 5
 
     def __attrs_post_init__(self):
         super().__attrs_post_init__()
@@ -77,15 +78,19 @@ def write_files(self, sources, target, partition, target_is_directory=True):
         Args:
             sources (List[str]): path(s) to the file(s) to be copied to the bound USB storage
                 partition.
-            target (str): target directory or file to copy to
+            target (PurePath): target directory or file to copy to
             partition (int): mount the specified partition or None to mount the whole disk
             target_is_directory (bool): Whether target is a directory
+
+        Raises:
+            Exception if anything goes wrong
         """
+        self._wait_for_medium(partition)
 
         self._start_wrapper()
 
         self.devpath = self._get_devpath(partition)
-        mount_path = self.proxy.mount(self.devpath)
+        mount_path = self.proxy.mount(self.devpath, self.MOUNT_RETRIES)
 
         try:
             # (pathlib.PurePath(...) / "/") == "/", so we turn absolute paths into relative
@@ -222,10 +227,14 @@ def get_size(self, partition=None):
                 getting the size of the root device (defaults to None)
 
         Returns:
-            int: size in bytes
+            int: size in bytes, or 0 if the partition is not found
         """
         args = ["cat", f"/sys/class/block/{self._get_devpath(partition)[5:]}/size"]
-        size = subprocess.check_output(self.storage.command_prefix + args)
+        try:
+            size = subprocess.check_output(self.storage.command_prefix + args)
+        except subprocess.CalledProcessError:
+            # while the medium is getting ready, the file does not yet exist
+            return 0
         try:
             return int(size) * 512
         except ValueError:

labgrid/util/agents/udisks2.py

@@ -15,25 +15,32 @@ class UDisks2Device:
     def __init__(self, devpath):
         self._logger = logging.getLogger("Device: ")
         self.devpath = devpath
-        client = UDisks.Client.new_sync(None)
+        self.fs = None
+
+    def _setup(self):
+        """Try to find the devpath
 
+        Raises:
+            ValueError: no udisks2 device or no filesystem found on devpath
+        """
+        client = UDisks.Client.new_sync(None)
         manager = client.get_object_manager()
         for obj in manager.get_objects():
             block = obj.get_block()
             if not block:
                 continue
 
             device_path = block.get_cached_property("Device").get_bytestring().decode('utf-8')
-            if device_path == devpath:
+            if device_path == self.devpath:
                 self.fs = obj.get_filesystem()
                 if self.fs is None:
-                    raise ValueError(f"no filesystem found on {devpath}")
+                    raise ValueError(f"no filesystem found on {self.devpath}")
 
                 return
 
-        raise ValueError(f"No udisks2 device found for {devpath}")
+        raise ValueError(f"No udisks2 device found for {self.devpath}")
 
-    def mount(self, readonly=False):
+    def mount(self, readonly=False, retries=0):
         opts = GLib.Variant('a{sv}', {'options': GLib.Variant('s', 'ro' if readonly else 'rw')})
 
         try:
@@ -83,17 +90,39 @@ def unmount(self, lazy=False):
 
 _devs = {}
 
-def _get_udisks2_dev(devpath):
+def _get_udisks2_dev(devpath, retries):
+    """Try to get the udisks2 device
+
+    Args:
+        devpath (str): Device name
+        retries (int): Number of retries to allow
+
+    Raises:
+        ValueError: Failed to obtain the device (e.g. does not exist)
+    """
     if devpath not in _devs:
-        _devs[devpath] = UDisks2Device(devpath=devpath)
+        dev = UDisks2Device(devpath=devpath)
+        while True:
+            try:
+                dev._setup()
+                break
+            except ValueError as exc:
+                if 'No udisks2 device' not in str(exc) or not retries:
+                    raise
+                retries -= 1
+            dev._logger.warning('udisks2: Retrying %s...', devpath)
+            time.sleep(1)
+
+        # Success, so record the new device
+        _devs[devpath] = dev
     return _devs[devpath]
 
-def handle_mount(devpath):
-    dev = _get_udisks2_dev(devpath)
+def handle_mount(devpath, retries=0):
+    dev = _get_udisks2_dev(devpath, retries)
     return dev.mount()
 
 def handle_unmount(devpath, lazy=False):
-    dev = _get_udisks2_dev(devpath)
+    dev = _get_udisks2_dev(devpath, 0)
     return dev.unmount(lazy=lazy)
 
 methods = {