core/homeassistant/block_async_io.py
J. Nick Koston 1fef4fa1f6
Prevent time.sleep calls from blocking the event loop (#118561)
* Prevent time.sleep calls from blocking the event loop

We have been warning on these since Jan 2022. 2+ years seems more than enough
time to give to fix these. see https://github.com/home-assistant/core/pull/63766

* Prevent time.sleep calls from blocking the event loop

We have been warning on these since Jan 2022. 2+ years seems more than enough
time to give to fix these. see https://github.com/home-assistant/core/pull/63766
2024-05-31 11:08:22 -04:00

80 lines
2.5 KiB
Python

"""Block blocking calls being done in asyncio."""
import builtins
from contextlib import suppress
from http.client import HTTPConnection
import importlib
import sys
import threading
import time
from typing import Any
from .helpers.frame import get_current_frame
from .util.loop import protect_loop
_IN_TESTS = "unittest" in sys.modules
ALLOWED_FILE_PREFIXES = ("/proc",)
def _check_import_call_allowed(mapped_args: dict[str, Any]) -> bool:
# If the module is already imported, we can ignore it.
return bool((args := mapped_args.get("args")) and args[0] in sys.modules)
def _check_file_allowed(mapped_args: dict[str, Any]) -> bool:
# If the file is in /proc we can ignore it.
args = mapped_args["args"]
path = args[0] if type(args[0]) is str else str(args[0]) # noqa: E721
return path.startswith(ALLOWED_FILE_PREFIXES)
def _check_sleep_call_allowed(mapped_args: dict[str, Any]) -> bool:
#
# Avoid extracting the stack unless we need to since it
# will have to access the linecache which can do blocking
# I/O and we are trying to avoid blocking calls.
#
# frame[0] is us
# frame[1] is raise_for_blocking_call
# frame[2] is protected_loop_func
# frame[3] is the offender
with suppress(ValueError):
return get_current_frame(4).f_code.co_filename.endswith("pydevd.py")
return False
def enable() -> None:
"""Enable the detection of blocking calls in the event loop."""
loop_thread_id = threading.get_ident()
# Prevent urllib3 and requests doing I/O in event loop
HTTPConnection.putrequest = protect_loop( # type: ignore[method-assign]
HTTPConnection.putrequest, loop_thread_id=loop_thread_id
)
# Prevent sleeping in event loop.
time.sleep = protect_loop(
time.sleep,
check_allowed=_check_sleep_call_allowed,
loop_thread_id=loop_thread_id,
)
if not _IN_TESTS:
# Prevent files being opened inside the event loop
builtins.open = protect_loop( # type: ignore[assignment]
builtins.open,
strict_core=False,
strict=False,
check_allowed=_check_file_allowed,
loop_thread_id=loop_thread_id,
)
# unittest uses `importlib.import_module` to do mocking
# so we cannot protect it if we are running tests
importlib.import_module = protect_loop(
importlib.import_module,
strict_core=False,
strict=False,
check_allowed=_check_import_call_allowed,
loop_thread_id=loop_thread_id,
)