91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import Counter
|
|
|
|
from devplacepy.database import get_int_setting
|
|
from devplacepy.services.nodes.session import NodeSession
|
|
from devplacepy.services.nodes.spec import SPEC
|
|
|
|
LIMITS = SPEC.limits
|
|
|
|
COUNTER_NAMES: tuple[str, ...] = (
|
|
"connected_total",
|
|
"rejected_capacity",
|
|
"rejected_ticket",
|
|
"rejected_consent",
|
|
"rejected_device",
|
|
"protocol_errors",
|
|
"evicted_heartbeat",
|
|
"kicked",
|
|
"replaced",
|
|
"jobs_assigned",
|
|
"jobs_done",
|
|
"jobs_failed",
|
|
"jobs_refused",
|
|
"integrity_failures",
|
|
"results_truncated",
|
|
)
|
|
|
|
|
|
class NodeCounters:
|
|
def __init__(self) -> None:
|
|
self._counts: Counter[str] = Counter({name: 0 for name in COUNTER_NAMES})
|
|
|
|
def inc(self, name: str) -> None:
|
|
if name not in COUNTER_NAMES:
|
|
raise KeyError(name)
|
|
self._counts[name] += 1
|
|
|
|
def snapshot(self) -> dict[str, int]:
|
|
return {name: self._counts[name] for name in COUNTER_NAMES}
|
|
|
|
|
|
class NodeRegistry:
|
|
def __init__(self, counters: NodeCounters) -> None:
|
|
self._sessions: dict[str, NodeSession] = {}
|
|
self.counters = counters
|
|
|
|
def capacity(self) -> int:
|
|
configured = get_int_setting("nodes_max_total", LIMITS.max_connected_nodes)
|
|
return max(0, min(configured, LIMITS.max_connected_nodes))
|
|
|
|
def count(self) -> int:
|
|
return len(self._sessions)
|
|
|
|
def sessions(self) -> list[NodeSession]:
|
|
return list(self._sessions.values())
|
|
|
|
def get(self, device_uid: str) -> NodeSession | None:
|
|
return self._sessions.get(device_uid)
|
|
|
|
def by_user(self, user_uid: str) -> list[NodeSession]:
|
|
return [item for item in self._sessions.values() if item.user_uid == user_uid]
|
|
|
|
def admit(self, session: NodeSession) -> tuple[bool, NodeSession | None]:
|
|
previous = self._sessions.get(session.device_uid)
|
|
if previous is None and self.count() >= self.capacity():
|
|
self.counters.inc("rejected_capacity")
|
|
return False, None
|
|
self._sessions[session.device_uid] = session
|
|
self.counters.inc("connected_total")
|
|
if previous is not None:
|
|
self.counters.inc("replaced")
|
|
return True, previous
|
|
|
|
def release(self, session: NodeSession) -> None:
|
|
if self._sessions.get(session.device_uid) is session:
|
|
del self._sessions[session.device_uid]
|
|
|
|
async def kick(self, session: NodeSession, reason: str) -> None:
|
|
self.release(session)
|
|
self.counters.inc("kicked")
|
|
await session.kick(reason)
|
|
|
|
async def kick_all(self, reason: str) -> int:
|
|
sessions = self.sessions()
|
|
for session in sessions:
|
|
await self.kick(session, reason)
|
|
return len(sessions)
|