1import asyncio
2import atexit
3import logging
4from _testcapi import INT_MAX
5
6from grpc import ChannelConnectivity
7from grpc.aio import insecure_channel
8
9from .base import MarzNodeBase
10from .database import MarzNodeDB
11from .marznode_pb2_grpc import MarzServiceStub
12from .marznode_pb2 import UserData, UsersData, Empty, User, Inbound
13from ..models.node import NodeStatus
14
15logger = logging.getLogger(__name__)
16
17channel_options = [
18 ("grpc.keepalive_time_ms", 8000),
19 ("grpc.keepalive_timeout_ms", 5000),
20 ("grpc.http2.max_pings_without_data", 0),
21 ("grpc.keepalive_permit_without_calls", 1),
22 ("grpc.max_connection_idle_ms", INT_MAX),
23 ("grpc.client_idle_timeout_ms", INT_MAX),
24 ("grpc.max_connection_age_ms", INT_MAX)
25]
26
27
28class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB):
29 def __init__(self, node_id: int, address: str, port: int, usage_coefficient: int = 1):
30 self.id = node_id
31 self._address = address
32 self._port = port
33
34 self._channel = insecure_channel(f"{self._address}:{self._port}", channel_options)
35 self._stub = MarzServiceStub(self._channel)
36 asyncio.create_task(self._monitor_channel())
37 self._streaming_task = None
38
39 self._updates_queue = asyncio.Queue(1)
40 self.synced = False
41 self.usage_coefficient = usage_coefficient
42 atexit.register(self._close_channel)
43
44 def _close_channel(self):
45 asyncio.run(self._channel.close())
46
47 async def _monitor_channel(self):
48 try:
49 await asyncio.wait_for(self._channel.channel_ready(), timeout=5)
50 except TimeoutError:
51 logger.info("timeout for node, id: %i", self.id)
52 self.set_status(NodeStatus.unhealthy, "timeout")
53 while state := self._channel.get_state():
54 logger.debug("node %i state: %s", self.id, state.value)
55 if state == ChannelConnectivity.READY:
56 await self._sync()
57 self._streaming_task = asyncio.create_task(self._stream_user_updates())
58 self.set_status(NodeStatus.healthy)
59 logger.info("Connected to node %i", self.id)
60 else:
61 self.synced = False
62 self.set_status(NodeStatus.unhealthy)
63 if self._streaming_task:
64 self._streaming_task.cancel()
65
66 await self._channel.wait_for_state_change(state)
67
68 async def _stream_user_updates(self):
69 logger.debug("opened the stream")
70 stream = self._stub.SyncUsers()
71 while True:
72 user_update = await self._updates_queue.get()
73 logger.debug("got something from queue")
74 user = user_update["user"]
75 await stream.write(
76 UserData(user=User(id=user.id, username=user.username, key=user.key),
77 inbounds=[Inbound(tag=t) for t in user_update["inbounds"]]))
78
79 async def update_user(self, user, inbounds: set[str] | None = None):
80 if inbounds is None:
81 inbounds = set()
82
83 if self.synced:
84 await self._updates_queue.put({"user": user, "inbounds": inbounds})
85
86 async def _repopulate_users(self, users_data: list[dict]) -> None:
87 updates = [UserData(user=User(id=u["id"], username=u["username"], key=u["key"]),
88 inbounds=[Inbound(tag=t) for t in u["inbounds"]]) for u in users_data]
89 await self._stub.RepopulateUsers(UsersData(users_data=updates))
90
91 async def fetch_users_stats(self):
92 response = await self._stub.FetchUsersStats(Empty())
93 return response.users_stats
94
95 async def _fetch_inbounds(self) -> list:
96 response = await self._stub.FetchInbounds(Empty())
97 return response.inbounds
98
99 async def _sync(self):
100 inbounds = await self._fetch_inbounds()
101 self.store_inbounds(inbounds)
102 users = self.list_users()
103 await self._repopulate_users(users)
104 self.synced = True