refactor api_server and post_parser: implement asynchronous JSON file operations with locking mechanism to ensure thread safety, and streamline media file ID persistence logic

2025-09-13 20:47:16 +03:00
parent 1464ea326e
commit 68a2dbfa04
5 changed files with 225 additions and 37 deletions
@@ -37,6 +37,7 @@ from config import get_settings, setup_logging
 from rss_generator import generate_channel_rss, generate_channel_html
 from post_parser import PostParser
 from url_signer import verify_media_digest, generate_media_digest
+from file_io import read_json_file_sync, write_json_file_sync, get_media_file_ids_lock

 # Global python-magic instance for MIME type detection
 magic_mime = magic.Magic(mime=True)
@@ -283,7 +284,8 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu
                            file_data.get('file_unique_id') == file_unique_id):
                            file_data['added'] = datetime.now().timestamp()
                            break
-                    await asyncio.to_thread(write_json_file_sync, file_ids_path, media_files)
+                    async with get_media_file_ids_lock():
+                        await asyncio.to_thread(write_json_file_sync, file_ids_path, media_files)
            except Exception as e:
                logger.error(f"Failed to update timestamp for {channel}/{post_id}/{file_unique_id}: {str(e)}")
            return cache_path, False
@@ -308,7 +310,8 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu
                    )
                ]
                if len(media_files_updated) < initial_count:
-                    await asyncio.to_thread(write_json_file_sync, file_ids_path, media_files_updated)
+                    async with get_media_file_ids_lock():
+                        await asyncio.to_thread(write_json_file_sync, file_ids_path, media_files_updated)
                    logger.info(f"Removed invalid entry for {channel}/{post_id}/{file_unique_id} from media_file_ids.json")
                else:
                    logger.warning(f"Entry for {channel}/{post_id}/{file_unique_id} not found in media_file_ids.json for removal")
@@ -515,7 +518,8 @@ async def cache_media_files() -> None:

            if files_removed > 0:
                try:
-                    await asyncio.to_thread(write_json_file_sync, file_path, updated_media_files)
+                    async with get_media_file_ids_lock():
+                        await asyncio.to_thread(write_json_file_sync, file_path, updated_media_files)
                    logger.info(f"Removed {files_removed} old files from cache")
                except Exception as e:
                    logger.error(f"Failed to update media_file_ids.json: {str(e)}")
@@ -591,16 +595,6 @@ def calculate_cache_stats() -> dict[str, Any]:
        "channels": channels_stats
    }

-# Synchronous JSON helpers to be used via asyncio.to_thread
-def read_json_file_sync(file_path: str) -> list:
-    with open(file_path, 'r', encoding='utf-8') as f:
-        return json.load(f)
-
-def write_json_file_sync(file_path: str, data: list) -> None:
-    temp_path = f"{file_path}.tmp"
-    with open(temp_path, 'w', encoding='utf-8') as f:
-        json.dump(data, f, ensure_ascii=False, indent=2)
-    os.replace(temp_path, file_path)

 def is_local_request(request: Request) -> bool:
    local_hosts = ["127.0.0.1", "localhost"]
@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# flake8: noqa
+# pylint: disable=missing-function-docstring
+
+import os
+import json
+import asyncio
+from typing import Any, List
+
+
+# One shared lock for operations that write to media_file_ids.json
+_media_file_ids_lock: asyncio.Lock | None = None
+
+
+def get_media_file_ids_lock() -> asyncio.Lock:
+    global _media_file_ids_lock
+    if _media_file_ids_lock is None:
+        _media_file_ids_lock = asyncio.Lock()
+    return _media_file_ids_lock
+
+
+def read_json_file_sync(file_path: str) -> List[dict[str, Any]]:
+    with open(file_path, 'r', encoding='utf-8') as f:
+        return json.load(f)
+
+
+def write_json_file_sync(file_path: str, data: List[dict[str, Any]]) -> None:
+    temp_path = f"{file_path}.tmp"
+    with open(temp_path, 'w', encoding='utf-8') as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+    os.replace(temp_path, file_path)
+
+
@@ -0,0 +1,136 @@
+
+### Где и почему может виснуть
+
+- Синхронный диск/JSON в обработчиках запросов
+  - `_save_media_file_ids` делает `open/json.load/json.dump` прямо в пути запроса, без `to_thread`. При большом `data/media_file_ids.json` это блокирует event loop.
+```950:979:post_parser.py
+                        if os.path.exists(file_path):
+                            with open(file_path, 'r', encoding='utf-8') as f:
+                                existing_data = json.load(f)
+...
+                        with open(file_path, 'w', encoding='utf-8') as f:
+                            json.dump(existing_data, f, ensure_ascii=False, indent=2)
+```
+  - Кеш истории также читает/пишет файлы синхронно:
+```71:99:tg_cache.py
+        with open(cache_file, 'rb') as f:
+            cache_data = pickle.load(f)
+...
+        with open(cache_file, 'wb') as f:
+            pickle.dump(cache_data, f)
+```
+  - Итог: под нагрузкой любой долгий sync I/O останавливает обработку всех запросов на время операции.
+
+- CPU‑тяжёлое формирование RSS/HTML в event loop
+  - Парсинг/санитизация/генерация RSS идут синхронно в корутинах (bleach, feedgen, склейка больших строк). На больших лимитах это “съедает” цикл.
+```185:299:rss_generator.py
+final_posts = await _render_messages_groups(...)
+...
+rss_feed = fg.rss_str(pretty=True)
+```
+```511:533:rss_generator.py
+final_posts = await _render_messages_groups(...)
+...
+html = '\n<hr class="post-divider">\n'.join(html_posts)
+```
+
+- Глобальный `python-magic` в потоках
+  - Один общий `Magic()` используется из разных потоков — это не потокобезопасно, возможны зависания внутри libmagic.
+```198:206:api_server.py
+media_type = await asyncio.to_thread(magic_mime.from_file, file_path)
+```
+
+- BaseHTTPMiddleware над стримингом файлов
+  - `BaseHTTPMiddleware` оборачивает `call_next`. Для `FileResponse` это может ломать/буферизовать стриминг и усиливать задержки.
+```48:62:api_server.py
+class RequestLoggingMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        ...
+        response = await call_next(request)
+```
+
+- Отсутствие таймаутов на вызовах Telegram API
+  - Если `get_messages/download_media` “подвисли”, запрос к `/media` висит долго.
+```748:758:api_server.py
+file_path, delete_after = await download_media_file(...)
+```
+```232:258:api_server.py
+message = await client.client.get_messages(...)
+file_path = await client.client.download_media(...)
+```
+
+- Шумное логирование в middleware
+  - Лог всех заголовков и запросов при каждом хите может создавать I/O‑бутылочное горлышко под нагрузкой.
+```49:59:api_server.py
+logger.info(f"Request: {request.method} {request.url}")
+logger.info(f"Headers: {dict(request.headers)}")
+```
+
+- Конкурентная запись и разрастание `media_file_ids.json`
+  - Записи без блокировок → гонки и порча JSON. В фоне потом ремонт json-repair (тяжёлый CPU).
+```503:511:api_server.py
+except json.JSONDecodeError:
+    media_files = await asyncio.to_thread(fix_corrupted_json, file_path)
+```
+
+- Один воркер uvicorn
+  - Любая тяжёлая задача “монополизирует” процесс — остальное, включая отдачу файлов, ждёт.
+```114:121:api_server.py
+uvicorn.run(..., reload=True, loop="uvloop")
+```
+
+### Что сделать (короткий чеклист)
+
+- [x] Убрать sync I/O из пути запроса
+  - [x] Обернуть чтение/запись JSON и pickle в `asyncio.to_thread`.
+  - [x] Для `media_file_ids.json` везде использовать единые sync‑хелперы с временным файлом и заменить прямые `open/json.dump`:
+    - [x] Вынести в отдельный модуль функции уже имеющихся `read_json_file_sync`/`write_json_file_sync` и звать их через `to_thread`.
+  - [x] Добавить один `asyncio.Lock` на операции записи `media_file_ids.json` чтобы не было гонок.
+
+- Разгрузить CPU из event loop
+  - В RSS/HTML:
+    - После того как данные получены из Telegram, вынести “склейку HTML”, bleach‑санитизацию и `fg.rss_str` в `asyncio.to_thread`.
+    - Ограничить `limit` по умолчанию до 50–80. Для больших значений — отдавать 429 или 400 с рекомендацией.
+    - Опционально — кэшировать итоговые RSS/HTML на пару минут.
+
+- Исправить `python-magic`
+  - Создавать `Magic()` per‑call внутри `to_thread` ИЛИ защищать глобальный объект `threading.Lock`.
+  - Если не критично — падать на `mimetypes` и не вызывать `magic` для известных расширений.
+
+- Middleware для логов
+  - В проде отключить текущий `RequestLoggingMiddleware` или перевести на ASGI‑middleware без вмешательства в стриминг.
+  - Снизить уровень и объём логов (заголовки — только на DEBUG, семплирование).
+
+- Таймауты телеграма
+  - Оборачивать `get_messages`/`download_media` в `asyncio.wait_for(..., timeout=30–60)` с понятной ошибкой/503.
+  - На FloodWait в `/media` — отдавать 429 с `Retry-After`.
+
+- Запись `media_file_ids.json`
+  - Причесать размер: периодически чистить/ротация по размеру/возрасту.
+  - В фоне перепись файла уже есть — не трогать event loop; убедиться, что на записи тоже используется temp‑файл + `os.replace` (как в `write_json_file_sync`).
+
+- Процессный параллелизм
+  - Запускать с несколькими воркерами, без `reload`:
+    - пример: `uvicorn api_server:app --host 0.0.0.0 --port 80 --workers 2 --loop uvloop --http httptools`
+  - Или за Gunicorn: `--workers 2 --worker-class uvicorn.workers.UvicornWorker`.
+
+- Пул потоков
+  - Явно ограничить/увеличить пул для `to_thread` если есть много фоновых задач, чтобы не выжирался из‑за долгих операций.
+
+### Минимальные целевые правки (по месту)
+
+- `_save_media_file_ids` — переписать на `await asyncio.to_thread(read_json_file_sync/ write_json_file_sync)` + общий `asyncio.Lock`. Убрать прямые `open/json.dump`.
+- `tg_cache` — обернуть `pickle.load/dump` в `to_thread`.
+- `prepare_file_response` — перестать шарить `magic_mime`; создавать локальный `magic.Magic(mime=True)` внутри `to_thread` или защитить глобальный объект `Lock`.
+- `RequestLoggingMiddleware` — удалить/заменить на ASGI‑middleware и логировать только метод+путь, без заголовков по умолчанию.
+- `get_messages/download_media` — таймауты через `asyncio.wait_for`.
+- `uvicorn.run` — убрать `reload=True`, добавить `workers=2+`.
+
+### Как быстро проверить
+- Нагрузить `/rss` c большим `limit` параллельно с `/media` и измерить задержку отдачи файлов до/после:
+  - до фиксов статик “замирает” во время генерации;
+  - после выноса CPU/I/O в `to_thread` и увеличения воркеров — файлы продолжают отдаваться.
+
+Коротко: основные причины — синхронный диск и CPU в корутинах, общий `python-magic`, и единичный воркер. Перенос I/O/CPU в `to_thread`, фиксы гонок JSON, таймауты на Telegram и 2–4 воркера решат подвисания.
+
+- Нашёл проблемные места: sync I/O и CPU в обработчиках, `python-magic`, middleware, таймауты и один воркер. Дал точечные действия для устранения подвисаний.
@@ -9,9 +9,9 @@
 # pylance: disable=reportMissingImports, reportMissingModuleSource

 import logging
+import asyncio
 import re
 import os
-import json
 import html
 import inspect
 from datetime import datetime
@@ -21,6 +21,7 @@ from pyrogram.enums import MessageMediaType
 from bleach.css_sanitizer import CSSSanitizer   
 from bleach import clean as HTMLSanitizer
 from config import get_settings
+from file_io import read_json_file_sync, write_json_file_sync, get_media_file_ids_lock
 from url_signer import generate_media_digest

 Config = get_settings()
@@ -913,6 +914,29 @@ class PostParser:
            logger.error(f"file_id_extraction_error: media_type {message.media}, error {str(e)}")
            return None

+    async def _persist_media_file_id_async(self, file_path: str, file_data: dict) -> None:
+        try:
+            existing_data = []
+            if os.path.exists(file_path):
+                existing_data = await asyncio.to_thread(read_json_file_sync, file_path)
+
+            found = False
+            for item in existing_data:
+                if (item.get('channel') == file_data['channel'] and 
+                        item.get('post_id') == file_data['post_id'] and 
+                        item.get('file_unique_id') == file_data['file_unique_id']):
+                    item['added'] = datetime.now().timestamp()
+                    found = True
+                    break
+
+            if not found:
+                existing_data.append(file_data)
+
+            async with get_media_file_ids_lock():
+                await asyncio.to_thread(write_json_file_sync, file_path, existing_data)
+        except Exception as e:
+            logger.error(f"file_id_save_error: error writing to {file_path}, error {str(e)}")
+
    def _save_media_file_ids(self, message: Message) -> None:
        try:
            file_data = {
@@ -949,27 +973,25 @@ class PostParser:

                    file_path = os.path.join(os.path.abspath("./data"), 'media_file_ids.json')
                    try:
-                        existing_data = []
-                        if os.path.exists(file_path):
-                            with open(file_path, 'r', encoding='utf-8') as f:
-                                existing_data = json.load(f)
-                        
-                        # Check if file already exists by all three fields
-                        found = False
-                        for item in existing_data:
-                            if (item.get('channel') == file_data['channel'] and 
-                                    item.get('post_id') == file_data['post_id'] and 
-                                    item.get('file_unique_id') == file_data['file_unique_id']):
-                                item['added'] = datetime.now().timestamp()
-                                found = True
-                                break
-                        
-                        # Add new entry if not found
-                        if not found:
-                            existing_data.append(file_data)
-                        
-                        with open(file_path, 'w', encoding='utf-8') as f:
-                            json.dump(existing_data, f, ensure_ascii=False, indent=2)
+                        loop = asyncio.get_running_loop()
+                        if loop.is_running():
+                            loop.create_task(self._persist_media_file_id_async(file_path, file_data))
+                        else:
+                            # Fallback sync write (should not occur during normal FastAPI runtime)
+                            existing_data = []
+                            if os.path.exists(file_path):
+                                existing_data = read_json_file_sync(file_path)  # type: ignore
+                            found = False
+                            for item in existing_data:
+                                if (item.get('channel') == file_data['channel'] and 
+                                        item.get('post_id') == file_data['post_id'] and 
+                                        item.get('file_unique_id') == file_data['file_unique_id']):
+                                    item['added'] = datetime.now().timestamp()
+                                    found = True
+                                    break
+                            if not found:
+                                existing_data.append(file_data)
+                            write_json_file_sync(file_path, existing_data)  # type: ignore

                    except Exception as e:
                        logger.error(f"file_id_save_error: error writing to {file_path}, error {str(e)}")
@@ -14,6 +14,7 @@ import json
 import pickle
 import logging
 import random
+import asyncio
 import time
 from datetime import datetime, timedelta
 from typing import Any, Optional, Union, List
@@ -115,7 +116,7 @@ async def cached_get_chat_history(client: Client, channel_id: Union[str, int], l
    Returns:
        List of messages, same as original client.get_chat_history()
    """
-    cached_messages = _get_history_from_cache(channel_id, limit)
+    cached_messages = await asyncio.to_thread(_get_history_from_cache, channel_id, limit)
    
    if cached_messages is not None:
        return cached_messages
@@ -125,7 +126,7 @@ async def cached_get_chat_history(client: Client, channel_id: Union[str, int], l
        messages = []
        async for message in client.get_chat_history(channel_id, limit=limit):
            messages.append(message)
-        _save_history_to_cache(channel_id, messages, limit)
+        await asyncio.to_thread(_save_history_to_cache, channel_id, messages, limit)
        
        return messages
    except Exception as e: