mirror of
https://github.com/HKUDS/OpenSpace.git
synced 2026-08-28 05:15:00 +00:00
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
"""
|
|
Query execution guard — ensures mutual exclusion for concurrent queries.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import enum
|
|
from typing import Optional
|
|
|
|
from openspace.utils.logging import Logger
|
|
|
|
logger = Logger.get_logger(__name__)
|
|
|
|
|
|
class QueryState(enum.Enum):
|
|
IDLE = "idle"
|
|
RUNNING = "running"
|
|
RESERVED = "reserved"
|
|
|
|
|
|
class QueryGuard:
|
|
"""Ensures only one query runs at a time with optional reservation."""
|
|
|
|
def __init__(self) -> None:
|
|
self._state = QueryState.IDLE
|
|
self._lock = asyncio.Lock()
|
|
self._cancel_event = asyncio.Event()
|
|
self._reserved_query: Optional[str] = None
|
|
self._running_task: Optional[asyncio.Task] = None
|
|
|
|
@property
|
|
def state(self) -> QueryState:
|
|
return self._state
|
|
|
|
@property
|
|
def is_idle(self) -> bool:
|
|
return self._state == QueryState.IDLE
|
|
|
|
@property
|
|
def is_running(self) -> bool:
|
|
return self._state == QueryState.RUNNING
|
|
|
|
@property
|
|
def is_cancelled(self) -> bool:
|
|
return self._cancel_event.is_set()
|
|
|
|
async def try_start(self) -> bool:
|
|
"""Attempt to transition from IDLE to RUNNING.
|
|
Returns True if the guard was acquired, False if already running."""
|
|
async with self._lock:
|
|
if self._state != QueryState.IDLE:
|
|
logger.warning("QueryGuard: cannot start — state is %s", self._state.value)
|
|
return False
|
|
self._state = QueryState.RUNNING
|
|
self._cancel_event.clear()
|
|
logger.debug("QueryGuard: IDLE → RUNNING")
|
|
return True
|
|
|
|
def end(self) -> None:
|
|
"""Transition back to IDLE after query completes."""
|
|
prev = self._state
|
|
self._state = QueryState.IDLE
|
|
self._running_task = None
|
|
self._cancel_event.clear()
|
|
logger.debug("QueryGuard: %s → IDLE", prev.value)
|
|
|
|
def reserve(self, query: str) -> None:
|
|
"""Queue a query for execution when the current one finishes."""
|
|
self._reserved_query = query
|
|
self._state = QueryState.RESERVED
|
|
logger.debug("QueryGuard: reserved query (%.60s...)", query)
|
|
|
|
def take_reserved(self) -> Optional[str]:
|
|
"""Consume the reserved query, returning it or None."""
|
|
q = self._reserved_query
|
|
self._reserved_query = None
|
|
if self._state == QueryState.RESERVED:
|
|
self._state = QueryState.IDLE
|
|
return q
|
|
|
|
def cancel(self) -> None:
|
|
"""Signal cancellation to the running query."""
|
|
if self._state == QueryState.RUNNING:
|
|
self._cancel_event.set()
|
|
logger.info("QueryGuard: cancel requested")
|
|
if self._running_task and not self._running_task.done():
|
|
self._running_task.cancel()
|
|
|
|
def bind_task(self, task: asyncio.Task) -> None:
|
|
"""Bind the currently running asyncio.Task for cancellation support."""
|
|
self._running_task = task
|
|
|
|
async def wait_for_cancel(self) -> None:
|
|
"""Await until cancel() is called. Useful for cooperative cancellation."""
|
|
await self._cancel_event.wait()
|
|
|
|
def check_cancelled(self) -> None:
|
|
"""Raise asyncio.CancelledError if cancellation was requested."""
|
|
if self._cancel_event.is_set():
|
|
raise asyncio.CancelledError("Query cancelled via QueryGuard")
|