From 7de3cc078b66633b571756fd5a46a79e3d003a90 Mon Sep 17 00:00:00 2001 From: Carter Tinney Date: Wed, 2 Sep 2026 08:22:31 -0700 Subject: [PATCH 1/3] fix: make async loop creation thread-safe Serialize lazy creation of the shared async client loops so concurrent first access cannot publish different event loops. This preserves Janus queue affinity and prevents handler runners from entering a permanent restart loop. Add concurrent initialization and queue-affinity regression coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../iot/device/iothub/aio/loop_management.py | 44 +++++++++++-------- tests/unit/iothub/aio/test_async_inbox.py | 8 ++++ tests/unit/iothub/aio/test_loop_management.py | 28 ++++++++++++ 3 files changed, 62 insertions(+), 18 deletions(-) diff --git a/azure-iot-device/azure/iot/device/iothub/aio/loop_management.py b/azure-iot-device/azure/iot/device/iothub/aio/loop_management.py index f2b73a9c2..c0dd31930 100644 --- a/azure-iot-device/azure/iot/device/iothub/aio/loop_management.py +++ b/azure-iot-device/azure/iot/device/iothub/aio/loop_management.py @@ -3,8 +3,8 @@ # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- -""" This module contains functions of managing event loops for the IoTHub client -""" +"""This module contains functions of managing event loops for the IoTHub client""" + import asyncio import threading import logging @@ -16,21 +16,24 @@ "CLIENT_INTERNAL_LOOP": None, "CLIENT_HANDLER_RUNNER_LOOP": None, } +_loop_creation_lock = threading.Lock() def _cleanup(): """Clear all running loops and end respective threads. ONLY FOR TESTING USAGE By using this function, you can wipe all global loops. + Do not call while clients or inboxes are still in use. DO NOT USE THIS IN PRODUCTION CODE """ - for loop_name, loop in loops.items(): - if loop is not None: - logger.debug("Stopping event loop - {}".format(loop_name)) - loop.call_soon_threadsafe(loop.stop) - # NOTE: Stopping the loop will also end the thread, because the only thing keeping - # the thread alive was the loop running - loops[loop_name] = None + with _loop_creation_lock: + for loop_name, loop in loops.items(): + if loop is not None: + logger.debug("Stopping event loop - {}".format(loop_name)) + loop.call_soon_threadsafe(loop.stop) + # NOTE: Stopping the loop will also end the thread, because the only thing keeping + # the thread alive was the loop running + loops[loop_name] = None def _make_new_loop(loop_name): @@ -45,23 +48,28 @@ def _make_new_loop(loop_name): loops[loop_name] = new_loop +def _get_or_create_loop(loop_name): + loop = loops[loop_name] + if loop is None: + with _loop_creation_lock: + loop = loops[loop_name] + if loop is None: + _make_new_loop(loop_name) + loop = loops[loop_name] + return loop + + def get_client_internal_loop(): """Return the loop for internal client operations""" - if loops["CLIENT_INTERNAL_LOOP"] is None: - _make_new_loop("CLIENT_INTERNAL_LOOP") - return loops["CLIENT_INTERNAL_LOOP"] + return _get_or_create_loop("CLIENT_INTERNAL_LOOP") def get_client_handler_runner_loop(): """Return the loop for handler runners""" - if loops["CLIENT_HANDLER_RUNNER_LOOP"] is None: - _make_new_loop("CLIENT_HANDLER_RUNNER_LOOP") - return loops["CLIENT_HANDLER_RUNNER_LOOP"] + return _get_or_create_loop("CLIENT_HANDLER_RUNNER_LOOP") def get_client_handler_loop(): """Return the loop for invoking user-provided handlers on the client""" # TODO: Try and store the user loop somehow - if loops["CLIENT_HANDLER_LOOP"] is None: - _make_new_loop("CLIENT_HANDLER_LOOP") - return loops["CLIENT_HANDLER_LOOP"] + return _get_or_create_loop("CLIENT_HANDLER_LOOP") diff --git a/tests/unit/iothub/aio/test_async_inbox.py b/tests/unit/iothub/aio/test_async_inbox.py index 1c0236bf9..246f49793 100644 --- a/tests/unit/iothub/aio/test_async_inbox.py +++ b/tests/unit/iothub/aio/test_async_inbox.py @@ -105,6 +105,14 @@ async def test_removes_item_from_inbox_if_already_there(self, mocker, inbox): assert retrieved_item is item assert inbox.empty() + @pytest.mark.it("Runs Janus async operations on the shared internal loop") + async def test_uses_shared_internal_loop(self, mocker, inbox): + inbox.put(mocker.MagicMock()) + + await asyncio.wait_for(inbox.get(), timeout=PROMPT_TIMEOUT) + + assert inbox._queue._loop is loop_management.get_client_internal_loop() + @pytest.mark.it( "Blocks on an empty inbox until an item is available to remove and return, if using blocking mode" ) diff --git a/tests/unit/iothub/aio/test_loop_management.py b/tests/unit/iothub/aio/test_loop_management.py index 2d7af158a..84f334630 100644 --- a/tests/unit/iothub/aio/test_loop_management.py +++ b/tests/unit/iothub/aio/test_loop_management.py @@ -6,8 +6,12 @@ import pytest import asyncio +import concurrent.futures import logging +import threading +import time from azure.iot.device.iothub.aio import loop_management +from tests.unit.helpers import BATCH_COMPLETION_TIMEOUT logging.basicConfig(level=logging.DEBUG) @@ -49,6 +53,30 @@ def test_same_loop(self, fn_under_test): loop2 = fn_under_test() assert loop1 is loop2 + @pytest.mark.it("Creates only one event loop when first called concurrently") + def test_threadsafe_first_call(self, mocker, fn_under_test): + start_barrier = threading.Barrier(3) + + def make_loop(loop_name): + time.sleep(0.05) + loop_management.loops[loop_name] = mocker.MagicMock() + + make_loop_mock = mocker.patch.object( + loop_management, "_make_new_loop", side_effect=make_loop + ) + + def get_loop(): + start_barrier.wait(timeout=BATCH_COMPLETION_TIMEOUT) + return fn_under_test() + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(get_loop) for _ in range(2)] + start_barrier.wait(timeout=BATCH_COMPLETION_TIMEOUT) + returned_loops = [future.result(timeout=BATCH_COMPLETION_TIMEOUT) for future in futures] + + assert make_loop_mock.call_count == 1 + assert returned_loops[0] is returned_loops[1] + @pytest.mark.describe(".get_client_internal_loop()") class TestGetClientInternalLoop(SharedCustomLoopTests): From 3f246fb03bd559828b93e82d16b0de7bddd942e7 Mon Sep 17 00:00:00 2001 From: Carter Tinney Date: Wed, 2 Sep 2026 08:36:39 -0700 Subject: [PATCH 2/3] docs: explain async loop creation lock Document the Janus loop-affinity invariant and why loop state must be checked again after acquiring the creation lock. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- azure-iot-device/azure/iot/device/iothub/aio/loop_management.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/azure-iot-device/azure/iot/device/iothub/aio/loop_management.py b/azure-iot-device/azure/iot/device/iothub/aio/loop_management.py index c0dd31930..e5a73af3d 100644 --- a/azure-iot-device/azure/iot/device/iothub/aio/loop_management.py +++ b/azure-iot-device/azure/iot/device/iothub/aio/loop_management.py @@ -16,6 +16,7 @@ "CLIENT_INTERNAL_LOOP": None, "CLIENT_HANDLER_RUNNER_LOOP": None, } +# Janus queues bind to the first loop they use, so concurrent callers must receive the same loop. _loop_creation_lock = threading.Lock() @@ -52,6 +53,7 @@ def _get_or_create_loop(loop_name): loop = loops[loop_name] if loop is None: with _loop_creation_lock: + # Another caller may have created the loop while this caller waited for the lock. loop = loops[loop_name] if loop is None: _make_new_loop(loop_name) From bd536009dd839c92c6a33e20a76d3c29e3944c20 Mon Sep 17 00:00:00 2001 From: Carter Tinney Date: Wed, 2 Sep 2026 10:07:56 -0700 Subject: [PATCH 3/3] tests: make async loop race regression deterministic Coordinate the first two loop-map reads so both worker threads observe the uninitialized state before either can publish a loop. This removes the scheduler-dependent sleep and guarantees the old implementation fails the test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/unit/iothub/aio/test_loop_management.py | 27 ++++++++++++------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/tests/unit/iothub/aio/test_loop_management.py b/tests/unit/iothub/aio/test_loop_management.py index 84f334630..f20d6f258 100644 --- a/tests/unit/iothub/aio/test_loop_management.py +++ b/tests/unit/iothub/aio/test_loop_management.py @@ -9,7 +9,6 @@ import concurrent.futures import logging import threading -import time from azure.iot.device.iothub.aio import loop_management from tests.unit.helpers import BATCH_COMPLETION_TIMEOUT @@ -55,23 +54,33 @@ def test_same_loop(self, fn_under_test): @pytest.mark.it("Creates only one event loop when first called concurrently") def test_threadsafe_first_call(self, mocker, fn_under_test): - start_barrier = threading.Barrier(3) + class CoordinatedLoopMap(dict): + def __init__(self, loops): + super().__init__(loops) + self._read_barrier = threading.Barrier(2) + self._read_lock = threading.Lock() + self._reads_to_coordinate = 2 + + def __getitem__(self, loop_name): + loop = super().__getitem__(loop_name) + with self._read_lock: + coordinate_read = self._reads_to_coordinate > 0 + if coordinate_read: + self._reads_to_coordinate -= 1 + if coordinate_read: + self._read_barrier.wait(timeout=BATCH_COMPLETION_TIMEOUT) + return loop def make_loop(loop_name): - time.sleep(0.05) loop_management.loops[loop_name] = mocker.MagicMock() + mocker.patch.object(loop_management, "loops", CoordinatedLoopMap(loop_management.loops)) make_loop_mock = mocker.patch.object( loop_management, "_make_new_loop", side_effect=make_loop ) - def get_loop(): - start_barrier.wait(timeout=BATCH_COMPLETION_TIMEOUT) - return fn_under_test() - with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: - futures = [executor.submit(get_loop) for _ in range(2)] - start_barrier.wait(timeout=BATCH_COMPLETION_TIMEOUT) + futures = [executor.submit(fn_under_test) for _ in range(2)] returned_loops = [future.result(timeout=BATCH_COMPLETION_TIMEOUT) for future in futures] assert make_loop_mock.call_count == 1