feat(api): add atomic temp file handling for media downloads

- Use unique temporary file paths with UUID to avoid race conditions during concurrent downloads.
- Perform atomic rename of the temp file to the final cache location and handle existing concurrent downloads.
- Clean up zero‑size and stale temporary files, including new “.tmp.<hex>” pattern.
- Extend cleanup logic to detect and remove race‑condition temp files.
- Validate and convert environment variables (TG_API_ID, TG_PROXY_PORT, API_PORT) with clear error messages and proper exit handling.
- Flush prints and replace `os._exit` with `sys.exit` for graceful termination.
- Fix URL regex in `PostParser` and guard against missing channel usernames when generating media URLs.
- Deep‑copy message lists in RSS generator to prevent mutation across calls.
- Propagate `FloodWait` exceptions to be handled by the API server instead of retrying internally.
- Enhance `TelegramClient` disconnect handling with a brief pause, shutdown check, and reconnection reset.
- Refine auth retry logic to check for actual `KeyError` instances.
This commit is contained in:
vvzvlad
2026-05-17 23:49:18 +03:00
parent dd5999b51f
commit ab8f15d49d
5 changed files with 177 additions and 134 deletions
+63 -32
View File
@@ -11,6 +11,8 @@
import logging
import os
import re
import uuid
import mimetypes
from typing import List, Union, Any
from urllib.parse import quote
@@ -438,29 +440,47 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu
logger.error(f"Failed to remove entry for {channel}/{post_id}/{file_unique_id} from SQLite: {str(e)}")
raise HTTPException(status_code=404, detail="File not found in message")
# Use a unique temp path to avoid race conditions when concurrent requests download the same file
temp_path = cache_path + f".tmp.{uuid.uuid4().hex}"
try:
file_path = await client.safe_download_media(file_id, cache_path)
file_path = await client.safe_download_media(file_id, temp_path)
except asyncio.TimeoutError:
logger.error(f"Timeout downloading media {file_unique_id}")
raise HTTPException(status_code=504, detail="Download timeout")
# Check if the downloaded file exists and has a size greater than zero
if not file_path or not os.path.exists(file_path) or os.path.getsize(file_path) == 0:
logger.error(f"download_failed_zero_size: Downloaded file {file_unique_id} for {channel}/{post_id} is zero size or missing.")
# Attempt to clean up the invalid file
if file_path and os.path.exists(file_path):
# Clean up the temp file if it was created
if os.path.exists(temp_path):
try:
os.remove(file_path)
logger.info(f"Removed zero-size file: {file_path}")
os.remove(temp_path)
except OSError:
pass
raise HTTPException(status_code=504, detail="Download timeout")
# Check if the downloaded temp file exists and has a size greater than zero
if not file_path or not os.path.exists(temp_path) or os.path.getsize(temp_path) == 0:
logger.error(f"download_failed_zero_size: Downloaded file {file_unique_id} for {channel}/{post_id} is zero size or missing.")
# Attempt to clean up the invalid temp file
if os.path.exists(temp_path):
try:
os.remove(temp_path)
logger.info(f"Removed zero-size temp file: {temp_path}")
except OSError as e:
logger.error(f"cleanup_error: Failed to remove zero-size file {file_path}: {e}")
# Raise an error to indicate download failure
# Raise specific error
logger.error(f"cleanup_error: Failed to remove zero-size temp file {temp_path}: {e}")
# Raise specific error to indicate download failure
raise ZeroSizeFileError(f"Downloaded file {file_unique_id} for {channel}/{post_id} is zero size or missing after download attempt.")
# Atomic rename: if another concurrent request already wrote the final file, discard our temp copy
if os.path.exists(cache_path) and os.path.getsize(cache_path) > 0:
logger.info(f"Concurrent download already completed for {file_unique_id}, discarding temp file")
try:
os.remove(temp_path)
except OSError as e:
logger.warning(f"cleanup_error: Failed to remove temp file {temp_path}: {e}")
return cache_path, False
# Atomically move temp file to final cache path (atomic on POSIX)
os.rename(temp_path, cache_path)
logger.info(f"Downloaded media file {file_unique_id} to {cache_path}")
return file_path, False
return cache_path, False
def remove_old_cached_files_sync(media_files: list, cache_dir: str) -> tuple[list, int]:
@@ -511,35 +531,46 @@ def remove_old_cached_files_sync(media_files: list, cache_dir: str) -> tuple[lis
logger.error(f"Error processing cache entry {file_data}: {str(e)}")
continue
# Additionally, remove temporary files with prefix "temp_" if they are older than a threshold (1 hour)
# Clean up temporary files: "temp_"-prefixed (large videos) and ".tmp."-suffixed (race-condition downloads)
temp_threshold = 3600 # 1 hour in seconds
for root, _, files in os.walk(cache_dir):
for file in files:
if file.startswith("temp_"):
file_path = os.path.join(root, file)
try:
file_mod_time = os.path.getmtime(file_path)
if time.time() - file_mod_time > temp_threshold:
# Get channel/post/file_id from path
is_large_video_temp = file.startswith("temp_")
# Match exactly the format generated in download_media_file: {name}.tmp.{32 hex chars}
is_race_temp = bool(re.match(r'^.+\.tmp\.[0-9a-f]{32}$', file))
if not (is_large_video_temp or is_race_temp):
continue
file_path = os.path.join(root, file)
try:
file_mod_time = os.path.getmtime(file_path)
if time.time() - file_mod_time > temp_threshold:
if is_large_video_temp and not is_race_temp:
# Extract channel/post/file_id from path and name
rel_path = os.path.relpath(os.path.dirname(file_path), cache_dir)
channel, post_id = rel_path.split(os.sep)
parts = rel_path.split(os.sep)
if len(parts) != 2:
logger.warning(f"Unexpected path depth for large-video temp file, skipping: {file_path}")
continue
channel, post_id = parts
file_unique_id = file[5:] # Remove 'temp_' prefix
# Remove file
os.remove(file_path)
files_removed += 1
logger.info(f"Removed temporary file: {file_path}")
logger.info(f"Removed temporary large video file: {file_path}")
# Also remove the temporary file entry from the in-memory list
updated_media_files = [
f for f in updated_media_files
if not (f.get('channel') == channel and
f.get('post_id') == int(post_id) and
f for f in updated_media_files
if not (f.get('channel') == channel and
f.get('post_id') == int(post_id) and
f.get('file_unique_id') == file_unique_id)
]
except Exception as e:
logger.error(f"Failed to remove temporary file {file_path}: {str(e)}")
elif is_race_temp:
# Race-condition temp file — just delete from disk, no SQLite entry to remove
os.remove(file_path)
files_removed += 1
logger.info(f"Removed stale race-condition temp file: {file_path}")
except Exception as e:
logger.error(f"Failed to remove temporary file {file_path}: {str(e)}")
return updated_media_files, files_removed
+27 -6
View File
@@ -9,6 +9,7 @@
# pylance: disable=reportMissingImports, reportMissingModuleSource
import os
import sys
import logging
from typing import Any
@@ -42,8 +43,15 @@ def get_settings() -> dict[str, Any]:
tg_api_id = os.getenv("TG_API_ID")
tg_api_hash = os.getenv("TG_API_HASH")
if not tg_api_id or not tg_api_hash:
print("TG_API_ID and TG_API_HASH must be set")
os._exit(1)
print("TG_API_ID and TG_API_HASH must be set", flush=True)
sys.exit(1)
# Validate and convert TG_API_ID to integer
try:
tg_api_id_int = int(tg_api_id)
except ValueError:
print(f"TG_API_ID must be a valid integer, got: {tg_api_id!r}", flush=True)
sys.exit(1)
log_level = os.getenv("LOG_LEVEL", "INFO")
@@ -53,23 +61,36 @@ def get_settings() -> dict[str, Any]:
proxy_host = os.getenv("TG_PROXY_HOST")
proxy: dict[str, Any] | None = None
if proxy_host:
# Validate and convert TG_PROXY_PORT to integer
try:
proxy_port = int(os.getenv("TG_PROXY_PORT") or 1080)
except ValueError:
print(f"TG_PROXY_PORT must be a valid integer, got: {os.getenv('TG_PROXY_PORT')!r}", flush=True)
sys.exit(1)
proxy = {
"scheme": "SOCKS5",
"hostname": proxy_host,
"port": int(os.getenv("TG_PROXY_PORT") or 1080),
"port": proxy_port,
"username": os.getenv("TG_PROXY_USERNAME") or None,
"password": os.getenv("TG_PROXY_PASSWORD") or None,
}
# Validate and convert API_PORT to integer
try:
api_port = int(os.getenv("API_PORT") or 8000)
except ValueError:
print(f"API_PORT must be a valid integer, got: {os.getenv('API_PORT')!r}", flush=True)
sys.exit(1)
return {
"tg_api_id": int(tg_api_id),
"tg_api_id": tg_api_id_int,
"tg_api_hash": tg_api_hash,
"session_path": os.getenv("SESSION_PATH", "data") or "data",
"api_host": os.getenv("API_HOST", "0.0.0.0"),
"api_port": int(os.getenv("API_PORT") or 8000),
"api_port": api_port,
"pyrogram_bridge_url": os.getenv("PYROGRAM_BRIDGE_URL", ""),
"log_level": log_level,
"debug": os.getenv("DEBUG", "False") == "True",
"debug": os.getenv("DEBUG", "False").strip() in ["True", "true"],
"token": os.getenv("TOKEN", ""),
"trusted_proxies": [ip.strip() for ip in os.getenv("TRUSTED_PROXIES", "").split(",") if ip.strip()],
"time_based_merge": os.getenv("TIME_BASED_MERGE", "False").strip() in ["True", "true"],
+59 -58
View File
@@ -225,7 +225,7 @@ class PostParser:
if text_stripped:
text_was_processed = True
text_has_urls = bool(re.search(r'https?://[^\\s<>\"\\\']+', text_stripped))
text_has_urls = bool(re.search(r'https?://[^\s<>"\']+', text_stripped))
clean_text = html.unescape(text) # Remove HTML entities
clean_text = re.sub(r'<[^>]+?>', '', clean_text) # Remove HTML tags
@@ -675,61 +675,60 @@ class PostParser:
logger.debug(f"File unique id not found for message {message.id}")
elif file_unique_id:
channel_username = self.get_channel_username(message)
file = f"{channel_username}/{message.id}/{file_unique_id}"
digest = generate_media_digest(file)
url = f"{base_url}/media/{file}/{digest}"
# Guard: channel_username may be None for private chats without a username
if not channel_username:
logger.warning(f"Could not generate media URL for message {message.id}: channel username is missing.")
content_media.append('</div>')
else:
file = f"{channel_username}/{message.id}/{file_unique_id}"
digest = generate_media_digest(file)
url = f"{base_url}/media/{file}/{digest}"
logger.debug(f"Collected media file: {channel_username}/{message.id}/{file_unique_id}")
# Check if document is a PDF file
if (message.media == MessageMediaType.DOCUMENT and
message.document is not None and hasattr(message.document, 'mime_type') and
message.document.mime_type == 'application/pdf'):
# Only attempt to create link if channel_username is available
if channel_username:
if channel_username.startswith('-100'):
logger.debug(f"Collected media file: {channel_username}/{message.id}/{file_unique_id}")
# Check if document is a PDF file
if (message.media == MessageMediaType.DOCUMENT and
message.document is not None and hasattr(message.document, 'mime_type') and
message.document.mime_type == 'application/pdf'):
# Only attempt to create link if channel_username is available
if channel_username.startswith('-100'):
tg_link = f"https://t.me/c/{channel_username[4:]}/{message.id}"
else:
else:
tg_link = f"https://t.me/{channel_username}/{message.id}"
content_media.append(f'<div class="document-pdf" style="padding: 10px;">')
content_media.append(f'<a href="{tg_link}" target="_blank">[PDF-файл]</a></div>')
else:
# Handle case where channel_username is None (e.g., log or add placeholder)
logger.warning(f"Could not generate PDF link for {message.id}: ch username is missing.")
# Add placeholder without link
content_media.append(f'<div class="document-pdf" style="padding: 10px;">[PDF-файл]</div>')
elif message.media in [MessageMediaType.PHOTO, MessageMediaType.DOCUMENT]:
content_media.append(f'<img src="{url}" style="max-width:100%; width:auto; height:auto;'
f'max-height:400px; object-fit:contain;">')
elif message.media == MessageMediaType.VIDEO:
content_media.append(f'<video controls src="{url}" style="max-width:100%; width:auto;'
f'height:auto; max-height:400px;"></video>')
elif message.media == MessageMediaType.ANIMATION:
content_media.append(f'<video controls src="{url}" style="max-width:100%; width:auto;'
f'height:auto; max-height:400px;"></video>')
elif message.media == MessageMediaType.VIDEO_NOTE:
content_media.append(f'<video controls src="{url}" style="max-width:100%; width:auto;'
f'height:auto; max-height:400px;"></video>')
elif message.media == MessageMediaType.AUDIO:
mime_type = getattr(message.audio, 'mime_type', 'audio/mpeg')
content_media.append(f'<audio controls style="width:100%; max-width:400px;">'
f'<source src="{url}" type="{mime_type}"></audio>')
content_media.append('<br>')
elif message.media == MessageMediaType.VOICE:
mime_type = getattr(message.voice, 'mime_type', 'audio/ogg')
content_media.append(f'<audio controls style="width:100%; max-width:400px;">'
f'<source src="{url}" type="{mime_type}"></audio>')
content_media.append('<br>')
elif message.media == MessageMediaType.STICKER:
emoji = getattr(message.sticker, 'emoji', '')
if getattr(message.sticker, 'is_video', False):
content_media.append(f'<video controls autoplay loop muted src="{url}"'
f'style="max-width:100%; width:auto; height:auto; max-height:200px;'
f'object-fit:contain;"></video>')
else:
content_media.append(f'<img src="{url}" alt="Sticker {emoji}" style="max-width:100%;'
f'width:auto; height:auto; max-height:200px; object-fit:contain;">')
content_media.append('</div>')
elif message.media in [MessageMediaType.PHOTO, MessageMediaType.DOCUMENT]:
content_media.append(f'<img src="{url}" style="max-width:100%; width:auto; height:auto;'
f'max-height:400px; object-fit:contain;">')
elif message.media == MessageMediaType.VIDEO:
content_media.append(f'<video controls src="{url}" style="max-width:100%; width:auto;'
f'height:auto; max-height:400px;"></video>')
elif message.media == MessageMediaType.ANIMATION:
content_media.append(f'<video controls src="{url}" style="max-width:100%; width:auto;'
f'height:auto; max-height:400px;"></video>')
elif message.media == MessageMediaType.VIDEO_NOTE:
content_media.append(f'<video controls src="{url}" style="max-width:100%; width:auto;'
f'height:auto; max-height:400px;"></video>')
elif message.media == MessageMediaType.AUDIO:
mime_type = getattr(message.audio, 'mime_type', 'audio/mpeg')
content_media.append(f'<audio controls style="width:100%; max-width:400px;">'
f'<source src="{url}" type="{mime_type}"></audio>')
content_media.append('<br>')
elif message.media == MessageMediaType.VOICE:
mime_type = getattr(message.voice, 'mime_type', 'audio/ogg')
content_media.append(f'<audio controls style="width:100%; max-width:400px;">'
f'<source src="{url}" type="{mime_type}"></audio>')
content_media.append('<br>')
elif message.media == MessageMediaType.STICKER:
emoji = getattr(message.sticker, 'emoji', '')
if getattr(message.sticker, 'is_video', False):
content_media.append(f'<video controls autoplay loop muted src="{url}"'
f'style="max-width:100%; width:auto; height:auto; max-height:200px;'
f'object-fit:contain;"></video>')
else:
content_media.append(f'<img src="{url}" alt="Sticker {emoji}" style="max-width:100%;'
f'width:auto; height:auto; max-height:200px; object-fit:contain;">')
content_media.append('</div>')
if webpage := getattr(message, "web_page", None): # Web page preview
if webpage_html := self._format_webpage(webpage, message):
@@ -789,13 +788,15 @@ class PostParser:
if photo := getattr(webpage, "photo", None):
if file_unique_id := getattr(photo, "file_unique_id", None):
channel_username = self.get_channel_username(message)
file = f"{channel_username}/{message.id}/{file_unique_id}"
digest = generate_media_digest(file)
url = f"{base_url}/media/{file}/{digest}"
html_parts.append(f'<div class="webpage-photo" style="margin-top:10px;">')
html_parts.append(f'<a href="{webpage.url}" target="_blank">')
html_parts.append(f'<img src="{url}" style="max-width:100%; width:auto;'
f'height:auto; max-height:200px; object-fit:contain;"></a></div>')
# Guard: skip photo if channel_username is unavailable to avoid broken URLs
if channel_username:
file = f"{channel_username}/{message.id}/{file_unique_id}"
digest = generate_media_digest(file)
url = f"{base_url}/media/{file}/{digest}"
html_parts.append(f'<div class="webpage-photo" style="margin-top:10px;">')
html_parts.append(f'<a href="{webpage.url}" target="_blank">')
html_parts.append(f'<img src="{url}" style="max-width:100%; width:auto;'
f'height:auto; max-height:200px; object-fit:contain;"></a></div>')
html_parts.append('</div>')
return '\n'.join(html_parts)
+10 -35
View File
@@ -9,6 +9,7 @@
# pylance: disable=reportMissingImports, reportMissingModuleSource
# mypy: disable-error-code="import-untyped"
import copy
import logging
import asyncio
import re
@@ -32,6 +33,9 @@ async def _create_time_based_media_groups(messages: list[Message], merge_seconds
"""
Create media groups based on time difference between messages
"""
# Deep-copy the input list to avoid mutating cached Message objects (the cache may
# reuse the same objects across calls with different merge_seconds values)
messages = copy.deepcopy(messages)
# Compute fallback once so all None-date messages get the same sort key (deterministic order)
_sort_fallback = datetime.now(timezone.utc)
messages_sorted = sorted(messages, key=lambda msg: msg.date or _sort_fallback) # type: ignore
@@ -350,24 +354,9 @@ async def generate_channel_rss(channel: str | int,
except (errors.UsernameInvalid, errors.UsernameNotOccupied) as e:
logger.warning(f"Channel not found error for {channel}: {str(e)}")
return create_error_feed(channel, base_url)
except errors.FloodWait as e:
# Handle Telegram flood wait - wait for the specified duration and retry
wait_time = e.value
logger.warning(f"Telegram flood wait detected for channel '{channel}'. Waiting {wait_time} seconds before retry...")
await asyncio.sleep(wait_time)
logger.info(f"Flood wait completed for channel '{channel}', retrying get_chat...")
# Retry the get_chat operation after waiting
try:
channel_info = await post_parser.client.get_chat(channel)
channel_title = channel_info.title or f"Telegram: {channel}"
channel_username = channel_info.username or (str(channel_info.id) if channel_info.id and str(channel_info.id).startswith('-100') else None)
if not channel_username:
logger.warning(f"Could not get username for channel {channel} after flood wait, using identifier for error feed.")
return create_error_feed(str(channel), base_url)
except Exception as retry_e:
logger.error(f"Retry after flood wait failed for channel '{channel}': {str(retry_e)}", exc_info=True)
raise ValueError(f"Failed to get chat info for {channel} after flood wait: {str(retry_e)}") from retry_e
except errors.FloodWait:
# Let FloodWait bubble up to api_server.py, which returns HTTP 429 with Retry-After header
raise
except Exception as e:
logger.error(f"Error during get_chat for channel '{channel}' (type: {type(channel)}): {str(e)}", exc_info=True) # Log error specifically for get_chat
# Re-raise the original exception to be caught by the outer handler if needed,
@@ -547,23 +536,9 @@ async def generate_channel_html(channel: str | int,
logger.warning(f"Channel not found error for {channel} in HTML generation: {str(e)}")
# Consider returning a dedicated HTML error page.
return create_error_feed(channel, base_url)
except errors.FloodWait as e:
# Handle Telegram flood wait - wait for the specified duration and retry
wait_time = e.value
logger.warning(f"Telegram flood wait detected for channel '{channel}' in HTML generation. Waiting {wait_time} seconds before retry...")
await asyncio.sleep(wait_time)
logger.info(f"Flood wait completed for channel '{channel}' in HTML generation, retrying get_chat...")
# Retry the get_chat operation after waiting
try:
channel_info = await post_parser.client.get_chat(channel)
channel_username = channel_info.username or (str(channel_info.id) if channel_info.id and str(channel_info.id).startswith('-100') else None)
if not channel_username:
logger.warning(f"Could not get username for channel {channel} after flood wait in HTML generation, returning error feed.")
return create_error_feed(str(channel), base_url)
except Exception as retry_e:
logger.error(f"Retry after flood wait failed for channel '{channel}' in HTML generation: {str(retry_e)}", exc_info=True)
raise ValueError(f"Failed to get chat info for {channel} after flood wait in HTML generation: {str(retry_e)}") from retry_e
except errors.FloodWait:
# Let FloodWait bubble up to api_server.py, which returns HTTP 429 with Retry-After header
raise
except Exception as e:
logger.error(f"Error during get_chat for channel '{channel}' (type: {type(channel)}) in HTML generation: {str(e)}", exc_info=True)
raise ValueError(f"Failed to get chat info for {channel} in HTML generation: {str(e)}") from e
+18 -3
View File
@@ -60,7 +60,22 @@ class TelegramClient:
return
self.disconnect_count += 1
logger.warning(f"connection_handler: connection lost (#{self.disconnect_count})")
# Wait briefly to allow Pyrogram's internal reconnect to complete
await asyncio.sleep(10)
# Re-check shutdown flag — it may have been set while we were sleeping
# (e.g. another concurrent _on_disconnect already triggered _restart_app)
if self._shutting_down:
logger.debug("connection_handler: shutdown started during sleep, suppressing further action")
return
# If the client reconnected successfully, reset the counter and return
if self.client.is_connected:
logger.info(f"connection_handler: reconnected successfully, resetting disconnect counter")
self.disconnect_count = 0
return
if self.disconnect_count >= self.max_disconnects:
logger.critical(f"connection_handler: reached disconnect limit ({self.max_disconnects})")
self._restart_app()
@@ -99,7 +114,7 @@ class TelegramClient:
timeout=30.0
)
except Exception as e:
if "KeyError" in str(e) and attempt < max_retries - 1:
if isinstance(e, KeyError) and attempt < max_retries - 1:
logger.warning(f"Auth error on attempt {attempt + 1}, retrying in 5s...")
await asyncio.sleep(5)
continue
@@ -114,7 +129,7 @@ class TelegramClient:
timeout=120.0
)
except Exception as e:
if "KeyError" in str(e) and attempt < max_retries - 1:
if isinstance(e, KeyError) and attempt < max_retries - 1:
logger.warning(f"Download auth error on attempt {attempt + 1}, retrying...")
await asyncio.sleep(5)
continue