wehub-resource-sync
146 lines
4.4 KiB
Python
146 lines
4.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from enum import Enum, auto
|
|
import asyncio
|
|
import threading
|
|
|
|
|
|
class EventType(Enum):
|
|
LOADING = auto()
|
|
|
|
|
|
class EventStatus(Enum):
|
|
ONGOING = auto()
|
|
DONE = auto()
|
|
NOT_FOUND = auto()
|
|
|
|
|
|
class EventManager:
|
|
"""
|
|
A thread-safe event manager to manage asynchronous events.
|
|
Each event is identified by its type and a unique id.
|
|
Events are organized by status for efficient counting.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
# Guard by lock
|
|
# Structure: events[event_type][event_status][event_id] = future
|
|
self.events: dict[EventType, dict[EventStatus, dict[str, asyncio.Future]]] = {
|
|
et: {es: {} for es in EventStatus} for et in EventType
|
|
}
|
|
self.lock = threading.Lock()
|
|
|
|
def add_event(
|
|
self,
|
|
event_type: EventType,
|
|
event_id: str,
|
|
future: asyncio.Future,
|
|
) -> None:
|
|
"""
|
|
Add an event with the given type and id.
|
|
"""
|
|
with self.lock:
|
|
status_dict = self.events.get(event_type, None)
|
|
assert status_dict is not None, (
|
|
f"Invalid event type {event_type} in EventManager."
|
|
)
|
|
status_dict[EventStatus.ONGOING][event_id] = future
|
|
|
|
def pop_event(
|
|
self,
|
|
event_type: EventType,
|
|
event_id: str,
|
|
) -> asyncio.Future:
|
|
"""
|
|
Pop and return the event with the given type and id.
|
|
"""
|
|
with self.lock:
|
|
status_dict = self.events.get(event_type, None)
|
|
assert status_dict is not None, (
|
|
f"Invalid event type {event_type} in EventManager."
|
|
)
|
|
done_events = status_dict[EventStatus.DONE]
|
|
assert event_id in done_events, (
|
|
f"Event {event_id} of type {event_type} is not done or not found."
|
|
)
|
|
return done_events.pop(event_id)
|
|
|
|
def update_event_status(
|
|
self,
|
|
event_type: EventType,
|
|
event_id: str,
|
|
status: EventStatus,
|
|
) -> None:
|
|
"""
|
|
Update the status of the event with the given type and id.
|
|
"""
|
|
with self.lock:
|
|
status_dict = self.events.get(event_type, None)
|
|
assert status_dict is not None, (
|
|
f"Invalid event type {event_type} in EventManager."
|
|
)
|
|
# Find the event in any status dict
|
|
event = None
|
|
for s in EventStatus:
|
|
if event_id in status_dict[s]:
|
|
event = status_dict[s].pop(event_id)
|
|
break
|
|
|
|
if event is None:
|
|
raise KeyError(f"Event {event_id} of type {event_type} not found.")
|
|
|
|
# Move to new status dict
|
|
status_dict[status][event_id] = event
|
|
|
|
def get_event_status(
|
|
self,
|
|
event_type: EventType,
|
|
event_id: str,
|
|
) -> EventStatus:
|
|
"""
|
|
Get the status of the event with the given type and id.
|
|
"""
|
|
with self.lock:
|
|
status_dict = self.events.get(event_type, None)
|
|
assert status_dict is not None, (
|
|
f"Invalid event type {event_type} in EventManager."
|
|
)
|
|
for status in EventStatus:
|
|
if event_id in status_dict[status]:
|
|
return status
|
|
return EventStatus.NOT_FOUND
|
|
|
|
def get_events_count_by_status(
|
|
self,
|
|
event_type: EventType,
|
|
status: EventStatus,
|
|
) -> int:
|
|
"""
|
|
Get the count of events for the given event type and status.
|
|
This is a lightweight O(1) operation using dict length.
|
|
"""
|
|
with self.lock:
|
|
status_dict = self.events.get(event_type, None)
|
|
if status_dict is None:
|
|
return 0
|
|
return len(status_dict[status])
|
|
|
|
def get_event_future(
|
|
self,
|
|
event_type: EventType,
|
|
event_id: str,
|
|
) -> asyncio.Future:
|
|
"""
|
|
Pop and return the event with the given type and id.
|
|
"""
|
|
with self.lock:
|
|
status_dict = self.events.get(event_type, None)
|
|
assert status_dict is not None, (
|
|
f"Invalid event type {event_type} in EventManager."
|
|
)
|
|
done_events = status_dict[EventStatus.DONE]
|
|
assert event_id in done_events, (
|
|
f"Event {event_id} of type {event_type} is not done or not found."
|
|
)
|
|
return done_events[event_id]
|