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