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.0
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 f6a32ea Dawsh 2y ago self._streaming_task = asyncio.create_task(self._stream_user_updates())
72 49457ba Dawsh 2y ago except RpcError:
73 f6a32ea Dawsh 2y ago self.synced = False
74 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.unhealthy)
75 f6a32ea Dawsh 2y ago if self._streaming_task:
76 f6a32ea Dawsh 2y ago self._streaming_task.cancel()
77 49457ba Dawsh 2y ago else:
78 49457ba Dawsh 2y ago self.set_status(NodeStatus.healthy)
79 49457ba Dawsh 2y ago logger.info("Connected to node %i", self.id)
80 f6a32ea Dawsh 2y ago
81 f6a32ea Dawsh 2y ago await self._channel.wait_for_state_change(state)
82 f6a32ea Dawsh 2y ago
83 f6a32ea Dawsh 2y ago async def _stream_user_updates(self):
84 f6a32ea Dawsh 2y ago logger.debug("opened the stream")
85 f6a32ea Dawsh 2y ago stream = self._stub.SyncUsers()
86 f6a32ea Dawsh 2y ago while True:
87 f6a32ea Dawsh 2y ago user_update = await self._updates_queue.get()
88 f6a32ea Dawsh 2y ago logger.debug("got something from queue")
89 f6a32ea Dawsh 2y ago user = user_update["user"]
90 49457ba Dawsh 2y ago try:
91 49457ba Dawsh 2y ago await stream.write(
92 49457ba Dawsh 2y ago UserData(
93 49457ba Dawsh 2y ago user=User(id=user.id, username=user.username, key=user.key),
94 49457ba Dawsh 2y ago inbounds=[Inbound(tag=t) for t in user_update["inbounds"]],
95 49457ba Dawsh 2y ago )
96 01cdf2e Dawsh 2y ago )
97 49457ba Dawsh 2y ago except RpcError:
98 49457ba Dawsh 2y ago self.synced = False
99 49457ba Dawsh 2y ago self.set_status(NodeStatus.unhealthy)
100 49457ba Dawsh 2y ago return
101 f6a32ea Dawsh 2y ago
102 f6a32ea Dawsh 2y ago async def update_user(self, user, inbounds: set[str] | None = None):
103 f6a32ea Dawsh 2y ago if inbounds is None:
104 f6a32ea Dawsh 2y ago inbounds = set()
105 f6a32ea Dawsh 2y ago
106 f6a32ea Dawsh 2y ago if self.synced:
107 f6a32ea Dawsh 2y ago await self._updates_queue.put({"user": user, "inbounds": inbounds})
108 f6a32ea Dawsh 2y ago
109 f6a32ea Dawsh 2y ago async def _repopulate_users(self, users_data: list[dict]) -> None:
110 01cdf2e Dawsh 2y ago updates = [
111 01cdf2e Dawsh 2y ago UserData(
112 01cdf2e Dawsh 2y ago user=User(id=u["id"], username=u["username"], key=u["key"]),
113 01cdf2e Dawsh 2y ago inbounds=[Inbound(tag=t) for t in u["inbounds"]],
114 01cdf2e Dawsh 2y ago )
115 01cdf2e Dawsh 2y ago for u in users_data
116 01cdf2e Dawsh 2y ago ]
117 f6a32ea Dawsh 2y ago await self._stub.RepopulateUsers(UsersData(users_data=updates))
118 f6a32ea Dawsh 2y ago
119 f6a32ea Dawsh 2y ago async def fetch_users_stats(self):
120 f6a32ea Dawsh 2y ago response = await self._stub.FetchUsersStats(Empty())
121 f6a32ea Dawsh 2y ago return response.users_stats
122 f6a32ea Dawsh 2y ago
123 f6a32ea Dawsh 2y ago async def _fetch_inbounds(self) -> list:
124 f6a32ea Dawsh 2y ago response = await self._stub.FetchInbounds(Empty())
125 f6a32ea Dawsh 2y ago return response.inbounds
126 f6a32ea Dawsh 2y ago
127 f6a32ea Dawsh 2y ago async def _sync(self):
128 f6a32ea Dawsh 2y ago inbounds = await self._fetch_inbounds()
129 f6a32ea Dawsh 2y ago self.store_inbounds(inbounds)
130 f6a32ea Dawsh 2y ago users = self.list_users()
131 f6a32ea Dawsh 2y ago await self._repopulate_users(users)
132 f6a32ea Dawsh 2y ago self.synced = True
133 512d183 Mrdn 2y ago
134 512d183 Mrdn 2y ago async def get_logs(self, include_buffer=True):
135 512d183 Mrdn 2y ago async for response in self._stub.StreamXrayLogs(
136 01cdf2e Dawsh 2y ago XrayLogsRequest(include_buffer=include_buffer)
137 512d183 Mrdn 2y ago ):
138 512d183 Mrdn 2y ago yield response.line
139 512d183 Mrdn 2y ago
140 512d183 Mrdn 2y ago async def restart_xray(self, config: str):
141 49457ba Dawsh 2y ago try:
142 49457ba Dawsh 2y ago await self._stub.RestartXray(XrayConfig(configuration=config))
143 49457ba Dawsh 2y ago await self._sync()
144 49457ba Dawsh 2y ago except RpcError:
145 49457ba Dawsh 2y ago self.synced = False
146 49457ba Dawsh 2y ago self.set_status(NodeStatus.unhealthy)
147 49457ba Dawsh 2y ago raise
148 49457ba Dawsh 2y ago else:
149 49457ba Dawsh 2y ago self.set_status(NodeStatus.healthy)
150 512d183 Mrdn 2y ago
151 512d183 Mrdn 2y ago async def get_xray_config(self):
152 512d183 Mrdn 2y ago response = await self._stub.FetchXrayConfig(Empty())
153 512d183 Mrdn 2y ago return response.configuration