Merge branch 'main' into agents/dps-group-registration-timeout

Author: cartertinney

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,25 @@
     "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()
 
 
 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 +49,29 @@ 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:
+            # 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)
+                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")

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"
     )

tests/unit/iothub/aio/test_loop_management.py

@@ -6,8 +6,11 @@
 
 import pytest
 import asyncio
+import concurrent.futures
 import logging
+import threading
 from azure.iot.device.iothub.aio import loop_management
+from tests.unit.helpers import BATCH_COMPLETION_TIMEOUT
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -49,6 +52,40 @@ 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):
+        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):
+            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
+        )
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+            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
+        assert returned_loops[0] is returned_loops[1]
+
 
 @pytest.mark.describe(".get_client_internal_loop()")
 class TestGetClientInternalLoop(SharedCustomLoopTests):