From 87460b99c9fb0a207fceda100dde5ded90ea117b Mon Sep 17 00:00:00 2001
From: David Lechner <david@pybricks.com>
Date: Fri, 25 Aug 2023 12:22:21 -0500
Subject: [PATCH] backends/bluezdbus/manager: race InterfaceRemoved in
 _wait_for_services_discovery()

This adds a 3rd race condition to _wait_for_services_discovery() to
handle the case where an interface is removed from BlueZ without
changing the "Connected" property. This can happen, e.g. when we
hit "retry due to le-connection-abort-by-local".
---
 CHANGELOG.rst                       |  1 +
 bleak/backends/bluezdbus/manager.py | 73 +++++++++++++++++++++++------
 2 files changed, 60 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 30afd8a4..ced4d631 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -32,6 +32,7 @@ Fixed
 * Fixed typing for ``BaseBleakScanner`` detection callback.
 * Fixed possible crash in ``_stopped_handler()`` in WinRT backend. Fixes #1330.
 * Reduced expensive logging in the BlueZ backend. Merged #1376.
+* Fixed race condition with ``"InterfaceRemoved"`` when getting services in BlueZ backend.
 
 `0.20.2`_ (2023-04-19)
 ======================
diff --git a/bleak/backends/bluezdbus/manager.py b/bleak/backends/bluezdbus/manager.py
index 85767987..c2f34a74 100644
--- a/bleak/backends/bluezdbus/manager.py
+++ b/bleak/backends/bluezdbus/manager.py
@@ -7,6 +7,7 @@
 """
 
 import asyncio
+import contextlib
 import logging
 import os
 from typing import (
@@ -692,22 +693,66 @@ async def _wait_for_services_discovery(self, device_path: str) -> None:
 
         If a disconnect happens before the completion a BleakError exception is raised.
         """
-        services_discovered_wait_task = asyncio.create_task(
-            self._wait_condition(device_path, "ServicesResolved", True)
-        )
-        device_disconnected_wait_task = asyncio.create_task(
-            self._wait_condition(device_path, "Connected", False)
-        )
-        done, pending = await asyncio.wait(
-            {services_discovered_wait_task, device_disconnected_wait_task},
-            return_when=asyncio.FIRST_COMPLETED,
-        )
+        with contextlib.ExitStack() as stack:
+            services_discovered_wait_task = asyncio.create_task(
+                self._wait_condition(device_path, "ServicesResolved", True)
+            )
+            stack.callback(services_discovered_wait_task.cancel)
+
+            device_disconnected_wait_task = asyncio.create_task(
+                self._wait_condition(device_path, "Connected", False)
+            )
+            stack.callback(device_disconnected_wait_task.cancel)
+
+            # in some cases, we can get "InterfaceRemoved" without the
+            # "Connected" property changing, so we need to race against both
+            # conditions
+            device_removed_wait_task = asyncio.create_task(
+                self._wait_removed(device_path)
+            )
+            stack.callback(device_removed_wait_task.cancel)
+
+            done, _ = await asyncio.wait(
+                {
+                    services_discovered_wait_task,
+                    device_disconnected_wait_task,
+                    device_removed_wait_task,
+                },
+                return_when=asyncio.FIRST_COMPLETED,
+            )
 
-        for p in pending:
-            p.cancel()
+            if not done.isdisjoint(
+                {device_disconnected_wait_task, device_removed_wait_task}
+            ):
+                raise BleakError("failed to discover services, device disconnected")
 
-        if device_disconnected_wait_task in done:
-            raise BleakError("failed to discover services, device disconnected")
+    async def _wait_removed(self, device_path: str) -> None:
+        """
+        Waits for the device interface to be removed.
+
+        If the device is not present in BlueZ, this returns immediately.
+
+        Args:
+            device_path: The D-Bus object path of a Bluetooth device.
+        """
+        if device_path not in self._properties:
+            return
+
+        event = asyncio.Event()
+
+        def callback(_: str):
+            event.set()
+
+        device_removed_callback_and_state = DeviceRemovedCallbackAndState(
+            callback, self._properties[device_path][defs.DEVICE_INTERFACE]["Adapter"]
+        )
+
+        with contextlib.ExitStack() as stack:
+            self._device_removed_callbacks.append(device_removed_callback_and_state)
+            stack.callback(
+                self._device_removed_callbacks.remove, device_removed_callback_and_state
+            )
+            await event.wait()
 
     async def _wait_condition(
         self, device_path: str, property_name: str, property_value: Any