1import asyncio
2from collections import defaultdict
3from datetime import datetime
4
5from grpclib import GRPCError
6from sqlalchemy import and_, select, insert, update, bindparam
7
8from app import marznode
9from app.db import GetDB
10from app.db.models import NodeUsage, NodeUserUsage, User
11from app.marznode import MarzNodeBase
12
13
14def record_user_stats(params: list, node_id: int, consumption_factor: int = 1):
15 if not params:
16 return
17
18 created_at = datetime.fromisoformat(datetime.utcnow().strftime("%Y-%m-%dT%H:00:00"))
19
20 with GetDB() as db:
21 # make user usage row if doesn't exist
22 select_stmt = select(NodeUserUsage.user_id).where(
23 and_(
24 NodeUserUsage.node_id == node_id, NodeUserUsage.created_at == created_at
25 )
26 )
27 existings = [r[0] for r in db.execute(select_stmt).fetchall()]
28 uids_to_insert = set()
29
30 for p in params:
31 uid = p["uid"]
32 if uid in existings:
33 continue
34 uids_to_insert.add(uid)
35
36 if uids_to_insert:
37 stmt = insert(NodeUserUsage).values(
38 user_id=bindparam("uid"),
39 created_at=created_at,
40 node_id=node_id,
41 used_traffic=0,
42 )
43 db.execute(stmt, [{"uid": uid} for uid in uids_to_insert])
44
45 # record
46 stmt = (
47 update(NodeUserUsage)
48 .values(
49 used_traffic=NodeUserUsage.used_traffic
50 + bindparam("value") * consumption_factor
51 )
52 .where(
53 and_(
54 NodeUserUsage.user_id == bindparam("uid"),
55 NodeUserUsage.node_id == node_id,
56 NodeUserUsage.created_at == created_at,
57 )
58 )
59 )
60 db.connection().execute(
61 stmt, params, execution_options={"synchronize_session": None}
62 )
63 db.commit()
64
65
66def record_node_stats(params: dict, node_id: int):
67 if not params:
68 return
69
70 created_at = datetime.fromisoformat(datetime.utcnow().strftime("%Y-%m-%dT%H:00:00"))
71
72 with GetDB() as db:
73
74 # make node usage row if doesn't exist
75 select_stmt = select(NodeUsage.node_id).where(
76 and_(NodeUsage.node_id == node_id, NodeUsage.created_at == created_at)
77 )
78 notfound = db.execute(select_stmt).first() is None
79 if notfound:
80 stmt = insert(NodeUsage).values(
81 created_at=created_at, node_id=node_id, uplink=0, downlink=0
82 )
83 db.execute(stmt)
84 db.commit()
85
86 # record
87 stmt = (
88 update(NodeUsage)
89 .values(
90 uplink=NodeUsage.uplink + bindparam("up"),
91 downlink=NodeUsage.downlink + bindparam("down"),
92 )
93 .where(
94 and_(NodeUsage.node_id == node_id, NodeUsage.created_at == created_at)
95 )
96 )
97
98 db.execute(stmt, params)
99
100
101async def get_users_stats(node_id: int, node: MarzNodeBase) -> tuple[int, list[dict]]:
102 try:
103 params = list()
104 for stat in await node.fetch_users_stats():
105 if stat.usage:
106 params.append({"uid": stat.uid, "value": stat.usage})
107 return node_id, params
108 except:
109 return node_id, []
110
111
112async def record_user_usages():
113 # usage_coefficient = {None: 1} # default usage coefficient for the main api instance
114
115 results = await asyncio.gather(
116 *[get_users_stats(node_id, node) for node_id, node in marznode.nodes.items()]
117 )
118 api_params = {node_id: params for node_id, params in list(results)}
119
120 users_usage = defaultdict(int)
121 for node_id, params in api_params.items():
122 coefficient = (
123 node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1
124 )
125 for param in params:
126 users_usage[param["uid"]] += (
127 param["value"] * coefficient
128 ) # apply the usage coefficient
129
130 users_usage = list(
131 {"id": uid, "value": value} for uid, value in users_usage.items()
132 )
133 if not users_usage:
134 return
135
136 # record users usage
137 with GetDB() as db:
138 stmt = update(User).values(
139 used_traffic=User.used_traffic + bindparam("value"),
140 online_at=datetime.utcnow(),
141 )
142
143 db.execute(stmt, users_usage, execution_options={"synchronize_session": None})
144 db.commit()
145
146 for node_id, params in api_params.items():
147 record_user_stats(
148 params,
149 node_id,
150 node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1,
151 )