8000 Add support for eager tasks (#111425) · home-assistant/core@67e3569 · GitHub
[go: up one dir, main page]
Skip to content

Commit 67e3569

Browse files
bdracobdraco
authored
Add support for eager tasks (#111425)
* Add support for eager tasks python 3.12 supports eager tasks reading: https://docs.python.org/3/library/asyncio-task.html#eager-task-factory python/cpython#97696 There are lots of places were we are unlikely to suspend, but we might suspend so creating a task makes sense * reduce * revert entity * revert * coverage * coverage * coverage * coverage * fix test
1 parent 93cc6e0 commit 67e3569

File tree

8 files changed

+162
-18
lines changed

8 files changed

+162
-18
lines changed

‎homeassistant/components/websocket_api/decorators.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ def schedule_handler(
4545
hass.async_create_background_task(
4646
_handle_async_response(func, hass, connection, msg),
4747
task_name,
48+
eager_start=True,
4849
)
4950

5051
return schedule_handler

‎homeassistant/config_entries.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -915,6 +915,7 @@ def async_create_task(
915915
hass: HomeAssistant,
916916
target: Coroutine[Any, Any, _R],
917917
name: str | None = None,
918+
eager_start: bool = False,
918919
) -> asyncio.Task[_R]:
919920
"""Create a task from within the event loop.
920921
@@ -923,7 +924,7 @@ def async_create_task(
923924
target: target to call.
924925
"""
925926
task = hass.async_create_task(
926-
target, f"{name} {self.title} {self.domain} {self.entry_id}"
927+
target, f"{name} {self.title} {self.domain} {self.entry_id}", eager_start
927928
)
928929
self._tasks.add(task)
929930
task.add_done_callback(self._tasks.remove)
@@ -932,15 +933,19 @@ def async_create_task(
932933

933934
@callback
934935
def async_create_background_task(
935-
self, hass: HomeAssistant, target: Coroutine[Any, Any, _R], name: str
936+
self,
937+
hass: HomeAssistant,
938+
target: Coroutine[Any, Any, _R],
939+
name: str,
940+
eager_start: bool = False,
936941
) -> asyncio.Task[_R]:
937942
"""Create a background task tied to the config entry lifecycle.
938943
939944
Background tasks are automatically canceled when config entry is unloaded.
940945
941946
target: target to call.
942947
"""
943-
task = hass.async_create_background_task(target, name)
948+
task = hass.async_create_background_task(target, name, eager_start)
944949
self._background_tasks.add(task)
945950
task.add_done_callback(self._background_tasks.remove)
946951
return task

‎homeassistant/core.py

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@
9191
from .util import dt as dt_util, location
9292
from .util.async_ import (
9393
cancelling,
94+
create_eager_task,
9495
run_callback_threadsafe,
9596
shutdown_run_callback_threadsafe,
9697
)
@@ -622,7 +623,10 @@ def create_task(
622623

623624
@callback
624625
def async_create_task(
625-
self, target: Coroutine[Any, Any, _R], name: str | None = None
626+
self,
627+
target: Coroutine[Any, Any, _R],
628+
name: str | None = None,
629+
eager_start: bool = False,
626630
) -> asyncio.Task[_R]:
627631
"""Create a task from within the event loop.
628632
@@ -631,16 +635,17 @@ def async_create_task(
631635
632636
target: target to call.
633637
"""
634-
task = self.loop.create_task(target, name=name)
638+
if eager_start:
639+
task = create_eager_task(target, name=name, loop=self.loop)
640+
else:
641+
task = self.loop.create_task(target, name=name)
635642
self._tasks.add(task)
636643
task.add_done_callback(self._tasks.remove)
637644
return task
638645

639646
@callback
640647
def async_create_background_task(
641-
self,
642-
target: Coroutine[Any, Any, _R],
643-
name: str,
648+
self, target: Coroutine[Any, Any, _R], name: str, eager_start: bool = False
644649
) -> asyncio.Task[_R]:
645650
"""Create a task from within the event loop.
646651
@@ -650,7 +655,10 @@ def async_create_background_task(
650655
651656
This method must be run in the event loop.
652657
"""
653-
task = self.loop.create_task(target, name=name)
658+
if eager_start:
659+
task = create_eager_task(target, name=name, loop=self.loop)
660+
else:
661+
task = self.loop.create_task(target, name=name)
654662
self._background_tasks.add(task)
655663
task.add_done_callback(self._background_tasks.remove)
656664
return task

‎homeassistant/util/async_.py

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
"""Asyncio utilities."""
22
from __future__ import annotations
33

4-
from asyncio import Future, Semaphore, gather, get_running_loop
5-
from asyncio.events import AbstractEventLoop
4+
from asyncio import AbstractEventLoop, Future, Semaphore, Task, gather, get_running_loop
65
from collections.abc import Awaitable, Callable
76
import concurrent.futures
87
from contextlib import suppress
98
import functools
109
import logging
10+
import sys
1111
import threading
1212
from traceback import extract_stack
1313
from typing import Any, ParamSpec, TypeVar, TypeVarTuple
@@ -23,6 +23,36 @@
2323
_P = ParamSpec("_P")
2424
_Ts = TypeVarTuple("_Ts")
2525

26+
if sys.version_info >= (3, 12, 0):
27+
28+
def create_eager_task(
29+
coro: Awaitable[_T],
30+
*,
31+
name: str | None = None,
32+
loop: AbstractEventLoop | None = None,
33+
) -> Task[_T]:
34+
"""Create a task from a coroutine and schedule it to run immediately."""
35+
return Task(
36+
coro,
37+
loop=loop or get_running_loop(),
38+
name=name,
39+
eager_start=True, # type: ignore[call-arg]
40+
)
41+
else:
42+
43+
def create_eager_task(
44+
coro: Awaitable[_T],
45+
*,
46+
name: str | None = None,
47+
loop: AbstractEventLoop | None = None,
48+
) -> Task[_T]:
49+
"""Create a task from a coroutine and schedule it to run immediately."""
50+
return Task(
51+
coro,
52+
loop=loop or get_running_loop(),
53+
name=name,
54+
)
55+
2656

2757
def cancelling(task: Future[Any]) -> bool:
2858
"""Return True if task is cancelling."""

‎tests/common.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -260,14 +260,14 @@ def async_add_executor_job(target, *args):
260260

261261
return orig_async_add_executor_job(target, *args)
262262

263-
def async_create_task(coroutine, name=None):
263+
def async_create_task(coroutine, name=None, eager_start=False):
264264
"""Create task."""
265265
if isinstance(coroutine, Mock) and not isinstance(coroutine, AsyncMock):
266266
fut = asyncio.Future()
267267
fut.set_result(None)
268268
return fut
269269

270-
return orig_async_create_task(coroutine, name)
270+
return orig_async_create_task(coroutine, name, eager_start)
271271

272272
hass.async_add_job = async_add_job
273273
hass.async_add_executor_job = async_add_executor_job

‎tests/test_config_entries.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4228,11 +4228,16 @@ async def test_unload() -> None:
42284228

42294229
entry.async_on_unload(test_unload)
42304230
entry.async_create_task(hass, test_task())
4231-
entry.async_create_background_task(hass, test_task(), "background-task-name")
4231+
entry.async_create_background_task(
4232+
hass, test_task(), "background-task-name", eager_start=True
4233+
)
4234+
entry.async_create_background_task(
4235+
hass, test_task(), "background-task-name", eager_start=False
4236+
)
42324237
await asyncio.sleep(0)
42334238
hass.loop.call_soon(event.set)
42344239
await entry._async_process_on_unload(hass)
4235-
assert results == ["on_unload", "background", "normal"]
4240+
assert results == ["on_unload", "background", "background", "normal"]
42364241

42374242

42384243
async def test_preview_supported(

‎tests/test_core.py

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import gc
99
import logging
1010
import os
11+
import sys
1112
from tempfile import TemporaryDirectory
1213
import threading
1314
import time
@@ -161,7 +162,9 @@ def job():
161162
assert len(hass.loop.run_in_executor.mock_calls) == 2
162163

163164

164-
def test_async_create_task_schedule_coroutine(event_loop) -> None:
165+
def test_async_create_task_schedule_coroutine(
166+
event_loop: asyncio.AbstractEventLoop,
167+
) -> None:
165168
"""Test that we schedule coroutines and add jobs to the job pool."""
166169
hass = MagicMock(loop=MagicMock(wraps=event_loop))
167170

@@ -174,6 +177,44 @@ async def job():
174177
assert len(hass.add_job.mock_calls) == 0
175178

176179

180+
@pytest.mark.skipif(
181+
sys.version_info < (3, 12), reason="eager_start is only supported for Python 3.12"
182+
)
183+
def test_async_create_task_eager_start_schedule_coroutine(
184+
event_loop: asyncio.AbstractEventLoop,
185+
) -> None:
186+
"""Test that we schedule coroutines and add jobs to the job pool."""
187+
hass = MagicMock(loop=MagicMock(wraps=event_loop))
188+
189+
async def job():
190+
pass
191+
192+
ha.HomeAssistant.async_create_task(hass, job(), eager_start=True)
193+
# Should create the task directly since 3.12 supports eager_start
194+
assert len(hass.loop.create_task.mock_calls) == 0
195+
assert len(hass.add_job.mock_calls) == 0
196+
197+
198+
@pytest.mark.skipif(
199+
sys.version_info >= (3, 12), reason="eager_start is not supported on < 3.12"
200+
)
201+
def test_async_create_task_eager_start_fallback_schedule_coroutine(
202+
event_loop: asyncio.AbstractEventLoop,
203+
) -> None:
204+
"""Test that we schedule coroutines and add jobs to the job pool."""
205+
hass = MagicMock(loop=MagicMock(wraps=event_loop))
206+
207+
async def job():
208+
pass
209+
210+
ha.HomeAssistant.async_create_task(hass, job(), eager_start=True)
211+
assert len(hass.loop.call_soon.mock_calls) == 1
212+
# Should fallback to loop.create_task since 3.11 does
213+
# not support eager_start
214+
assert len(hass.loop.create_task.mock_calls) == 0
215+
assert len(hass.add_job.mock_calls) == 0
216+
217+
177218
def test_async_create_task_schedule_coroutine_with_name(event_loop) -> None:
178219
"""Test that we schedule coroutines and add jobs to the job pool with a name."""
179220
hass = MagicMock(loop=MagicMock(wraps=event_loop))
@@ -2598,7 +2639,8 @@ async def test_state_changed_events_to_not_leak_contexts(hass: HomeAssistant) ->
25982639
assert len(_get_by_type("homeassistant.core.Context")) == init_count
25992640

26002641

2601-
async def test_background_task(hass: HomeAssistant) -> None:
2642+
@pytest.mark.parametrize("eager_start", (True, False))
2643+
async def test_background_task(hass: HomeAssistant, eager_start: bool) -> None:
26022644
"""Test background tasks being quit."""
26032645
result = asyncio.Future()
26042646

@@ -2609,7 +2651,9 @@ async def test_task():
26092651
result.set_result(hass.state)
26102652
raise
26112653

2612-
task = hass.async_create_background_task(test_task(), "happy task")
2654+
task = hass.async_create_background_task(
2655+
test_task(), "happy task", eager_start=eager_start
2656+
)
26132657
assert "happy task" in str(task)
26142658
await asyncio.sleep(0)
26152659
await hass.async_stop()

‎tests/util/test_async.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"""Tests for async util methods from Python source."""
22
import asyncio
3+
import sys
34
import time
45
from unittest.mock import MagicMock, Mock, patch
56

@@ -246,3 +247,53 @@ async def test_callback_is_always_scheduled(hass: HomeAssistant) -> None:
246247
hasync.run_callback_threadsafe(hass.loop, callback)
247248

248249
mock_call_soon_threadsafe.assert_called_once()
250+
251+
252+
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Test requires Python 3.12+")
253+
async def test_create_eager_task_312(hass: HomeAssistant) -> None:
254+
"""Test create_eager_task schedules a task eagerly in the event loop.
255+
256+
For Python 3.12+, the task is scheduled eagerly in the event loop.
257+
"""
258+
events = []
259+
260+
async def _normal_task():
261+
events.append("normal")
262+
263+
async def _eager_task():
264+
events.append("eager")
265+
266+
task1 = hasync.create_eager_task(_eager_task())
267+
task2 = asyncio.create_task(_normal_task())
268+
269+
assert events == ["eager"]
270+
271+
await asyncio.sleep(0)
272+
assert events == ["eager", "normal"]
273+
await task1
274+
await task2
275+
276+
277+
@pytest.mark.skipif(sys.version_info >= (3, 12), reason="Test requires < Python 3.12")
278+
async def test_create_eager_task_pre_312(hass: HomeAssistant) -> None:
279+
"""Test create_eager_task schedules a task in the event loop.
280+
281+
For older python versions, the task is scheduled normally.
282+
"""
283+
events = []
284+
285+
async def _normal_task():
286+
events.append("normal")
287+
288+
async def _eager_task():
289+
events.append("eager")
290+
291+
task1 = hasync.create_eager_task(_eager_task())
292+
task2 = asyncio.create_task(_normal_task())
293+
294+
assert events == []
295+
296+
await asyncio.sleep(0)
297+
assert events == ["eager", "normal"]
298+
await task1
299+
await task2

0 commit comments

Comments
 (0)
0