7
7
8
8
import logging
9
9
from asyncio import Event
10
- from typing import Any , Callable , Generator
10
+ from typing import Callable
11
11
12
12
from frequenz .channels import Receiver
13
13
from frequenz .client .dispatch import Client
14
- from frequenz .sdk .actor import Actor
14
+ from frequenz .sdk .actor import Actor , BackgroundService
15
+ from typing_extensions import override
15
16
16
17
from ._actor_dispatcher import ActorDispatcher , DispatchInfo
17
18
from ._bg_service import DispatchScheduler , MergeStrategy
22
23
_logger = logging .getLogger (__name__ )
23
24
24
25
25
- class Dispatcher :
26
+ class Dispatcher ( BackgroundService ) :
26
27
"""A highlevel interface for the dispatch API.
27
28
28
29
This class provides a highlevel interface to the dispatch API.
@@ -183,6 +184,8 @@ def __init__(
183
184
server_url: The URL of the dispatch service.
184
185
key: The key to access the service.
185
186
"""
187
+ super ().__init__ (name = "Dispatcher" )
188
+
186
189
self ._client = Client (server_url = server_url , key = key )
187
190
self ._bg_service = DispatchScheduler (
188
191
microgrid_id ,
@@ -192,10 +195,32 @@ def __init__(
192
195
self ._empty_event = Event ()
193
196
self ._empty_event .set ()
194
197
198
+ @override
195
199
def start (self ) -> None :
196
200
"""Start the local dispatch service."""
197
201
self ._bg_service .start ()
198
202
203
+ @property
204
+ @override
205
+ def is_running (self ) -> bool :
206
+ """Whether the local dispatch service is running."""
207
+ return self ._bg_service .is_running
208
+
209
+ @override
210
+ async def wait (self ) -> None :
211
+ """Wait until all actor dispatches are stopped."""
212
+ await self ._empty_event .wait ()
213
+
214
+ @override
215
+ def cancel (self , msg : str | None = None ) -> None :
216
+ """Stop the local dispatch service."""
217
+ self ._bg_service .cancel (msg )
218
+
219
+ for instance in self ._actor_dispatchers .values ():
220
+ instance .cancel ()
221
+
222
+ self ._actor_dispatchers .clear ()
223
+
199
224
async def start_dispatching (
200
225
self ,
201
226
dispatch_type : str ,
@@ -256,10 +281,6 @@ async def stop_dispatching(self, dispatch_type: str) -> None:
256
281
if not self ._actor_dispatchers :
257
282
self ._empty_event .set ()
258
283
259
- def __await__ (self ) -> Generator [Any , None , bool ]:
260
- """Wait until all actor dispatches are stopped."""
261
- return self ._empty_event .wait ().__await__ ()
262
-
263
284
@property
264
285
def client (self ) -> Client :
265
286
"""Return the client."""
0 commit comments