From 3115967bb69499e4e24f3439fe1231d3fea1d95f Mon Sep 17 00:00:00 2001
From: mndza <diego.hdmp@gmail.com>
Date: Tue, 16 Jul 2024 16:51:24 +0200
Subject: [PATCH] apollo_fpga: Add optional timeout to
 ApolloDebugger._find_device

As PyUSB does not provide a way to wait for a device, we simulate it
with periodic sleeps. If we move to libusb1 at some point we can get rid
of this.
---
 apollo_fpga/__init__.py              | 22 ++++++++++++++--------
 apollo_fpga/commands/cli.py          |  3 ---
 apollo_fpga/gateware/flash_bridge.py |  3 ++-
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/apollo_fpga/__init__.py b/apollo_fpga/__init__.py
index bfe2061..60d71c5 100644
--- a/apollo_fpga/__init__.py
+++ b/apollo_fpga/__init__.py
@@ -111,8 +111,7 @@ def __init__(self, force_offline=False):
                 raise DebuggerNotFound(f"Handoff request failed: {e.strerror}")
 
             # Wait for Apollo to enumerate and try again
-            time.sleep(2)
-            device = self._find_device(self.APOLLO_USB_IDS, custom_match=self._device_has_apollo_id)
+            device = self._find_device(self.APOLLO_USB_IDS, custom_match=self._device_has_apollo_id, timeout=5000)
             if device is None:
                 raise DebuggerNotFound("Handoff was requested, but Apollo is not available")
 
@@ -149,7 +148,7 @@ def _request_handoff(cls, device):
         device.ctrl_transfer(request_type, REQUEST_APOLLO_ADV_STOP, wIndex=intf_number, timeout=5000)
 
     @staticmethod
-    def _find_device(ids, custom_match=None):
+    def _find_device(ids, custom_match=None, timeout=0):
         import usb.backend.libusb1
 
         # In Windows, we need to specify the libusb library location to create a backend.
@@ -168,11 +167,18 @@ def _find_device(ids, custom_match=None):
             backend = usb.backend.libusb1.get_backend()
 
         # Find the device.
-        for vid, pid in ids:
-            device = usb.core.find(backend=backend, idVendor=vid, idProduct=pid, custom_match=custom_match)
-            if device is not None:
-                return device
-
+        wait = 0
+        while True:
+            for vid, pid in ids:
+                device = usb.core.find(backend=backend, idVendor=vid, idProduct=pid, custom_match=custom_match)
+                if device is not None:
+                    return device
+            # Should we wait and try again?
+            if wait >= timeout:
+                break
+            time.sleep(0.5)
+            wait += 500
+            
         return None
 
     @staticmethod
diff --git a/apollo_fpga/commands/cli.py b/apollo_fpga/commands/cli.py
index 538b18a..8dfa510 100755
--- a/apollo_fpga/commands/cli.py
+++ b/apollo_fpga/commands/cli.py
@@ -204,9 +204,6 @@ def program_flash_fast(device, args):
     # Let the LUNA gateware take over in devices with shared USB port
     device.allow_fpga_takeover_usb()
 
-    # Wait for flash bridge enumeration
-    time.sleep(2)
-
     # Program SPI flash memory using the configured bridge
     bridge = FlashBridgeConnection()
     programmer = ECP5FlashBridgeProgrammer(bridge=bridge)
diff --git a/apollo_fpga/gateware/flash_bridge.py b/apollo_fpga/gateware/flash_bridge.py
index b7f2723..013d066 100644
--- a/apollo_fpga/gateware/flash_bridge.py
+++ b/apollo_fpga/gateware/flash_bridge.py
@@ -356,7 +356,8 @@ def __init__(self):
         # Try to create a connection to our configuration flash bridge.
         device = ApolloDebugger._find_device(
             ids=[(VENDOR_ID, PRODUCT_ID)],
-            custom_match=self._find_cfg_flash_bridge
+            custom_match=self._find_cfg_flash_bridge,
+            timeout=5000
         )
 
         # If we couldn't find the bridge, bail out.