ComfyUI-Ranking-API / media_file_response.py
ZHIWEI666's picture
Upload 11 files
e12a676 verified
Raw History Blame Contribute Delete
5.32 kB
"""媒体缓存文件的有界读取、Range交付与真实线程关闭。"""
import asyncio
import os
from secrets import token_hex
import stat
from fastapi.responses import FileResponse
from starlette.datastructures import MutableHeaders
from owned_thread import run_owned_thread
from workflow_stream_response import OwnedDelivery
class OwnedMediaFileResponse(OwnedDelivery, FileResponse):
total_timeout = 300
operation_timeout = 15
eof_timeout = 5
def __init__(self, path, **kwargs):
snapshot = os.stat(path)
if not stat.S_ISREG(snapshot.st_mode):
raise ValueError('invalid media cache file')
super().__init__(path, stat_result=snapshot, **kwargs)
self._file = None
self._init_delivery(self._close_file)
async def send_owned(self, scope, receive, send):
async def bounded_send(message):
final = message['type'] == 'http.response.body' and not message.get('more_body', False)
async with asyncio.timeout(self.eof_timeout if final else self.operation_timeout):
await send(message)
async with asyncio.timeout(self.total_timeout):
return await super().send_owned(scope, receive, bounded_send)
async def _file_work(self, function, *args):
async with asyncio.timeout(self.operation_timeout):
return await run_owned_thread(function, *args)
async def _open_file(self):
def open_file():
self._file = open(self.path, 'rb')
self._check_snapshot(os.fstat(self._file.fileno()))
await self._file_work(open_file)
def _check_snapshot(self, current):
fields = ('st_dev', 'st_ino', 'st_size', 'st_mtime_ns')
if os.name != 'nt':
fields += ('st_ctime_ns',)
if (not stat.S_ISREG(current.st_mode) or not current.st_ino or not self.stat_result.st_ino
or any(getattr(current, key) != getattr(self.stat_result, key) for key in fields)):
raise RuntimeError('media cache file changed')
async def _prepare_file(self, header_only):
if header_only:
current = await self._file_work(os.stat, self.path)
self._check_snapshot(current)
else:
await self._open_file()
async def _close_file(self):
if self._file is None:
return
def close_file():
try:
self._file.close()
finally:
self._file = None
await run_owned_thread(close_file)
async def _send_part(self, send, start, end, *, final):
await self._file_work(self._file.seek, start)
remaining = end - start
if remaining == 0:
await send({'type': 'http.response.body', 'body': b'', 'more_body': not final})
while remaining:
chunk = await self._file_work(self._file.read, min(self.chunk_size, remaining))
if not chunk:
raise RuntimeError('media cache file truncated')
remaining -= len(chunk)
await send({'type': 'http.response.body', 'body': chunk, 'more_body': remaining > 0 or not final})
async def _handle_simple(self, send, send_header_only, send_pathsend):
await self._prepare_file(send_header_only)
await send({'type': 'http.response.start', 'status': self.status_code, 'headers': self.raw_headers})
if send_header_only:
await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
return
await self._send_part(send, 0, self.stat_result.st_size, final=True)
async def _handle_single_range(self, send, start, end, file_size, send_header_only):
await self._prepare_file(send_header_only)
headers = MutableHeaders(raw=list(self.raw_headers))
headers['content-range'] = f'bytes {start}-{end - 1}/{file_size}'
headers['content-length'] = str(end - start)
await send({'type': 'http.response.start', 'status': 206, 'headers': headers.raw})
if send_header_only:
await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
return
await self._send_part(send, start, end, final=True)
async def _handle_multiple_ranges(self, send, ranges, file_size, send_header_only):
await self._prepare_file(send_header_only)
boundary = token_hex(13)
length, header = self.generate_multipart(ranges, boundary, file_size, self.headers['content-type'])
headers = MutableHeaders(raw=list(self.raw_headers))
headers['content-type'] = f'multipart/byteranges; boundary={boundary}'
headers['content-length'] = str(length)
await send({'type': 'http.response.start', 'status': 206, 'headers': headers.raw})
if send_header_only:
await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
return
for start, end in ranges:
await send({'type': 'http.response.body', 'body': header(start, end), 'more_body': True})
await self._send_part(send, start, end, final=False)
await send({'type': 'http.response.body', 'body': b'\r\n', 'more_body': True})
await send({'type': 'http.response.body', 'body': f'--{boundary}--'.encode('latin-1'), 'more_body': False})