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