"""媒体缓存文件的有界读取、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})