diff --git a/client.py b/client.py index f31967ab..4cd8b744 100644 --- a/client.py +++ b/client.py @@ -47,7 +47,9 @@ async def add_client( return {"status": "failed", "error": str(e)} -async def extend_client_key(xui, email: str, new_expiry_time: int, client_id: str): +async def extend_client_key( + xui, email: str, new_expiry_time: int, client_id: str, total_gb: int +): """ Функция для обновления срока действия ключа клиента по email. """ @@ -71,12 +73,7 @@ async def extend_client_key(xui, email: str, new_expiry_time: int, client_id: st client.expiry_time = new_expiry_time client.flow = "xtls-rprx-vision" client.sub_id = email - - if TOTAL_GB > 0: - client.total_gb = TOTAL_GB - logger.info( - f"Установлен объем трафика для клиента {client.email}: {TOTAL_GB} ГБ" - ) + client.total_gb = total_gb await xui.client.update(client.id, client) logger.info( @@ -91,23 +88,26 @@ async def delete_client( xui, email: str, client_id: str, -): +) -> bool: """ - Функция для удаления клиента с сервера 3x-ui + Функция для удаления клиента с сервера 3x-ui. + Возвращает True при успешном удалении, иначе False. """ await xui.login() try: client = await xui.client.get_by_email(email) - client.id = client_id if not client: - logger.warning(f"Клиент с email {client_id} не найден.") - return + logger.warning(f"Клиент с email {email} и ID {client_id} не найден.") + return False + client.id = client_id inbound_id = 1 - await xui.client.delete(inbound_id, client.id) + await xui.client.delete(inbound_id, client.id) logger.info(f"Клиент с ID {client_id} был удален успешно.") + return True except Exception as e: logger.error(f"Ошибка при удалении клиента с ID {client_id}: {e}") + return False diff --git a/handlers/admin/admin_user_editor.py b/handlers/admin/admin_user_editor.py index 3d10cab5..f23c2a76 100644 --- a/handlers/admin/admin_user_editor.py +++ b/handlers/admin/admin_user_editor.py @@ -10,11 +10,11 @@ from aiogram.utils.keyboard import InlineKeyboardBuilder from loguru import logger from bot import bot -from config import DATABASE_URL, SERVERS +from config import CLUSTERS, DATABASE_URL from database import get_client_id_by_email, restore_trial, update_key_expiry from filters.admin import IsAdminFilter from handlers.admin.admin_panel import back_to_admin_menu -from handlers.keys.key_utils import delete_key_from_server, renew_server_key +from handlers.keys.key_utils import delete_key_from_cluster, renew_key_in_cluster from handlers.utils import sanitize_key_name router = Router() @@ -178,9 +178,14 @@ async def process_key_edit(callback_query: CallbackQuery): key = record["key"] expiry_time = record["expiry_time"] server_id = record["server_id"] - server_name = SERVERS.get(server_id, {}).get( - "name", "Неизвестный сервер" - ) + server_name = "Неизвестный сервер" + + for cluster in CLUSTERS.values(): + if server_id in cluster: + server_name = cluster[server_id].get( + "name", "Неизвестный сервер" + ) + break expiry_date = datetime.utcfromtimestamp(expiry_time / 1000) current_date = datetime.utcnow() @@ -284,7 +289,12 @@ async def handle_key_name_input(message: types.Message, state: FSMContext): key = record["key"] expiry_time = record["expiry_time"] server_id = record["server_id"] - server_name = SERVERS.get(server_id, {}).get("name", "Неизвестный сервер") + server_name = "Неизвестный сервер" + + for cluster in CLUSTERS.values(): + if server_id in cluster: + server_name = cluster[server_id].get("name", "Неизвестный сервер") + break expiry_date = datetime.utcfromtimestamp(expiry_time / 1000).strftime( "%d %B %Y" @@ -370,10 +380,12 @@ async def handle_expiry_time_input(message: types.Message, state: FSMContext): async def update_key_on_all_servers(): tasks = [] - for server_id in SERVERS: + for cluster_id in CLUSTERS: # Разворачиваем на cluster_id и cluster tasks.append( asyncio.create_task( - renew_server_key(server_id, email, client_id, expiry_time) + renew_key_in_cluster( + cluster_id, email, client_id, expiry_time + ) ) ) await asyncio.gather(*tasks) @@ -468,15 +480,15 @@ async def process_callback_confirm_delete(callback_query: types.CallbackQuery): InlineKeyboardButton(text="⬅️ Назад", callback_data="view_keys") ) - async def delete_key_from_servers(): + async def delete_key_from_servers(email, client_id): tasks = [] - for server_id in SERVERS: + for cluster_id in CLUSTERS: tasks.append( - delete_key_from_server(server_id, email, client_id) + delete_key_from_cluster(cluster_id, email, client_id) ) await asyncio.gather(*tasks) - await delete_key_from_servers() + await delete_key_from_servers(email, client_id) await delete_key_from_db(client_id) await bot.edit_message_text( diff --git a/handlers/keys/key_management.py b/handlers/keys/key_management.py index ba6a8c92..1a069c4f 100644 --- a/handlers/keys/key_management.py +++ b/handlers/keys/key_management.py @@ -10,13 +10,13 @@ from aiogram.types import CallbackQuery, InlineKeyboardButton, InlineKeyboardMar from loguru import logger from bot import bot, dp -from config import CONNECT_ANDROID, CONNECT_IOS, DATABASE_URL, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK, SERVERS +from config import CONNECT_ANDROID, CONNECT_IOS, DATABASE_URL, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK from database import add_connection, get_balance, store_key, update_balance from handlers.instructions.instructions import send_instructions -from handlers.keys.key_utils import create_key_on_server +from handlers.keys.key_utils import create_key_on_cluster from handlers.profile import process_callback_view_profile from handlers.texts import KEY, KEY_TRIAL, NULL_BALANCE, RENEWAL_PLANS, key_message_success -from handlers.utils import sanitize_key_name +from handlers.utils import get_least_loaded_cluster, sanitize_key_name router = Router() @@ -238,15 +238,20 @@ async def handle_key_name_input(message: Message, state: FSMContext): ) try: + least_loaded_cluster = await get_least_loaded_cluster() + tasks = [] - for server_id in SERVERS: - tasks.append( - asyncio.create_task( - create_key_on_server( - server_id, tg_id, client_id, email, expiry_timestamp - ) + tasks.append( + asyncio.create_task( + create_key_on_cluster( + least_loaded_cluster, + tg_id, + client_id, + email, + expiry_timestamp, ) ) + ) await asyncio.gather(*tasks) @@ -267,7 +272,7 @@ async def handle_key_name_input(message: Message, state: FSMContext): logger.info(f"Storing key for user {tg_id} in the database.") await store_key( - tg_id, client_id, email, expiry_timestamp, public_link, "all_servers" + tg_id, client_id, email, expiry_timestamp, public_link, least_loaded_cluster ) except Exception as e: diff --git a/handlers/keys/key_utils.py b/handlers/keys/key_utils.py index 23abfd35..01fa97e1 100644 --- a/handlers/keys/key_utils.py +++ b/handlers/keys/key_utils.py @@ -1,61 +1,88 @@ +import asyncio + import asyncpg from loguru import logger from py3xui import AsyncApi from client import add_client, delete_client, extend_client_key -from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL, SERVERS +from config import ADMIN_PASSWORD, ADMIN_USERNAME, CLUSTERS, DATABASE_URL -async def create_key_on_server(server_id, tg_id, client_id, email, expiry_timestamp): +async def create_key_on_cluster(cluster_id, tg_id, client_id, email, expiry_timestamp): try: - xui = AsyncApi( - SERVERS[server_id]["API_URL"], - username=ADMIN_USERNAME, - password=ADMIN_PASSWORD, - ) + tasks = [] + cluster = CLUSTERS.get(cluster_id) - conn = await asyncpg.connect(DATABASE_URL) - existing_key = await conn.fetchrow("SELECT 1 FROM keys WHERE email = $1", email) + if not cluster: + raise ValueError(f"Кластер с ID {cluster_id} не найден.") - if existing_key: - raise ValueError(f"Email {email} уже существует в базе данных.") + for server_info in cluster.values(): + xui = AsyncApi( + server_info["API_URL"], + username=ADMIN_USERNAME, + password=ADMIN_PASSWORD, + ) - await add_client( - xui, - client_id, - email, - tg_id, - limit_ip=1, - total_gb=0, - expiry_time=expiry_timestamp, - enable=True, - flow="xtls-rprx-vision", - ) + conn = await asyncpg.connect(DATABASE_URL) + existing_key = await conn.fetchrow( + "SELECT 1 FROM keys WHERE email = $1", email + ) - await conn.close() + if existing_key: + raise ValueError(f"Email {email} уже существует в базе данных.") + + tasks.append( + add_client( + xui, + client_id, + email, + tg_id, + limit_ip=1, + total_gb=0, + expiry_time=expiry_timestamp, + enable=True, + flow="xtls-rprx-vision", + ) + ) + await conn.close() + + await asyncio.gather(*tasks) except Exception as e: - logger.error(f"Ошибка на сервере {server_id}: {e}") + logger.error(f"Ошибка при создании ключа: {e}") raise e -async def renew_server_key(server_id, email, client_id, new_expiry_time): +async def renew_key_in_cluster(cluster_id, email, client_id, new_expiry_time, total_gb): """ - Функция для продления срока действия ключа на сервере и сброса трафика, если необходимо. + Функция для продления срока действия ключа на всех серверах в кластере + и обновления объема трафика. """ try: - xui = AsyncApi( - SERVERS[server_id]["API_URL"], - username=ADMIN_USERNAME, - password=ADMIN_PASSWORD, - ) + cluster = CLUSTERS.get(cluster_id) - await extend_client_key(xui, email, new_expiry_time, client_id) + if not cluster: + raise ValueError(f"Кластер с ID {cluster_id} не найден.") + + tasks = [] + for server_info in cluster.values(): + xui = AsyncApi( + server_info["API_URL"], + username=ADMIN_USERNAME, + password=ADMIN_PASSWORD, + ) + + tasks.append( + extend_client_key(xui, email, new_expiry_time, client_id, total_gb) + ) + + await asyncio.gather(*tasks) except Exception as e: logger.error( - f"Не удалось продлить ключ {client_id} и сбросить трафик на сервере {server_id}: {e}" + f"Не удалось продлить ключ {client_id} в кластере {cluster_id}: {e}" ) + raise e async def delete_key_from_db(client_id): @@ -69,43 +96,70 @@ async def delete_key_from_db(client_id): await conn.close() -async def delete_key_from_server(server_id, email, client_id): - """Удаление ключа с сервера""" +async def delete_key_from_cluster(cluster_id, email, client_id): + """Удаление ключа с серверов в кластере""" try: - xui = AsyncApi( - SERVERS[server_id]["API_URL"], - username=ADMIN_USERNAME, - password=ADMIN_PASSWORD, - ) + cluster = CLUSTERS.get(cluster_id) - await delete_client(xui, email, client_id) + if not cluster: + raise ValueError(f"Кластер с ID {cluster_id} не найден.") - except Exception as e: - logger.error(f"Не удалось удалить ключ {client_id} на сервере {server_id}: {e}") + tasks = [] + for server_info in cluster.values(): + xui = AsyncApi( + server_info["API_URL"], + username=ADMIN_USERNAME, + password=ADMIN_PASSWORD, + ) + tasks.append(delete_client(xui, email, client_id)) -async def update_key_on_server(tg_id, client_id, email, expiry_time, server_id): - try: - xui = AsyncApi( - SERVERS[server_id]["API_URL"], - username=ADMIN_USERNAME, - password=ADMIN_PASSWORD, - ) - await add_client( - xui, - client_id, - email, - tg_id, - limit_ip=1, - total_gb=0, - expiry_time=expiry_time, - enable=True, - flow="xtls-rprx-vision", - ) - - logger.info(f"Ключ успешно обновлен на сервере {server_id} для {client_id}") + await asyncio.gather(*tasks) except Exception as e: logger.error( - f"Ошибка при обновлении ключа на сервере {server_id} для {client_id}: {e}" + f"Не удалось удалить ключ {client_id} в кластере {cluster_id}: {e}" ) + raise e + + +async def update_key_on_cluster(tg_id, client_id, email, expiry_time, cluster_id): + try: + cluster = CLUSTERS.get(cluster_id) + + if not cluster: + raise ValueError(f"Кластер с ID {cluster_id} не найден.") + + tasks = [] + for server_info in cluster.values(): + xui = AsyncApi( + server_info["API_URL"], + username=ADMIN_USERNAME, + password=ADMIN_PASSWORD, + ) + + tasks.append( + add_client( + xui, + client_id, + email, + tg_id, + limit_ip=1, + total_gb=0, + expiry_time=expiry_time, + enable=True, + flow="xtls-rprx-vision", + ) + ) + + await asyncio.gather(*tasks) + + logger.info( + f"Ключ успешно обновлен для {client_id} на всех серверах в кластере {cluster_id}" + ) + + except Exception as e: + logger.error( + f"Ошибка при обновлении ключа на серверах кластера {cluster_id} для {client_id}: {e}" + ) + raise e diff --git a/handlers/keys/keys.py b/handlers/keys/keys.py index 161359c2..bc4a1f00 100644 --- a/handlers/keys/keys.py +++ b/handlers/keys/keys.py @@ -9,11 +9,11 @@ from aiogram.types import BufferedInputFile from loguru import logger from bot import bot -from config import CONNECT_ANDROID, CONNECT_IOS, DATABASE_URL, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK, SERVERS +from config import CLUSTERS, CONNECT_ANDROID, CONNECT_IOS, DATABASE_URL, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK, TOTAL_GB from database import delete_key, get_balance, store_key, update_balance, update_key_expiry -from handlers.keys.key_utils import delete_key_from_db, delete_key_from_server, renew_server_key, update_key_on_server +from handlers.keys.key_utils import delete_key_from_cluster, delete_key_from_db, renew_key_in_cluster, update_key_on_cluster from handlers.texts import INSUFFICIENT_FUNDS_MSG, KEY_NOT_FOUND_MSG, NO_KEYS, PLAN_SELECTION_MSG, RENEWAL_PLANS, SUCCESS_RENEWAL_MSG, key_message -from handlers.utils import handle_error +from handlers.utils import get_least_loaded_cluster, handle_error locale.setlocale(locale.LC_TIME, "ru_RU.UTF-8") @@ -53,7 +53,7 @@ async def process_callback_view_keys(callback_query: types.CallbackQuery): inline_keyboard = types.InlineKeyboardMarkup(inline_keyboard=buttons) response_message = ( "🔑 Список ваших устройств\n\n" - "👆 Выберите устройство для управления подпиской:" + "👇 Выберите устройство для управления подпиской:" ) image_path = os.path.join(os.path.dirname(__file__), "pic_keys.jpg") @@ -167,7 +167,14 @@ async def process_callback_view_key(callback_query: types.CallbackQuery): expiry_time = record["expiry_time"] server_id = record["server_id"] - server_name = SERVERS.get(server_id, {}).get("name", "мультисервер") + if server_id == "все кластеры": + server_name = "мультикластер" + else: + cluster_name = CLUSTERS.get(server_id, {}).get( + "name", "кластер неизвестен" + ) + server_name = cluster_name + expiry_date = datetime.utcfromtimestamp(expiry_time / 1000) current_date = datetime.utcnow() time_left = expiry_date - current_date @@ -303,13 +310,14 @@ async def process_callback_update_subscription(callback_query: types.CallbackQue ) return + least_loaded_cluster_id = await get_least_loaded_cluster() + tasks = [] - for server_id in SERVERS: - tasks.append( - update_key_on_server( - tg_id, client_id, email, expiry_time, server_id - ) + tasks.append( + update_key_on_cluster( + tg_id, client_id, email, expiry_time, least_loaded_cluster_id ) + ) await asyncio.gather(*tasks) @@ -319,7 +327,7 @@ async def process_callback_update_subscription(callback_query: types.CallbackQue email, expiry_time, public_link, - server_id="все сервера", + server_id=least_loaded_cluster_id, ) try: @@ -527,9 +535,9 @@ async def process_callback_confirm_delete(callback_query: types.CallbackQuery): async def delete_key_from_servers(): try: tasks = [] - for server_id in SERVERS: + for cluster_id, cluster in CLUSTERS.items(): tasks.append( - delete_key_from_server(server_id, email, client_id) + delete_key_from_cluster(cluster_id, email, client_id) ) await asyncio.gather(*tasks) @@ -577,14 +585,10 @@ async def process_callback_renew_plan(callback_query: types.CallbackQuery): ) days_to_extend = 30 * int(plan) - try: - try: - await bot.delete_message( - chat_id=tg_id, message_id=callback_query.message.message_id - ) - except Exception: - pass + gb_multiplier = {"1": 1, "3": 3, "6": 6, "12": 12} + total_gb = TOTAL_GB * gb_multiplier.get(plan, 1) if TOTAL_GB > 0 else 0 + try: conn = await asyncpg.connect(DATABASE_URL) try: record = await conn.fetchrow( @@ -643,10 +647,10 @@ async def process_callback_renew_plan(callback_query: types.CallbackQuery): async def renew_key_on_servers(): tasks = [] - for server_id in SERVERS: + for cluster_id in CLUSTERS: task = asyncio.create_task( - renew_server_key( - server_id, email, client_id, new_expiry_time + renew_key_in_cluster( + cluster_id, email, client_id, new_expiry_time, total_gb ) ) tasks.append(task) diff --git a/handlers/keys/subscriptions.py b/handlers/keys/subscriptions.py index 734d6893..c49fc5e9 100644 --- a/handlers/keys/subscriptions.py +++ b/handlers/keys/subscriptions.py @@ -4,7 +4,7 @@ import aiohttp from aiohttp import web from loguru import logger -from config import SERVERS +from config import CLUSTERS async def fetch_url_content(url): @@ -50,9 +50,10 @@ async def handle_subscription(request): logger.info(f"Получен запрос на подписку для email: {email}") urls = [] - for server in SERVERS.values(): - server_subscription_url = f"{server['SUBSCRIPTION']}/{email}" - urls.append(server_subscription_url) + for cluster in CLUSTERS.values(): + for server in cluster.values(): + server_subscription_url = f"{server['SUBSCRIPTION']}/{email}" + urls.append(server_subscription_url) query_string = request.query_string logger.debug(f"Извлечен query string: {query_string}") diff --git a/handlers/keys/trial_key.py b/handlers/keys/trial_key.py index c4138b93..50d372a1 100644 --- a/handlers/keys/trial_key.py +++ b/handlers/keys/trial_key.py @@ -6,10 +6,10 @@ import asyncpg from py3xui import AsyncApi from client import add_client -from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL, PUBLIC_LINK, SERVERS, TRIAL_TIME +from config import ADMIN_PASSWORD, ADMIN_USERNAME, CLUSTERS, DATABASE_URL, PUBLIC_LINK, TRIAL_TIME from database import store_key from handlers.texts import INSTRUCTIONS -from handlers.utils import generate_random_email +from handlers.utils import generate_random_email, get_least_loaded_cluster async def create_trial_key(tg_id: int): @@ -39,13 +39,20 @@ async def generate_and_store_keys( conn = await asyncpg.connect(DATABASE_URL) try: current_time = datetime.utcnow() - expiry_time = current_time + timedelta(days={TRIAL_TIME}, hours=3) + expiry_time = current_time + timedelta(days=TRIAL_TIME, hours=3) expiry_timestamp = int(expiry_time.timestamp() * 1000) + least_loaded_cluster = await get_least_loaded_cluster() + tasks = [] - for server_id in SERVERS: + for server_id, server in CLUSTERS[least_loaded_cluster].items(): task = create_key_on_server( - server_id, client_id, email, tg_id, expiry_timestamp + least_loaded_cluster, + server_id, + client_id, + email, + tg_id, + expiry_timestamp, ) tasks.append(task) @@ -57,15 +64,15 @@ async def generate_and_store_keys( email, expiry_timestamp, public_link, - server_id="all_servers", + server_id=least_loaded_cluster, ) await conn.execute( """ - INSERT INTO connections (tg_id, trial) - VALUES ($1, 1) - ON CONFLICT (tg_id) - DO UPDATE SET trial = 1 + INSERT INTO connections (tg_id, trial) + VALUES ($1, 1) + ON CONFLICT (tg_id) + DO UPDATE SET trial = 1 """, tg_id, ) @@ -74,12 +81,19 @@ async def generate_and_store_keys( async def create_key_on_server( - server_id: str, client_id: str, email: str, tg_id: int, expiry_timestamp: int + cluster_id: str, + server_id: str, + client_id: str, + email: str, + tg_id: int, + expiry_timestamp: int, ): - """Создает ключ на сервере и возвращает результат.""" + """Создает ключ на сервере в указанном кластере и возвращает результат.""" xui = AsyncApi( - SERVERS[server_id]["API_URL"], username=ADMIN_USERNAME, password=ADMIN_PASSWORD + CLUSTERS[cluster_id][server_id]["API_URL"], + username=ADMIN_USERNAME, + password=ADMIN_PASSWORD, ) response = await add_client( diff --git a/handlers/notifications.py b/handlers/notifications.py index c98a71fc..0aa0bf77 100644 --- a/handlers/notifications.py +++ b/handlers/notifications.py @@ -7,7 +7,7 @@ from loguru import logger from py3xui import AsyncApi from client import delete_client, extend_client_key -from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL, SERVERS +from config import ADMIN_PASSWORD, ADMIN_USERNAME, CLUSTERS, DATABASE_URL from database import delete_key, get_balance, update_balance, update_key_expiry from handlers.texts import KEY_EXPIRY_10H, KEY_EXPIRY_24H, KEY_RENEWED, RENEWAL_PLANS @@ -259,23 +259,28 @@ async def handle_expired_keys(bot: Bot, conn: asyncpg.Connection, current_time: ) await update_key_expiry(client_id, new_expiry_time) - all_success = True - for server_id in SERVERS: - xui = AsyncApi( - SERVERS[server_id]["API_URL"], - username=ADMIN_USERNAME, - password=ADMIN_PASSWORD, - ) - success = await extend_client_key( - xui, email, new_expiry_time, client_id - ) - if not success: - all_success = False - logger.error( - f"Не удалось продлить ключ для пользователя {tg_id} на сервере {server_id}." + renewal_success = False + for cluster_id, cluster in CLUSTERS.items(): + for server_id, server in cluster.items(): + xui = AsyncApi( + server["API_URL"], + username=ADMIN_USERNAME, + password=ADMIN_PASSWORD, ) + success = await extend_client_key( + xui, email, new_expiry_time, client_id + ) + if success: + renewal_success = True + logger.info( + f"Ключ для пользователя {tg_id} успешно продлен на сервере {server_id} в кластере {cluster_id}." + ) + else: + logger.error( + f"Не удалось продлить ключ для пользователя {tg_id} на сервере {server_id} в кластере {cluster_id}." + ) - if all_success: + if renewal_success: try: await bot.send_message( tg_id, KEY_RENEWED, reply_markup=keyboard @@ -299,21 +304,14 @@ async def handle_expired_keys(bot: Bot, conn: asyncpg.Connection, current_time: await delete_key(client_id) - for server_id in SERVERS: - xui = AsyncApi( - SERVERS[server_id]["API_URL"], - username=ADMIN_USERNAME, - password=ADMIN_PASSWORD, - ) - success = await delete_client(xui, email, client_id) - if success: - logger.info( - f"Ключ для клиента {tg_id} успешно удален с сервера {server_id}." - ) - else: - logger.error( - f"Не удалось удалить ключ для клиента {tg_id} на сервере {server_id}." + for cluster_id, cluster in CLUSTERS.items(): + for server_id, server in cluster.items(): + xui = AsyncApi( + server["API_URL"], + username=ADMIN_USERNAME, + password=ADMIN_PASSWORD, ) + await delete_client(xui, email, client_id) except Exception as e: logger.error(f"Ошибка при обработке ключа для клиента {tg_id}: {e}") diff --git a/handlers/utils.py b/handlers/utils.py index 9ed382c6..98c9145e 100644 --- a/handlers/utils.py +++ b/handlers/utils.py @@ -1,10 +1,11 @@ import random import re +import asyncpg from loguru import logger from bot import bot -from config import SERVERS +from config import CLUSTERS, DATABASE_URL def sanitize_key_name(key_name: str) -> str: @@ -17,21 +18,52 @@ def generate_random_email(): return random_string -async def get_least_loaded_server(conn): - """Находит сервер с наименьшей загрузкой.""" - least_loaded_server_id = None - min_load_percentage = float("inf") +async def get_least_loaded_cluster(): + """ + Функция для получения кластера с наименьшей загрузкой (по количеству ключей). + Возвращает идентификатор кластера с наименьшей загрузкой или первый кластер из конфигурации, + если загруженность не определяется. В случае отсутствия кластеров с номером, возвращает 'cluster1'. + """ + cluster_loads = {} - for server_id, server in SERVERS.items(): - count = await conn.fetchval( - "SELECT COUNT(*) FROM keys WHERE server_id = $1", server_id - ) - percent_full = (count / 60) * 100 if count <= 60 else 100 - if percent_full < min_load_percentage: - min_load_percentage = percent_full - least_loaded_server_id = server_id + conn = await asyncpg.connect(DATABASE_URL) + try: + keys = await conn.fetch("SELECT * FROM keys") + for key in keys: + cluster_id = key["server_id"] - return least_loaded_server_id + if re.match(r"^cluster\d+$", cluster_id): + if cluster_id not in cluster_loads: + cluster_loads[cluster_id] = 0 + cluster_loads[cluster_id] += 1 + finally: + await conn.close() + + logger.info(f"Cluster loads: {cluster_loads}") + + if not cluster_loads: + available_clusters = [ + cluster_id + for cluster_id in CLUSTERS.keys() + if re.match(r"^cluster\d+$", cluster_id) + ] + + logger.info(f"Available clusters from config: {available_clusters}") + + if available_clusters: + logger.info( + f"Returning the first available cluster: {available_clusters[0]}" + ) + return available_clusters[0] + else: + logger.warning("No valid clusters found in config, returning 'cluster1'.") + return "cluster1" + + least_loaded_cluster = min(cluster_loads, key=cluster_loads.get) + + logger.info(f"Least loaded cluster selected: {least_loaded_cluster}") + + return least_loaded_cluster async def handle_error(tg_id, callback_query, message):