mirror of
https://github.com/Viren070/AIOStremio.git
synced 2025-12-01 23:24:19 +01:00
Viren070
118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
import asyncio
|
|
import io
|
|
from collections import deque
|
|
from typing import AsyncGenerator, Dict
|
|
|
|
import aiohttp
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
from utils.config import config
|
|
from utils.logger import logger
|
|
|
|
|
|
class StreamManager:
|
|
def __init__(
|
|
self,
|
|
chunk_size: int = config.chunk_size_mb * 1024 * 1024,
|
|
buffer_size: int = config.buffer_size_mb * 1024 * 1024,
|
|
):
|
|
self.chunk_size = chunk_size
|
|
self.buffer_size = buffer_size
|
|
self.default_headers = {
|
|
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
|
|
"Accept": "*/*",
|
|
"Connection": "keep-alive",
|
|
}
|
|
|
|
async def create_streaming_response(
|
|
self, url: str, request_headers: Dict
|
|
) -> StreamingResponse:
|
|
headers = self.default_headers.copy()
|
|
if "range" in request_headers:
|
|
headers["range"] = request_headers["range"]
|
|
|
|
timeout = aiohttp.ClientTimeout(total=180)
|
|
async with aiohttp.ClientSession(timeout=timeout) as session:
|
|
async with session.get(url, headers=headers) as response:
|
|
content_range = response.headers.get("Content-Range", "")
|
|
content_length = response.headers.get("Content-Length", "")
|
|
|
|
response_headers = {
|
|
"Accept-Ranges": "bytes",
|
|
"Content-Range": content_range if content_range else None,
|
|
"Content-Length": content_length if content_length else None,
|
|
"Connection": "keep-alive",
|
|
"Cache-Control": "no-cache",
|
|
}
|
|
response_headers = {k: v for k, v in response_headers.items() if v is not None}
|
|
media_type = response.headers.get("Content-Type", "video/mp4")
|
|
|
|
return StreamingResponse(
|
|
self._stream_content(url, headers),
|
|
media_type=media_type,
|
|
status_code=206 if "range" in request_headers else 200,
|
|
headers=response_headers,
|
|
)
|
|
|
|
async def _stream_content(
|
|
self, url: str, headers: Dict
|
|
) -> AsyncGenerator[bytes, None]:
|
|
timeout = aiohttp.ClientTimeout(total=None, connect=120, sock_read=120)
|
|
buffer = deque()
|
|
current_buffer_size = 0
|
|
buffer_low_threshold = self.buffer_size * 0.2 # 20% of buffer size
|
|
|
|
connector = aiohttp.TCPConnector(
|
|
limit=0,
|
|
ttl_dns_cache=300,
|
|
force_close=False,
|
|
enable_cleanup_closed=True,
|
|
)
|
|
|
|
async with aiohttp.ClientSession(
|
|
timeout=timeout, connector=connector
|
|
) as session:
|
|
async with session.get(url, headers=headers) as response:
|
|
|
|
async def fill_buffer():
|
|
nonlocal current_buffer_size
|
|
while True:
|
|
try:
|
|
if current_buffer_size < self.buffer_size:
|
|
chunk = await response.content.read(self.chunk_size)
|
|
if not chunk:
|
|
break
|
|
buffer.append(chunk)
|
|
current_buffer_size += len(chunk)
|
|
else:
|
|
await asyncio.sleep(0.1)
|
|
except (asyncio.TimeoutError, aiohttp.ClientError) as e:
|
|
logger.error(f"Buffer filling error: {str(e)}")
|
|
await asyncio.sleep(1)
|
|
continue
|
|
|
|
buffer_task = asyncio.create_task(fill_buffer())
|
|
|
|
try:
|
|
while not buffer:
|
|
await asyncio.sleep(0.1)
|
|
|
|
while True:
|
|
if not buffer and current_buffer_size == 0:
|
|
break
|
|
|
|
if buffer:
|
|
chunk = buffer.popleft()
|
|
current_buffer_size -= len(chunk)
|
|
yield chunk
|
|
|
|
if current_buffer_size < buffer_low_threshold:
|
|
await asyncio.sleep(0.1)
|
|
|
|
finally:
|
|
buffer_task.cancel()
|
|
try:
|
|
await buffer_task
|
|
except asyncio.CancelledError:
|
|
pass
|