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-alpha1
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 from _testcapi import INT_MAX
5 f6a32ea Dawsh 2y ago
6 f6a32ea Dawsh 2y ago from grpc import ChannelConnectivity
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 f6a32ea Dawsh 2y ago from .marznode_pb2_grpc import MarzServiceStub
12 01cdf2e Dawsh 2y ago from .marznode_pb2 import (
13 01cdf2e Dawsh 2y ago UserData,
14 01cdf2e Dawsh 2y ago UsersData,
15 01cdf2e Dawsh 2y ago Empty,
16 01cdf2e Dawsh 2y ago User,
17 01cdf2e Dawsh 2y ago Inbound,
18 01cdf2e Dawsh 2y ago XrayLogsRequest,
19 01cdf2e Dawsh 2y ago XrayConfig,
20 01cdf2e Dawsh 2y ago )
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 f6a32ea Dawsh 2y ago self._updates_queue = asyncio.Queue(1)
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 f6a32ea Dawsh 2y ago if state == ChannelConnectivity.READY:
68 f6a32ea Dawsh 2y ago await self._sync()
69 f6a32ea Dawsh 2y ago self._streaming_task = asyncio.create_task(self._stream_user_updates())
70 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.healthy)
71 f6a32ea Dawsh 2y ago logger.info("Connected to node %i", self.id)
72 f6a32ea Dawsh 2y ago else:
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 f6a32ea Dawsh 2y ago
78 f6a32ea Dawsh 2y ago await self._channel.wait_for_state_change(state)
79 f6a32ea Dawsh 2y ago
80 f6a32ea Dawsh 2y ago async def _stream_user_updates(self):
81 f6a32ea Dawsh 2y ago logger.debug("opened the stream")
82 f6a32ea Dawsh 2y ago stream = self._stub.SyncUsers()
83 f6a32ea Dawsh 2y ago while True:
84 f6a32ea Dawsh 2y ago user_update = await self._updates_queue.get()
85 f6a32ea Dawsh 2y ago logger.debug("got something from queue")
86 f6a32ea Dawsh 2y ago user = user_update["user"]
87 f6a32ea Dawsh 2y ago await stream.write(
88 01cdf2e Dawsh 2y ago UserData(
89 01cdf2e Dawsh 2y ago user=User(id=user.id, username=user.username, key=user.key),
90 01cdf2e Dawsh 2y ago inbounds=[Inbound(tag=t) for t in user_update["inbounds"]],
91 01cdf2e Dawsh 2y ago )
92 01cdf2e Dawsh 2y ago )
93 f6a32ea Dawsh 2y ago
94 f6a32ea Dawsh 2y ago async def update_user(self, user, inbounds: set[str] | None = None):
95 f6a32ea Dawsh 2y ago if inbounds is None:
96 f6a32ea Dawsh 2y ago inbounds = set()
97 f6a32ea Dawsh 2y ago
98 f6a32ea Dawsh 2y ago if self.synced:
99 f6a32ea Dawsh 2y ago await self._updates_queue.put({"user": user, "inbounds": inbounds})
100 f6a32ea Dawsh 2y ago
101 f6a32ea Dawsh 2y ago async def _repopulate_users(self, users_data: list[dict]) -> None:
102 01cdf2e Dawsh 2y ago updates = [
103 01cdf2e Dawsh 2y ago UserData(
104 01cdf2e Dawsh 2y ago user=User(id=u["id"], username=u["username"], key=u["key"]),
105 01cdf2e Dawsh 2y ago inbounds=[Inbound(tag=t) for t in u["inbounds"]],
106 01cdf2e Dawsh 2y ago )
107 01cdf2e Dawsh 2y ago for u in users_data
108 01cdf2e Dawsh 2y ago ]
109 f6a32ea Dawsh 2y ago await self._stub.RepopulateUsers(UsersData(users_data=updates))
110 f6a32ea Dawsh 2y ago
111 f6a32ea Dawsh 2y ago async def fetch_users_stats(self):
112 f6a32ea Dawsh 2y ago response = await self._stub.FetchUsersStats(Empty())
113 f6a32ea Dawsh 2y ago return response.users_stats
114 f6a32ea Dawsh 2y ago
115 f6a32ea Dawsh 2y ago async def _fetch_inbounds(self) -> list:
116 f6a32ea Dawsh 2y ago response = await self._stub.FetchInbounds(Empty())
117 f6a32ea Dawsh 2y ago return response.inbounds
118 f6a32ea Dawsh 2y ago
119 f6a32ea Dawsh 2y ago async def _sync(self):
120 f6a32ea Dawsh 2y ago inbounds = await self._fetch_inbounds()
121 f6a32ea Dawsh 2y ago self.store_inbounds(inbounds)
122 f6a32ea Dawsh 2y ago users = self.list_users()
123 f6a32ea Dawsh 2y ago await self._repopulate_users(users)
124 f6a32ea Dawsh 2y ago self.synced = True
125 512d183 Mrdn 2y ago
126 512d183 Mrdn 2y ago async def get_logs(self, include_buffer=True):
127 512d183 Mrdn 2y ago async for response in self._stub.StreamXrayLogs(
128 01cdf2e Dawsh 2y ago XrayLogsRequest(include_buffer=include_buffer)
129 512d183 Mrdn 2y ago ):
130 512d183 Mrdn 2y ago yield response.line
131 512d183 Mrdn 2y ago
132 512d183 Mrdn 2y ago async def restart_xray(self, config: str):
133 512d183 Mrdn 2y ago await self._stub.RestartXray(XrayConfig(configuration=config))
134 512d183 Mrdn 2y ago await self._sync()
135 512d183 Mrdn 2y ago
136 512d183 Mrdn 2y ago async def get_xray_config(self):
137 512d183 Mrdn 2y ago response = await self._stub.FetchXrayConfig(Empty())
138 512d183 Mrdn 2y ago return response.configuration