From 3d3a2e096b7e53ea94b28eb487d7e01ac7229473 Mon Sep 17 00:00:00 2001
From: Simon Holesch <simon@holesch.de>
Date: Sun, 18 Aug 2024 13:02:51 +0200
Subject: [PATCH] Agent: Remove USB Detach Delay

This delay was introduced for the `not-my-board edit` command, where
devices get detached and immediately attached again. Instead of delaying
detaching, just retry the first attach.
---
 not_my_board/_agent.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/not_my_board/_agent.py b/not_my_board/_agent.py
index 2ae3dfd..f97ee7c 100644
--- a/not_my_board/_agent.py
+++ b/not_my_board/_agent.py
@@ -98,10 +98,6 @@ async def usbip_attach(self, proxy, target, port_num, usbid):
     @staticmethod
     async def usbip_detach(vhci_port):
         usbip.detach(vhci_port)
-        # Unfortunately it takes ~ 0.5 seconds for the connection to close and
-        # for the remote device to be available again. Wait a bit, so an
-        # immediate attach after the detach succeeds.
-        await asyncio.sleep(2)
 
     async def port_forward(self, ready_event, proxy, target, local_port):
         connection_handler = functools.partial(
@@ -475,14 +471,13 @@ async def close(self):
         await super().close()
         if self._vhci_port is not None:
             await self._io.usbip_detach(self._vhci_port)
+            self._vhci_port = None
 
     async def _task_func(self):
         retry_timeout = 1
         while True:
             try:
-                self._vhci_port = await self._io.usbip_attach(
-                    self._proxy, USBIP_REMOTE, self.port_num, self.usbid
-                )
+                await self._attach()
                 logger.debug("%s: USB device attached", self.name)
                 self._ready_event.set()
                 retry_timeout = 1
@@ -491,6 +486,26 @@ async def _task_func(self):
                 await asyncio.sleep(retry_timeout)
                 retry_timeout = min(2 * retry_timeout, 30)
 
+    async def _attach(self):
+        # Only retry attach, if the tunnel was closed. This fixes an issue,
+        # when devices are detached and immediately attached again: When the
+        # connection is closed, then it takes ~ 0.5 seconds until the remote
+        # makes the device available again.
+        # When the first attach succeeds, then we send another attach request
+        # which blocks until the device is available again. This blocking is
+        # expected, so the attach timeout doesn't make sense anymore.
+        attach_timeout = 1 if self._vhci_port is None else None
+
+        while True:
+            try:
+                async with util.timeout(attach_timeout):
+                    self._vhci_port = await self._io.usbip_attach(
+                        self._proxy, USBIP_REMOTE, self.port_num, self.usbid
+                    )
+                    break
+            except TimeoutError:
+                attach_timeout = min(2 * attach_timeout, 30)
+
     def is_attached(self):
         if self._vhci_port is None:
             return False