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