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.1.3
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 f6a32ea Dawsh 2y ago from .marznode_pb2 import UserData, UsersData, Empty, User, Inbound
13 f6a32ea Dawsh 2y ago from ..models.node import NodeStatus
14 f6a32ea Dawsh 2y ago
15 f6a32ea Dawsh 2y ago logger = logging.getLogger(__name__)
16 f6a32ea Dawsh 2y ago
17 f6a32ea Dawsh 2y ago channel_options = [
18 f6a32ea Dawsh 2y ago ("grpc.keepalive_time_ms", 8000),
19 f6a32ea Dawsh 2y ago ("grpc.keepalive_timeout_ms", 5000),
20 f6a32ea Dawsh 2y ago ("grpc.http2.max_pings_without_data", 0),
21 f6a32ea Dawsh 2y ago ("grpc.keepalive_permit_without_calls", 1),
22 f6a32ea Dawsh 2y ago ("grpc.max_connection_idle_ms", INT_MAX),
23 f6a32ea Dawsh 2y ago ("grpc.client_idle_timeout_ms", INT_MAX),
24 f6a32ea Dawsh 2y ago ("grpc.max_connection_age_ms", INT_MAX)
25 f6a32ea Dawsh 2y ago ]
26 f6a32ea Dawsh 2y ago
27 f6a32ea Dawsh 2y ago
28 f6a32ea Dawsh 2y ago class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB):
29 f6a32ea Dawsh 2y ago def __init__(self, node_id: int, address: str, port: int, usage_coefficient: int = 1):
30 f6a32ea Dawsh 2y ago self.id = node_id
31 f6a32ea Dawsh 2y ago self._address = address
32 f6a32ea Dawsh 2y ago self._port = port
33 f6a32ea Dawsh 2y ago
34 f6a32ea Dawsh 2y ago self._channel = insecure_channel(f"{self._address}:{self._port}", channel_options)
35 f6a32ea Dawsh 2y ago self._stub = MarzServiceStub(self._channel)
36 f6a32ea Dawsh 2y ago asyncio.create_task(self._monitor_channel())
37 f6a32ea Dawsh 2y ago self._streaming_task = None
38 f6a32ea Dawsh 2y ago
39 f6a32ea Dawsh 2y ago self._updates_queue = asyncio.Queue(1)
40 f6a32ea Dawsh 2y ago self.synced = False
41 f6a32ea Dawsh 2y ago self.usage_coefficient = usage_coefficient
42 f6a32ea Dawsh 2y ago atexit.register(self._close_channel)
43 f6a32ea Dawsh 2y ago
44 f6a32ea Dawsh 2y ago def _close_channel(self):
45 f6a32ea Dawsh 2y ago asyncio.run(self._channel.close())
46 f6a32ea Dawsh 2y ago
47 f6a32ea Dawsh 2y ago async def _monitor_channel(self):
48 f6a32ea Dawsh 2y ago try:
49 f6a32ea Dawsh 2y ago await asyncio.wait_for(self._channel.channel_ready(), timeout=5)
50 f6a32ea Dawsh 2y ago except TimeoutError:
51 f6a32ea Dawsh 2y ago logger.info("timeout for node, id: %i", self.id)
52 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.unhealthy, "timeout")
53 f6a32ea Dawsh 2y ago while state := self._channel.get_state():
54 f6a32ea Dawsh 2y ago logger.debug("node %i state: %s", self.id, state.value)
55 f6a32ea Dawsh 2y ago if state == ChannelConnectivity.READY:
56 f6a32ea Dawsh 2y ago await self._sync()
57 f6a32ea Dawsh 2y ago self._streaming_task = asyncio.create_task(self._stream_user_updates())
58 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.healthy)
59 f6a32ea Dawsh 2y ago logger.info("Connected to node %i", self.id)
60 f6a32ea Dawsh 2y ago else:
61 f6a32ea Dawsh 2y ago self.synced = False
62 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.unhealthy)
63 f6a32ea Dawsh 2y ago if self._streaming_task:
64 f6a32ea Dawsh 2y ago self._streaming_task.cancel()
65 f6a32ea Dawsh 2y ago
66 f6a32ea Dawsh 2y ago await self._channel.wait_for_state_change(state)
67 f6a32ea Dawsh 2y ago
68 f6a32ea Dawsh 2y ago async def _stream_user_updates(self):
69 f6a32ea Dawsh 2y ago logger.debug("opened the stream")
70 f6a32ea Dawsh 2y ago stream = self._stub.SyncUsers()
71 f6a32ea Dawsh 2y ago while True:
72 f6a32ea Dawsh 2y ago user_update = await self._updates_queue.get()
73 f6a32ea Dawsh 2y ago logger.debug("got something from queue")
74 f6a32ea Dawsh 2y ago user = user_update["user"]
75 f6a32ea Dawsh 2y ago await stream.write(
76 f6a32ea Dawsh 2y ago UserData(user=User(id=user.id, username=user.username, key=user.key),
77 f6a32ea Dawsh 2y ago inbounds=[Inbound(tag=t) for t in user_update["inbounds"]]))
78 f6a32ea Dawsh 2y ago
79 f6a32ea Dawsh 2y ago async def update_user(self, user, inbounds: set[str] | None = None):
80 f6a32ea Dawsh 2y ago if inbounds is None:
81 f6a32ea Dawsh 2y ago inbounds = set()
82 f6a32ea Dawsh 2y ago
83 f6a32ea Dawsh 2y ago if self.synced:
84 f6a32ea Dawsh 2y ago await self._updates_queue.put({"user": user, "inbounds": inbounds})
85 f6a32ea Dawsh 2y ago
86 f6a32ea Dawsh 2y ago async def _repopulate_users(self, users_data: list[dict]) -> None:
87 f6a32ea Dawsh 2y ago updates = [UserData(user=User(id=u["id"], username=u["username"], key=u["key"]),
88 f6a32ea Dawsh 2y ago inbounds=[Inbound(tag=t) for t in u["inbounds"]]) for u in users_data]
89 f6a32ea Dawsh 2y ago await self._stub.RepopulateUsers(UsersData(users_data=updates))
90 f6a32ea Dawsh 2y ago
91 f6a32ea Dawsh 2y ago async def fetch_users_stats(self):
92 f6a32ea Dawsh 2y ago response = await self._stub.FetchUsersStats(Empty())
93 f6a32ea Dawsh 2y ago return response.users_stats
94 f6a32ea Dawsh 2y ago
95 f6a32ea Dawsh 2y ago async def _fetch_inbounds(self) -> list:
96 f6a32ea Dawsh 2y ago response = await self._stub.FetchInbounds(Empty())
97 f6a32ea Dawsh 2y ago return response.inbounds
98 f6a32ea Dawsh 2y ago
99 f6a32ea Dawsh 2y ago async def _sync(self):
100 f6a32ea Dawsh 2y ago inbounds = await self._fetch_inbounds()
101 f6a32ea Dawsh 2y ago self.store_inbounds(inbounds)
102 f6a32ea Dawsh 2y ago users = self.list_users()
103 f6a32ea Dawsh 2y ago await self._repopulate_users(users)
104 f6a32ea Dawsh 2y ago self.synced = True