From d3b8a698c2e9ad9ee175e623caafcdd91ac6547b Mon Sep 17 00:00:00 2001 From: Johnny Deuss Date: Wed, 22 Nov 2023 10:39:32 +0000 Subject: [PATCH] Clean up channel subscription on exception --- broadcaster/_base.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/broadcaster/_base.py b/broadcaster/_base.py index 4de1417..8b64bef 100644 --- a/broadcaster/_base.py +++ b/broadcaster/_base.py @@ -88,13 +88,13 @@ async def subscribe(self, channel: str) -> AsyncIterator["Subscriber"]: self._subscribers[channel] = set([queue]) else: self._subscribers[channel].add(queue) - - yield Subscriber(queue) - - self._subscribers[channel].remove(queue) - if not self._subscribers.get(channel): - del self._subscribers[channel] - await self._backend.unsubscribe(channel) + try: + yield Subscriber(queue) + finally: + self._subscribers[channel].remove(queue) + if not self._subscribers.get(channel): + del self._subscribers[channel] + await self._backend.unsubscribe(channel) finally: await queue.put(None)