Skip to content

Commit ecaca06

Browse files
authored
Enhancements from the last PR (#114)
- **Use shorter variant of the same code** - **Clean up retry-tasks on shutdown**
2 parents f316eb2 + 2350b4c commit ecaca06

File tree

1 file changed

+15
-26
lines changed

1 file changed

+15
-26
lines changed

src/frequenz/dispatch/_actor_dispatcher.py

+15-26
Original file line numberDiff line numberDiff line change
@@ -133,19 +133,25 @@ async def main():
133133
```
134134
"""
135135

136-
class RetryFailedDispatches:
137-
"""Manages the retry of failed dispatches."""
136+
class FailedDispatchesRetrier(BackgroundService):
137+
"""Manages the retring of failed dispatches."""
138138

139139
def __init__(self, retry_interval: timedelta) -> None:
140140
"""Initialize the retry manager.
141141
142142
Args:
143143
retry_interval: The interval between retries.
144144
"""
145+
super().__init__()
145146
self._retry_interval = retry_interval
146147
self._channel = Broadcast[Dispatch](name="retry_channel")
147148
self._sender = self._channel.new_sender()
148-
self._tasks: set[asyncio.Task[None]] = set()
149+
150+
def start(self) -> None:
151+
"""Start the background service.
152+
153+
This is a no-op.
154+
"""
149155

150156
def new_receiver(self) -> Receiver[Dispatch]:
151157
"""Create a new receiver for dispatches to retry.
@@ -187,7 +193,7 @@ def __init__( # pylint: disable=too-many-arguments, too-many-positional-argumen
187193
],
188194
running_status_receiver: Receiver[Dispatch],
189195
dispatch_identity: Callable[[Dispatch], int] | None = None,
190-
retry_interval: timedelta | None = timedelta(seconds=60),
196+
retry_interval: timedelta = timedelta(seconds=60),
191197
) -> None:
192198
"""Initialize the dispatch handler.
193199
@@ -197,7 +203,7 @@ def __init__( # pylint: disable=too-many-arguments, too-many-positional-argumen
197203
running_status_receiver: The receiver for dispatch running status changes.
198204
dispatch_identity: A function to identify to which actor a dispatch refers.
199205
By default, it uses the dispatch ID.
200-
retry_interval: The interval between retries. If `None`, retries are disabled.
206+
retry_interval: The interval between retries.
201207
"""
202208
super().__init__()
203209
self._dispatch_identity: Callable[[Dispatch], int] = (
@@ -211,11 +217,7 @@ def __init__( # pylint: disable=too-many-arguments, too-many-positional-argumen
211217
name="dispatch_updates_channel", resend_latest=True
212218
)
213219
self._updates_sender = self._updates_channel.new_sender()
214-
self._retrier = (
215-
ActorDispatcher.RetryFailedDispatches(retry_interval)
216-
if retry_interval
217-
else None
218-
)
220+
self._retrier = ActorDispatcher.FailedDispatchesRetrier(retry_interval)
219221

220222
def start(self) -> None:
221223
"""Start the background service."""
@@ -258,12 +260,7 @@ async def _start_actor(self, dispatch: Dispatch) -> None:
258260
dispatch.type,
259261
exc_info=e,
260262
)
261-
if self._retrier:
262-
self._retrier.retry(dispatch)
263-
else:
264-
_logger.error(
265-
"No retry mechanism enabled, dispatch %r failed", dispatch
266-
)
263+
self._retrier.retry(dispatch)
267264
else:
268265
# No exception occurred, so we can add the actor to the list
269266
self._actors[identity] = actor
@@ -275,26 +272,18 @@ async def _stop_actor(self, stopping_dispatch: Dispatch, msg: str) -> None:
275272
stopping_dispatch: The dispatch that is stopping the actor.
276273
msg: The message to be passed to the actors being stopped.
277274
"""
278-
actor: Actor | None = None
279275
identity = self._dispatch_identity(stopping_dispatch)
280276

281-
actor = self._actors.get(identity)
282-
283-
if actor:
277+
if actor := self._actors.pop(identity, None):
284278
await actor.stop(msg)
285-
286-
del self._actors[identity]
287279
else:
288280
_logger.warning(
289281
"Actor for dispatch type %r is not running", stopping_dispatch.type
290282
)
291283

292284
async def _run(self) -> None:
293285
"""Run the background service."""
294-
if not self._retrier:
295-
async for dispatch in self._dispatch_rx:
296-
await self._handle_dispatch(dispatch)
297-
else:
286+
async with self._retrier:
298287
retry_recv = self._retrier.new_receiver()
299288

300289
async for selected in select(retry_recv, self._dispatch_rx):

0 commit comments

Comments
 (0)