1import asyncio
2import atexit
3import logging
4
5from _testcapi import INT_MAX
6from grpc import ChannelConnectivity, RpcError
7from grpc.aio import insecure_channel
8
9from .base import MarzNodeBase
10from .database import MarzNodeDB
11from .marznode_pb2 import (
12 UserData,
13 UsersData,
14 Empty,
15 User,
16 Inbound,
17 BackendLogsRequest,
18 RestartBackendRequest,
19 BackendConfig,
20 Backend,
21 BackendStats,
22)
23from .marznode_pb2_grpc import MarzServiceStub
24from ..models.node import NodeStatus
25
26logger = logging.getLogger(__name__)
27
28channel_options = [
29 ("grpc.keepalive_time_ms", 8000),
30 ("grpc.keepalive_timeout_ms", 5000),
31 ("grpc.http2.max_pings_without_data", 0),
32 ("grpc.keepalive_permit_without_calls", 1),
33 ("grpc.max_connection_idle_ms", INT_MAX),
34 ("grpc.client_idle_timeout_ms", INT_MAX),
35 ("grpc.max_connection_age_ms", INT_MAX),
36]
37
38
39class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB):
40 def __init__(
41 self, node_id: int, address: str, port: int, usage_coefficient: int = 1
42 ):
43 self.id = node_id
44 self._address = address
45 self._port = port
46
47 self._channel = insecure_channel(
48 f"{self._address}:{self._port}", channel_options
49 )
50 self._stub = MarzServiceStub(self._channel)
51 self._monitor_task = asyncio.create_task(self._monitor_channel())
52 self._streaming_task = None
53
54 self._updates_queue = asyncio.Queue(5)
55 self.synced = False
56 self.usage_coefficient = usage_coefficient
57 atexit.register(self._close_channel)
58
59 async def stop(self):
60 await self._channel.close()
61 self._monitor_task.cancel()
62
63 def _close_channel(self):
64 asyncio.run(self._channel.close())
65
66 async def _monitor_channel(self):
67 try:
68 await asyncio.wait_for(self._channel.channel_ready(), timeout=5)
69 except TimeoutError:
70 logger.info("timeout for node, id: %i", self.id)
71 self.set_status(NodeStatus.unhealthy, "timeout")
72 while state := self._channel.get_state():
73 logger.debug("node %i state: %s", self.id, state.value)
74 try:
75 if state != ChannelConnectivity.READY:
76 raise RpcError
77 await self._sync()
78 self._streaming_task = asyncio.create_task(
79 self._stream_user_updates()
80 )
81 except RpcError:
82 self.synced = False
83 self.set_status(NodeStatus.unhealthy)
84 if self._streaming_task:
85 self._streaming_task.cancel()
86 else:
87 self.set_status(NodeStatus.healthy)
88 logger.info("Connected to node %i", self.id)
89
90 await self._channel.wait_for_state_change(state)
91
92 async def _stream_user_updates(self):
93 logger.debug("opened the stream")
94 stream = self._stub.SyncUsers()
95 while True:
96 user_update = await self._updates_queue.get()
97 logger.debug("got something from queue")
98 user = user_update["user"]
99 try:
100 await stream.write(
101 UserData(
102 user=User(
103 id=user.id, username=user.username, key=user.key
104 ),
105 inbounds=[
106 Inbound(tag=t) for t in user_update["inbounds"]
107 ],
108 )
109 )
110 except RpcError:
111 self.synced = False
112 self.set_status(NodeStatus.unhealthy)
113 return
114
115 async def update_user(self, user, inbounds: set[str] | None = None):
116 if inbounds is None:
117 inbounds = set()
118
119 await self._updates_queue.put({"user": user, "inbounds": inbounds})
120
121 async def _repopulate_users(self, users_data: list[dict]) -> None:
122 updates = [
123 UserData(
124 user=User(id=u["id"], username=u["username"], key=u["key"]),
125 inbounds=[Inbound(tag=t) for t in u["inbounds"]],
126 )
127 for u in users_data
128 ]
129 await self._stub.RepopulateUsers(UsersData(users_data=updates))
130
131 async def fetch_users_stats(self):
132 response = await self._stub.FetchUsersStats(Empty())
133 return response.users_stats
134
135 async def _fetch_backends(self) -> list:
136 response = await self._stub.FetchBackends(Empty())
137 return response.backends
138
139 async def _sync(self):
140 backends = await self._fetch_backends()
141 self.store_backends(backends)
142 users = self.list_users()
143 await self._repopulate_users(users)
144 self.synced = True
145
146 async def get_logs(self, name: str = "xray", include_buffer=True):
147 async for response in self._stub.StreamBackendLogs(
148 BackendLogsRequest(
149 backend_name=name, include_buffer=include_buffer
150 )
151 ):
152 yield response.line
153
154 async def restart_backend(
155 self, name: str, config: str, config_format: int
156 ):
157 try:
158 await self._stub.RestartBackend(
159 RestartBackendRequest(
160 backend_name=name,
161 config=BackendConfig(
162 configuration=config, config_format=config_format
163 ),
164 )
165 )
166 await self._sync()
167 except RpcError:
168 self.synced = False
169 self.set_status(NodeStatus.unhealthy)
170 raise
171 else:
172 self.set_status(NodeStatus.healthy)
173
174 async def get_backend_config(self, name: str = "xray"):
175 response = await self._stub.FetchBackendConfig(Backend(name=name))
176 return response.configuration, response.config_format
177
178 async def get_backend_stats(self, name: str):
179 response: BackendStats = await self._stub.GetBackendStats(
180 Backend(name=name)
181 )
182 return response