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.3.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 01cdf2e Dawsh 2y ago XrayLogsRequest,
18 01cdf2e Dawsh 2y ago XrayConfig,
19 01cdf2e Dawsh 2y ago )
20 49457ba Dawsh 2y ago from .marznode_pb2_grpc import MarzServiceStub
21 f6a32ea Dawsh 2y ago from ..models.node import NodeStatus
22 f6a32ea Dawsh 2y ago
23 f6a32ea Dawsh 2y ago logger = logging.getLogger(__name__)
24 f6a32ea Dawsh 2y ago
25 f6a32ea Dawsh 2y ago channel_options = [
26 f6a32ea Dawsh 2y ago ("grpc.keepalive_time_ms", 8000),
27 f6a32ea Dawsh 2y ago ("grpc.keepalive_timeout_ms", 5000),
28 f6a32ea Dawsh 2y ago ("grpc.http2.max_pings_without_data", 0),
29 f6a32ea Dawsh 2y ago ("grpc.keepalive_permit_without_calls", 1),
30 f6a32ea Dawsh 2y ago ("grpc.max_connection_idle_ms", INT_MAX),
31 f6a32ea Dawsh 2y ago ("grpc.client_idle_timeout_ms", INT_MAX),
32 01cdf2e Dawsh 2y ago ("grpc.max_connection_age_ms", INT_MAX),
33 f6a32ea Dawsh 2y ago ]
34 f6a32ea Dawsh 2y ago
35 f6a32ea Dawsh 2y ago
36 f6a32ea Dawsh 2y ago class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB):
37 01cdf2e Dawsh 2y ago def __init__(
38 01cdf2e Dawsh 2y ago self, node_id: int, address: str, port: int, usage_coefficient: int = 1
39 01cdf2e Dawsh 2y ago ):
40 f6a32ea Dawsh 2y ago self.id = node_id
41 f6a32ea Dawsh 2y ago self._address = address
42 f6a32ea Dawsh 2y ago self._port = port
43 f6a32ea Dawsh 2y ago
44 01cdf2e Dawsh 2y ago self._channel = insecure_channel(
45 01cdf2e Dawsh 2y ago f"{self._address}:{self._port}", channel_options
46 01cdf2e Dawsh 2y ago )
47 f6a32ea Dawsh 2y ago self._stub = MarzServiceStub(self._channel)
48 f6a32ea Dawsh 2y ago asyncio.create_task(self._monitor_channel())
49 f6a32ea Dawsh 2y ago self._streaming_task = None
50 f6a32ea Dawsh 2y ago
51 49457ba Dawsh 2y ago self._updates_queue = asyncio.Queue(5)
52 f6a32ea Dawsh 2y ago self.synced = False
53 f6a32ea Dawsh 2y ago self.usage_coefficient = usage_coefficient
54 f6a32ea Dawsh 2y ago atexit.register(self._close_channel)
55 f6a32ea Dawsh 2y ago
56 f6a32ea Dawsh 2y ago def _close_channel(self):
57 f6a32ea Dawsh 2y ago asyncio.run(self._channel.close())
58 f6a32ea Dawsh 2y ago
59 f6a32ea Dawsh 2y ago async def _monitor_channel(self):
60 f6a32ea Dawsh 2y ago try:
61 f6a32ea Dawsh 2y ago await asyncio.wait_for(self._channel.channel_ready(), timeout=5)
62 f6a32ea Dawsh 2y ago except TimeoutError:
63 f6a32ea Dawsh 2y ago logger.info("timeout for node, id: %i", self.id)
64 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.unhealthy, "timeout")
65 f6a32ea Dawsh 2y ago while state := self._channel.get_state():
66 f6a32ea Dawsh 2y ago logger.debug("node %i state: %s", self.id, state.value)
67 49457ba Dawsh 2y ago try:
68 49457ba Dawsh 2y ago if state != ChannelConnectivity.READY:
69 49457ba Dawsh 2y ago raise RpcError
70 f6a32ea Dawsh 2y ago await self._sync()
71 f1ab16c Dawsh 2y ago self._streaming_task = asyncio.create_task(
72 f1ab16c Dawsh 2y ago self._stream_user_updates()
73 f1ab16c Dawsh 2y ago )
74 49457ba Dawsh 2y ago except RpcError:
75 f6a32ea Dawsh 2y ago self.synced = False
76 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.unhealthy)
77 f6a32ea Dawsh 2y ago if self._streaming_task:
78 f6a32ea Dawsh 2y ago self._streaming_task.cancel()
79 49457ba Dawsh 2y ago else:
80 49457ba Dawsh 2y ago self.set_status(NodeStatus.healthy)
81 49457ba Dawsh 2y ago logger.info("Connected to node %i", self.id)
82 f6a32ea Dawsh 2y ago
83 f6a32ea Dawsh 2y ago await self._channel.wait_for_state_change(state)
84 f6a32ea Dawsh 2y ago
85 f6a32ea Dawsh 2y ago async def _stream_user_updates(self):
86 f6a32ea Dawsh 2y ago logger.debug("opened the stream")
87 f6a32ea Dawsh 2y ago stream = self._stub.SyncUsers()
88 f6a32ea Dawsh 2y ago while True:
89 f6a32ea Dawsh 2y ago user_update = await self._updates_queue.get()
90 f6a32ea Dawsh 2y ago logger.debug("got something from queue")
91 f6a32ea Dawsh 2y ago user = user_update["user"]
92 49457ba Dawsh 2y ago try:
93 49457ba Dawsh 2y ago await stream.write(
94 49457ba Dawsh 2y ago UserData(
95 f1ab16c Dawsh 2y ago user=User(
96 f1ab16c Dawsh 2y ago id=user.id, username=user.username, key=user.key
97 f1ab16c Dawsh 2y ago ),
98 f1ab16c Dawsh 2y ago inbounds=[
99 f1ab16c Dawsh 2y ago Inbound(tag=t) for t in user_update["inbounds"]
100 f1ab16c Dawsh 2y ago ],
101 49457ba Dawsh 2y ago )
102 01cdf2e Dawsh 2y ago )
103 49457ba Dawsh 2y ago except RpcError:
104 49457ba Dawsh 2y ago self.synced = False
105 49457ba Dawsh 2y ago self.set_status(NodeStatus.unhealthy)
106 49457ba Dawsh 2y ago return
107 f6a32ea Dawsh 2y ago
108 f6a32ea Dawsh 2y ago async def update_user(self, user, inbounds: set[str] | None = None):
109 f6a32ea Dawsh 2y ago if inbounds is None:
110 f6a32ea Dawsh 2y ago inbounds = set()
111 f6a32ea Dawsh 2y ago
112 4a2faf0 Dawsh 2y ago await self._updates_queue.put({"user": user, "inbounds": inbounds})
113 f6a32ea Dawsh 2y ago
114 f6a32ea Dawsh 2y ago async def _repopulate_users(self, users_data: list[dict]) -> None:
115 01cdf2e Dawsh 2y ago updates = [
116 01cdf2e Dawsh 2y ago UserData(
117 01cdf2e Dawsh 2y ago user=User(id=u["id"], username=u["username"], key=u["key"]),
118 01cdf2e Dawsh 2y ago inbounds=[Inbound(tag=t) for t in u["inbounds"]],
119 01cdf2e Dawsh 2y ago )
120 01cdf2e Dawsh 2y ago for u in users_data
121 01cdf2e Dawsh 2y ago ]
122 f6a32ea Dawsh 2y ago await self._stub.RepopulateUsers(UsersData(users_data=updates))
123 f6a32ea Dawsh 2y ago
124 f6a32ea Dawsh 2y ago async def fetch_users_stats(self):
125 f6a32ea Dawsh 2y ago response = await self._stub.FetchUsersStats(Empty())
126 f6a32ea Dawsh 2y ago return response.users_stats
127 f6a32ea Dawsh 2y ago
128 f6a32ea Dawsh 2y ago async def _fetch_inbounds(self) -> list:
129 f6a32ea Dawsh 2y ago response = await self._stub.FetchInbounds(Empty())
130 f6a32ea Dawsh 2y ago return response.inbounds
131 f6a32ea Dawsh 2y ago
132 f6a32ea Dawsh 2y ago async def _sync(self):
133 f6a32ea Dawsh 2y ago inbounds = await self._fetch_inbounds()
134 f6a32ea Dawsh 2y ago self.store_inbounds(inbounds)
135 f6a32ea Dawsh 2y ago users = self.list_users()
136 f6a32ea Dawsh 2y ago await self._repopulate_users(users)
137 f6a32ea Dawsh 2y ago self.synced = True
138 512d183 Mrdn 2y ago
139 512d183 Mrdn 2y ago async def get_logs(self, include_buffer=True):
140 512d183 Mrdn 2y ago async for response in self._stub.StreamXrayLogs(
141 01cdf2e Dawsh 2y ago XrayLogsRequest(include_buffer=include_buffer)
142 512d183 Mrdn 2y ago ):
143 512d183 Mrdn 2y ago yield response.line
144 512d183 Mrdn 2y ago
145 512d183 Mrdn 2y ago async def restart_xray(self, config: str):
146 49457ba Dawsh 2y ago try:
147 49457ba Dawsh 2y ago await self._stub.RestartXray(XrayConfig(configuration=config))
148 49457ba Dawsh 2y ago await self._sync()
149 49457ba Dawsh 2y ago except RpcError:
150 49457ba Dawsh 2y ago self.synced = False
151 49457ba Dawsh 2y ago self.set_status(NodeStatus.unhealthy)
152 49457ba Dawsh 2y ago raise
153 49457ba Dawsh 2y ago else:
154 49457ba Dawsh 2y ago self.set_status(NodeStatus.healthy)
155 512d183 Mrdn 2y ago
156 512d183 Mrdn 2y ago async def get_xray_config(self):
157 512d183 Mrdn 2y ago response = await self._stub.FetchXrayConfig(Empty())
158 512d183 Mrdn 2y ago return response.configuration