diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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")