diff --git a/bellows/multicast.py b/bellows/multicast.py index a1fe3634..80208491 100644 --- a/bellows/multicast.py +++ b/bellows/multicast.py @@ -1,3 +1,4 @@ +import asyncio import logging from bellows import types as t @@ -14,6 +15,9 @@ def __init__(self, ezsp): tuple[int, int], tuple[t.EmberMulticastTableEntry, int] ] = {} self._available = set() + # The cache is only updated after the firmware write completes, so concurrent + # calls for the same group would otherwise each claim a table slot + self._lock = asyncio.Lock() async def _initialize(self) -> None: self._multicast = {} @@ -77,6 +81,10 @@ async def _set_multicast_entry( return status, entry async def subscribe(self, group_id: int, endpoint_id: int = 1) -> t.sl_Status: + async with self._lock: + return await self._subscribe(group_id, endpoint_id) + + async def _subscribe(self, group_id: int, endpoint_id: int) -> t.sl_Status: if (group_id, endpoint_id) in self._multicast: LOGGER.debug("%s is already subscribed", t.EmberMulticastId(group_id)) return t.sl_Status.OK @@ -99,6 +107,10 @@ async def subscribe(self, group_id: int, endpoint_id: int = 1) -> t.sl_Status: return status async def unsubscribe(self, group_id: int, endpoint_id: int = 1) -> t.sl_Status: + async with self._lock: + return await self._unsubscribe(group_id, endpoint_id) + + async def _unsubscribe(self, group_id: int, endpoint_id: int) -> t.sl_Status: try: _entry, idx = self._multicast[group_id, endpoint_id] except KeyError: diff --git a/tests/test_multicast.py b/tests/test_multicast.py index ef44bf6d..6b422374 100644 --- a/tests/test_multicast.py +++ b/tests/test_multicast.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock, MagicMock, sentinel import pytest @@ -192,3 +193,44 @@ async def test_unsubscribe_fail(multicast): assert set_entry.call_count == 1 assert (grp_id, 1) in multicast._multicast assert len(multicast._available) == 0 + + +async def test_subscribe_concurrent_same_group(multicast): + """Concurrent subscriptions to one group must consume a single table slot.""" + grp_id = 0x0200 + multicast._available.update({1, 2}) + + async def mock_set(*args): + await asyncio.sleep(0) + return [t.EmberStatus.SUCCESS] + + multicast._ezsp.setMulticastTableEntry = MagicMock(side_effect=mock_set) + + results = await asyncio.gather( + multicast.subscribe(grp_id), multicast.subscribe(grp_id) + ) + + assert results == [t.sl_Status.OK, t.sl_Status.OK] + assert multicast._ezsp.setMulticastTableEntry.call_count == 1 + assert len(multicast._available) == 1 + + # The cached slot is the one in use, so unsubscribing frees it + await multicast.unsubscribe(grp_id) + assert len(multicast._available) == 2 + + +async def test_subscribe_unsubscribe_serialized(multicast): + """An unsubscribe issued while a subscribe is in flight waits for it.""" + grp_id = 0x0200 + multicast._available.add(1) + + async def mock_set(*args): + await asyncio.sleep(0) + return [t.EmberStatus.SUCCESS] + + multicast._ezsp.setMulticastTableEntry = MagicMock(side_effect=mock_set) + + await asyncio.gather(multicast.subscribe(grp_id), multicast.unsubscribe(grp_id)) + + assert (grp_id, 1) not in multicast._multicast + assert multicast._available == {1}