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 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 def _close_channel(self):
60 asyncio.run(self._channel.close())
61
62 async def _monitor_channel(self):
63 try:
64 await asyncio.wait_for(self._channel.channel_ready(), timeout=5)
65 except TimeoutError:
66 logger.info("timeout for node, id: %i", self.id)
67 self.set_status(NodeStatus.unhealthy, "timeout")
68 while state := self._channel.get_state():
69 logger.debug("node %i state: %s", self.id, state.value)
70 try:
71 if state != ChannelConnectivity.READY:
72 raise RpcError
73 await self._sync()
74 self._streaming_task = asyncio.create_task(
75 self._stream_user_updates()
76 )
77 except RpcError:
78 self.synced = False
79 self.set_status(NodeStatus.unhealthy)
80 if self._streaming_task:
81 self._streaming_task.cancel()
82 else:
83 self.set_status(NodeStatus.healthy)
84 logger.info("Connected to node %i", self.id)
85
86 await self._channel.wait_for_state_change(state)
87
88 async def _stream_user_updates(self):
89 logger.debug("opened the stream")
90 stream = self._stub.SyncUsers()
91 while True:
92 user_update = await self._updates_queue.get()
93 logger.debug("got something from queue")
94 user = user_update["user"]
95 try:
96 await stream.write(
97 UserData(
98 user=User(
99 id=user.id, username=user.username, key=user.key
100 ),
101 inbounds=[
102 Inbound(tag=t) for t in user_update["inbounds"]
103 ],
104 )
105 )
106 except RpcError:
107 self.synced = False
108 self.set_status(NodeStatus.unhealthy)
109 return
110
111 async def update_user(self, user, inbounds: set[str] | None = None):
112 if inbounds is None:
113 inbounds = set()
114
115 await self._updates_queue.put({"user": user, "inbounds": inbounds})
116
117 async def _repopulate_users(self, users_data: list[dict]) -> None:
118 updates = [
119 UserData(
120 user=User(id=u["id"], username=u["username"], key=u["key"]),
121 inbounds=[Inbound(tag=t) for t in u["inbounds"]],
122 )
123 for u in users_data
124 ]
125 await self._stub.RepopulateUsers(UsersData(users_data=updates))
126
127 async def fetch_users_stats(self):
128 response = await self._stub.FetchUsersStats(Empty())
129 return response.users_stats
130
131 async def _fetch_backends(self) -> list:
132 response = await self._stub.FetchBackends(Empty())
133 return response.backends
134
135 async def _sync(self):
136 backends = await self._fetch_backends()
137 self.store_backends(backends)
138 users = self.list_users()
139 await self._repopulate_users(users)
140 self.synced = True
141
142 async def get_logs(self, name: str = "xray", include_buffer=True):
143 async for response in self._stub.StreamBackendLogs(
144 BackendLogsRequest(
145 backend_name=name, include_buffer=include_buffer
146 )
147 ):
148 yield response.line
149
150 async def restart_backend(
151 self, name: str, config: str, config_format: int
152 ):
153 try:
154 await self._stub.RestartBackend(
155 RestartBackendRequest(
156 backend_name=name,
157 config=BackendConfig(
158 configuration=config, config_format=config_format
159 ),
160 )
161 )
162 await self._sync()
163 except RpcError:
164 self.synced = False
165 self.set_status(NodeStatus.unhealthy)
166 raise
167 else:
168 self.set_status(NodeStatus.healthy)
169
170 async def get_backend_config(self, name: str = "xray"):
171 response = await self._stub.FetchBackendConfig(Backend(name=name))
172 return response.configuration, response.config_format
173
174 async def get_backend_stats(self, name: str):
175 response: BackendStats = await self._stub.GetBackendStats(
176 Backend(name=name)
177 )
178 return response