154 lines
5.1 KiB
Python
154 lines
5.1 KiB
Python
import ipaddress
|
|
import time
|
|
from config import settings
|
|
|
|
|
|
class LoadBalancer:
|
|
def __init__(self):
|
|
self._rr_index = 0
|
|
self._conn_counts: dict[str, int] = {}
|
|
self._node_health: dict[str, dict] = {}
|
|
|
|
def get_enabled_nodes(self):
|
|
return [n for n in settings.edge_nodes if n.get("enabled", True)]
|
|
|
|
def get_node_by_id(self, node_id: str):
|
|
for n in settings.edge_nodes:
|
|
if n["id"] == node_id:
|
|
return n
|
|
return None
|
|
|
|
def select_best_node(self, client_ip: str | None = None):
|
|
nodes = self.get_enabled_nodes()
|
|
if not nodes:
|
|
return None
|
|
|
|
if client_ip:
|
|
best = self._find_subnet_match(client_ip, nodes)
|
|
if best:
|
|
return best
|
|
|
|
strategy = settings.load_balance_strategy
|
|
if strategy == "round-robin":
|
|
online = [n for n in nodes if self.is_online(n["id"])]
|
|
if not online:
|
|
online = nodes
|
|
chosen = online[self._rr_index % len(online)]
|
|
self._rr_index += 1
|
|
return chosen
|
|
elif strategy == "weighted":
|
|
return self._weighted_select(nodes)
|
|
elif strategy == "least-conn":
|
|
return self._least_conn_select(nodes)
|
|
else:
|
|
return self._weighted_select(nodes)
|
|
|
|
def _find_subnet_match(self, client_ip: str, nodes: list[dict]):
|
|
try:
|
|
ip = ipaddress.ip_address(client_ip)
|
|
except ValueError:
|
|
return None
|
|
|
|
best = None
|
|
best_prefix = -1
|
|
|
|
for node in nodes:
|
|
if not self.is_online(node["id"]):
|
|
continue
|
|
|
|
subnet_str = node.get("subnet", "0.0.0.0/0")
|
|
if not subnet_str:
|
|
continue
|
|
|
|
try:
|
|
subnet = ipaddress.ip_network(subnet_str, strict=False)
|
|
except ValueError:
|
|
continue
|
|
|
|
if ip in subnet:
|
|
if subnet.prefixlen > best_prefix:
|
|
best_prefix = subnet.prefixlen
|
|
best = node
|
|
|
|
return best
|
|
|
|
def _weighted_select(self, nodes: list[dict]):
|
|
online = [n for n in nodes if self.is_online(n["id"])]
|
|
if not online:
|
|
online = nodes
|
|
total = sum(n["weight"] for n in online)
|
|
if total == 0:
|
|
return online[0] if online else None
|
|
import random
|
|
r = random.uniform(0, total)
|
|
cumulative = 0
|
|
for n in online:
|
|
cumulative += n["weight"]
|
|
if r <= cumulative:
|
|
return n
|
|
return online[-1]
|
|
|
|
def _least_conn_select(self, nodes: list[dict]):
|
|
online = [n for n in nodes if self.is_online(n["id"])]
|
|
if not online:
|
|
online = nodes
|
|
sorted_nodes = sorted(online, key=lambda n: self._conn_counts.get(n["id"], 0))
|
|
return sorted_nodes[0] if sorted_nodes else None
|
|
|
|
def report_health(self, node_id: str, metrics: dict):
|
|
self._node_health[node_id] = {
|
|
"last_seen": time.time(),
|
|
"online": True,
|
|
**metrics,
|
|
}
|
|
|
|
def is_online(self, node_id: str) -> bool:
|
|
h = self._node_health.get(node_id)
|
|
if h is None:
|
|
if any(n["id"] == node_id and n.get("is_origin") for n in settings.edge_nodes):
|
|
return True
|
|
return False
|
|
if time.time() - h.get("last_seen", 0) > 30:
|
|
return False
|
|
return h.get("online", False)
|
|
|
|
def get_health_summary(self):
|
|
result = {}
|
|
for node in settings.edge_nodes:
|
|
nid = node["id"]
|
|
h = self._node_health.get(nid)
|
|
result[nid] = {
|
|
"name": node["name"],
|
|
"host": node["host"],
|
|
"online": self.is_online(nid),
|
|
"weight": node["weight"],
|
|
"enabled": node.get("enabled", True),
|
|
"subnet": node.get("subnet"),
|
|
"is_origin": node.get("is_origin", False),
|
|
"last_seen": h.get("last_seen") if h else None,
|
|
"metrics": {
|
|
"cpu_percent": h.get("cpu_percent"),
|
|
"memory_percent": h.get("memory_percent"),
|
|
"bandwidth_mbps": h.get("bandwidth_mbps"),
|
|
"active_connections": h.get("active_connections"),
|
|
"hls_readers": h.get("hls_readers"),
|
|
"network_in_mbps": h.get("network_in_mbps"),
|
|
"network_out_mbps": h.get("network_out_mbps"),
|
|
"viewer_count": h.get("viewer_count"),
|
|
"frame_drop_indicators": h.get("frame_drop_indicators"),
|
|
"packet_loss_indicators": h.get("packet_loss_indicators"),
|
|
"media_metrics_available": h.get("media_metrics_available", False),
|
|
} if h else None,
|
|
}
|
|
return result
|
|
|
|
def increment_connection(self, node_id: str):
|
|
self._conn_counts[node_id] = self._conn_counts.get(node_id, 0) + 1
|
|
|
|
def decrement_connection(self, node_id: str):
|
|
if self._conn_counts.get(node_id, 0) > 0:
|
|
self._conn_counts[node_id] -= 1
|
|
|
|
|
|
lb = LoadBalancer()
|