diff --git a/app/db/migrations/versions/058eb11c5b17_user_reset_in_table.py b/app/db/migrations/versions/058eb11c5b17_user_reset_in_table.py index 0f64b74..8f18a01 100644 --- a/app/db/migrations/versions/058eb11c5b17_user_reset_in_table.py +++ b/app/db/migrations/versions/058eb11c5b17_user_reset_in_table.py @@ -20,9 +20,23 @@ def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### op.drop_table("user_usage_logs") op.add_column( - "users", sa.Column("lifetime_used_traffic", sa.BigInteger(), nullable=True) + "users", + sa.Column( + "lifetime_used_traffic", + sa.BigInteger(), + server_default=sa.sql.null(), + nullable=True, + ), + ) + op.add_column( + "users", + sa.Column( + "traffic_reset_at", + sa.DateTime(), + server_default=sa.sql.null(), + nullable=True, + ), ) - op.add_column("users", sa.Column("traffic_reset_at", sa.DateTime(), nullable=True)) # ### end Alembic commands ### diff --git a/app/marznode/grpcio.py b/app/marznode/grpcio.py index 8d5202a..e3889db 100644 --- a/app/marznode/grpcio.py +++ b/app/marznode/grpcio.py @@ -1,14 +1,13 @@ import asyncio import atexit import logging -from _testcapi import INT_MAX -from grpc import ChannelConnectivity +from _testcapi import INT_MAX +from grpc import ChannelConnectivity, RpcError from grpc.aio import insecure_channel from .base import MarzNodeBase from .database import MarzNodeDB -from .marznode_pb2_grpc import MarzServiceStub from .marznode_pb2 import ( UserData, UsersData, @@ -18,6 +17,7 @@ from .marznode_pb2 import ( XrayLogsRequest, XrayConfig, ) +from .marznode_pb2_grpc import MarzServiceStub from ..models.node import NodeStatus logger = logging.getLogger(__name__) @@ -48,7 +48,7 @@ class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB): asyncio.create_task(self._monitor_channel()) self._streaming_task = None - self._updates_queue = asyncio.Queue(1) + self._updates_queue = asyncio.Queue(5) self.synced = False self.usage_coefficient = usage_coefficient atexit.register(self._close_channel) @@ -64,16 +64,19 @@ class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB): self.set_status(NodeStatus.unhealthy, "timeout") while state := self._channel.get_state(): logger.debug("node %i state: %s", self.id, state.value) - if state == ChannelConnectivity.READY: + try: + if state != ChannelConnectivity.READY: + raise RpcError await self._sync() self._streaming_task = asyncio.create_task(self._stream_user_updates()) - self.set_status(NodeStatus.healthy) - logger.info("Connected to node %i", self.id) - else: + except RpcError: self.synced = False self.set_status(NodeStatus.unhealthy) if self._streaming_task: self._streaming_task.cancel() + else: + self.set_status(NodeStatus.healthy) + logger.info("Connected to node %i", self.id) await self._channel.wait_for_state_change(state) @@ -84,12 +87,17 @@ class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB): user_update = await self._updates_queue.get() logger.debug("got something from queue") user = user_update["user"] - await stream.write( - UserData( - user=User(id=user.id, username=user.username, key=user.key), - inbounds=[Inbound(tag=t) for t in user_update["inbounds"]], + try: + await stream.write( + UserData( + user=User(id=user.id, username=user.username, key=user.key), + inbounds=[Inbound(tag=t) for t in user_update["inbounds"]], + ) ) - ) + except RpcError: + self.synced = False + self.set_status(NodeStatus.unhealthy) + return async def update_user(self, user, inbounds: set[str] | None = None): if inbounds is None: @@ -130,8 +138,15 @@ class MarzNodeGRPCIO(MarzNodeBase, MarzNodeDB): yield response.line async def restart_xray(self, config: str): - await self._stub.RestartXray(XrayConfig(configuration=config)) - await self._sync() + try: + await self._stub.RestartXray(XrayConfig(configuration=config)) + await self._sync() + except RpcError: + self.synced = False + self.set_status(NodeStatus.unhealthy) + raise + else: + self.set_status(NodeStatus.healthy) async def get_xray_config(self): response = await self._stub.FetchXrayConfig(Empty()) diff --git a/app/marznode/operations.py b/app/marznode/operations.py index 5060903..1443a64 100644 --- a/app/marznode/operations.py +++ b/app/marznode/operations.py @@ -1,3 +1,4 @@ +import asyncio from collections import defaultdict from typing import TYPE_CHECKING @@ -30,7 +31,9 @@ async def update_user(user: "DBUser", old_inbounds: set | None = None): for node_id, tags in node_inbounds.items(): if marznode.nodes.get(node_id): - await marznode.nodes[node_id].update_user(user=user, inbounds=tags) + asyncio.create_task( + marznode.nodes[node_id].update_user(user=user, inbounds=tags) + ) async def remove_user(user: "DBUser"): @@ -49,7 +52,12 @@ async def remove_node(node_id: int): async def add_node(db_node, certificate): await remove_node(db_node.id) if db_node.connection_backend == NodeConnectionBackend.grpcio: - node = MarzNodeGRPCIO(db_node.id, db_node.address, db_node.port) + node = MarzNodeGRPCIO( + db_node.id, + db_node.address, + db_node.port, + usage_coefficient=db_node.usage_coefficient, + ) else: node = MarzNodeGRPCLIB( db_node.id, @@ -57,6 +65,7 @@ async def add_node(db_node, certificate): db_node.port, certificate.key, certificate.certificate, + usage_coefficient=db_node.usage_coefficient, ) marznode.nodes[db_node.id] = node diff --git a/app/models/node.py b/app/models/node.py index 87d1a61..4928d92 100644 --- a/app/models/node.py +++ b/app/models/node.py @@ -24,7 +24,7 @@ class NodeBase(BaseModel): id: Optional[int] = Field(None) name: str address: str - port: int = 62050 + port: int = 53042 connection_backend: NodeConnectionBackend = Field( default=NodeConnectionBackend.grpclib ) @@ -58,7 +58,7 @@ class NodeModify(Node): port: Optional[int] = Field(None) connection_backend: Optional[NodeConnectionBackend] = Field(None) status: Optional[NodeStatus] = Field(None) - usage_coefficient: Optional[float] = Field(None) + usage_coefficient: Optional[float] = Field(None, ge=0) model_config = ConfigDict( json_schema_extra={ "example": { diff --git a/app/routes/service.py b/app/routes/service.py index a990a0e..60f3dfe 100644 --- a/app/routes/service.py +++ b/app/routes/service.py @@ -1,7 +1,6 @@ import sqlalchemy from fastapi import APIRouter, Depends, Query from fastapi import HTTPException -from fastapi_pagination.customization import UseParamsFields, CustomizedPage from fastapi_pagination.ext.sqlalchemy import paginate from fastapi_pagination.links import Page @@ -62,7 +61,6 @@ async def modify_service(id: int, modification: ServiceModify, db: DBDep): - **inbounds** list of inbound ids. if not specified no change will be applied; in case of an empty list all inbounds would be removed. """ - # TODO: Update all affected users in nodes dbservice = crud.get_service(db, id) if not dbservice: raise HTTPException(status_code=404, detail="Service not found")