44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
import secrets
|
|
import threading
|
|
|
|
import anyio
|
|
|
|
|
|
class PingManagerBase:
|
|
def _generate_id(self) -> bytes:
|
|
return secrets.token_bytes()
|
|
|
|
|
|
class PingManager(PingManagerBase):
|
|
def __init__(self) -> None:
|
|
self._pings: dict[bytes, threading.Event] = {}
|
|
|
|
def create(self, ping_id: bytes | None = None) -> tuple[bytes, threading.Event]:
|
|
ping_id = self._generate_id() if not ping_id else ping_id
|
|
event = threading.Event()
|
|
self._pings[ping_id] = event
|
|
return ping_id, event
|
|
|
|
def ack(self, ping_id: bytes | bytearray) -> None:
|
|
event = self._pings.pop(bytes(ping_id), None)
|
|
if event is not None:
|
|
event.set()
|
|
|
|
|
|
class AsyncPingManager(PingManagerBase):
|
|
def __init__(self) -> None:
|
|
self._pings: dict[bytes, anyio.Event] = {}
|
|
|
|
def create(self, ping_id: bytes | None = None) -> tuple[bytes, anyio.Event]:
|
|
ping_id = self._generate_id() if not ping_id else ping_id
|
|
event = anyio.Event()
|
|
self._pings[ping_id] = event
|
|
return ping_id, event
|
|
|
|
def ack(self, ping_id: bytes | bytearray) -> None:
|
|
event = self._pings.pop(bytes(ping_id), None)
|
|
if event is not None:
|
|
event.set()
|