jeans/supervisor
jeans/supervisor
1
0
Fork 0
mirror of https://github.com/home-assistant/supervisor.git synced 2025-07-20 23:56:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use context manager for concurrency control

Browse Source 
Instead of manually managing concurrency cleanup use a context manager
to ensure that resources are properly released even if an error occurs.
This commit is contained in:
Stefan Agner 2025-07-14 22:49:04 +02:00
parent 79964fd405
commit 7f584e386d
No known key found for this signature in database
GPG Key ID: AE01353D1E44747D
1 changed files with 38 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
supervisor/jobs/decorator.py
View File

@ -1,8 +1,8 @@
"""Job decorator.""" """Job decorator."""
import asyncio import asyncio
from collections.abc import Awaitable, Callable from collections.abc import AsyncIterator, Awaitable, Callable
from contextlib import suppress from contextlib import asynccontextmanager, suppress
from datetime import datetime, timedelta from datetime import datetime, timedelta
from functools import wraps from functools import wraps
import logging import logging
@ -318,37 +318,34 @@ class Job(CoreSysAttributes):
except JobConditionException as err: except JobConditionException as err:
return self._handle_job_condition_exception(err) return self._handle_job_condition_exception(err)
# Handle execution limits # Handle execution limits using context manager
await self._handle_concurrency_control(job_group, job) async with self._concurrency_control(job_group, job):
if not await self._handle_throttling(group_name): if not await self._handle_throttling(group_name):
self._release_concurrency_control(job_group) return # Job was throttled, exit early
return # Job was throttled, exit early
# Execute Job # Execute Job
with job.start(): with job.start():
try: try:
self.set_last_call(datetime.now(), group_name) self.set_last_call(datetime.now(), group_name)
if self._rate_limited_calls is not None: if self._rate_limited_calls is not None:
self.add_rate_limited_call( self.add_rate_limited_call(
self.last_call(group_name), group_name self.last_call(group_name), group_name
) )
return await method(obj, *args, **kwargs) return await method(obj, *args, **kwargs)
# If a method has a conditional JobCondition, they must check it in the method # If a method has a conditional JobCondition, they must check it in the method
# These should be handled like normal JobConditions as much as possible # These should be handled like normal JobConditions as much as possible
except JobConditionException as err: except JobConditionException as err:
return self._handle_job_condition_exception(err) return self._handle_job_condition_exception(err)
except HassioError as err: except HassioError as err:
job.capture_error(err) job.capture_error(err)
raise err raise err
except Exception as err: except Exception as err:
_LOGGER.exception("Unhandled exception: %s", err) _LOGGER.exception("Unhandled exception: %s", err)
job.capture_error() job.capture_error()
await async_capture_exception(err) await async_capture_exception(err)
raise JobException() from err raise JobException() from err
finally:
self._release_concurrency_control(job_group)
# Jobs that weren't started are always cleaned up. Also clean up done jobs if required # Jobs that weren't started are always cleaned up. Also clean up done jobs if required
finally: finally:
@ -533,6 +530,17 @@ class Job(CoreSysAttributes):
raise self.on_condition(str(err)) from err raise self.on_condition(str(err)) from err
raise err raise err
@asynccontextmanager
async def _concurrency_control(
self, job_group: JobGroup | None, job: SupervisorJob
) -> AsyncIterator[None]:
"""Context manager for concurrency control that ensures locks are always released."""
await self._handle_concurrency_control(job_group, job)
try:
yield
finally:
self._release_concurrency_control(job_group)
async def _handle_throttling(self, group_name: str | None) -> bool: async def _handle_throttling(self, group_name: str | None) -> bool:
"""Handle throttling limits. Returns True if job should continue, False if throttled.""" """Handle throttling limits. Returns True if job should continue, False if throttled."""
if self.throttle in (JobThrottle.THROTTLE, JobThrottle.GROUP_THROTTLE): if self.throttle in (JobThrottle.THROTTLE, JobThrottle.GROUP_THROTTLE):