266 lines
12 KiB
Python
266 lines
12 KiB
Python
import asyncio
|
||
import os
|
||
from typing import Annotated, Optional
|
||
|
||
import modal
|
||
from loguru import logger
|
||
from modal import current_function_call_id
|
||
import sentry_sdk
|
||
from fastapi import APIRouter, Depends, UploadFile, HTTPException, File, Form
|
||
from fastapi.responses import JSONResponse, RedirectResponse
|
||
from starlette import status
|
||
from ..config import WorkerConfig
|
||
from ..middleware.authorization import verify_token
|
||
from ..models.media_model import (MediaSources,
|
||
CacheResult,
|
||
MediaSource,
|
||
MediaCacheStatus,
|
||
DownloadResult,
|
||
UploadResultResponse,
|
||
UploadBase64Request
|
||
)
|
||
from ..models.web_model import SentryTransactionInfo
|
||
from ..utils.KVCache import MediaSourceKVCache
|
||
from ..utils.SentryUtils import SentryUtils
|
||
|
||
config = WorkerConfig()
|
||
|
||
router = APIRouter(prefix="/cache")
|
||
if not config.CF_KV_namespace_id:
|
||
raise ValueError("未配置Cloudflare KV namespace ID")
|
||
|
||
modal_kv_cache = MediaSourceKVCache(kv_name=config.modal_kv_name,
|
||
cf_kv_id=config.CF_KV_namespace_id,
|
||
environment=config.modal_environment)
|
||
|
||
|
||
@router.post("/",
|
||
tags=["缓存"],
|
||
summary="缓存视频文件",
|
||
description="异步缓存视频文件到S3存储桶和Modal Dict(KV)",
|
||
dependencies=[Depends(verify_token)])
|
||
async def cache(medias: MediaSources) -> CacheResult:
|
||
fn_id = current_function_call_id()
|
||
caches: MediaSources
|
||
sentry_trace = SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(),
|
||
x_baggage=sentry_sdk.get_baggage())
|
||
|
||
@SentryUtils.sentry_tracker(name="同步视频缓存", op="cache.get", fn_id=fn_id,
|
||
sentry_trace_id=None, sentry_baggage=None)
|
||
async def cache_handler(media: MediaSource):
|
||
cache_span = sentry_sdk.get_current_span()
|
||
cache_span.set_data("runner_id", fn_id)
|
||
cache_span.set_data("cache.key", [media.urn])
|
||
cached_media = modal_kv_cache.get_cache(media.urn)
|
||
cache_hit: bool = False
|
||
if not cached_media:
|
||
# start new download task
|
||
with cache_span.start_child(name="视频缓存任务入队",
|
||
op="queue.publish") as queue_publish_span:
|
||
fn = modal.Function.from_name(config.modal_app_name, 'cache_submit')
|
||
fn_task = fn.spawn(media, sentry_trace)
|
||
queue_publish_span.set_data("cache.key", media.urn)
|
||
queue_publish_span.set_data("messaging.message.id", fn_task.object_id)
|
||
queue_publish_span.set_data("messaging.destination.name",
|
||
"video-downloader.cache_submit")
|
||
queue_publish_span.set_data("messaging.message.body.size", 0)
|
||
media.status = MediaCacheStatus.downloading
|
||
media.downloader_id = fn_task.object_id
|
||
modal_kv_cache.set_cache(media)
|
||
else:
|
||
media = cached_media
|
||
match media.status:
|
||
case MediaCacheStatus.ready:
|
||
cache_hit = True
|
||
case MediaCacheStatus.downloading: # 下载任务已经在进行
|
||
cache_hit = True
|
||
case _:
|
||
# start new download task
|
||
with cache_span.start_child(name="视频缓存任务入队",
|
||
op="queue.publish") as queue_publish_span:
|
||
fn = modal.Function.from_name(config.modal_app_name, 'cache_submit')
|
||
fn_task = fn.spawn(media, sentry_trace)
|
||
queue_publish_span.set_data("cache.key", media.urn)
|
||
queue_publish_span.set_data("messaging.message.id", fn_task.object_id)
|
||
queue_publish_span.set_data("messaging.destination.name",
|
||
"video-downloader.cache_submit")
|
||
queue_publish_span.set_data("messaging.message.body.size", 0)
|
||
media.status = MediaCacheStatus.downloading
|
||
media.downloader_id = fn_task.object_id
|
||
modal_kv_cache.set_cache(media)
|
||
cache_hit = False
|
||
|
||
cache_span.set_data("cache.hit", cache_hit)
|
||
return media
|
||
|
||
async with asyncio.TaskGroup() as group:
|
||
tasks = [group.create_task(cache_handler(media)) for media in medias.inputs]
|
||
|
||
cache_task_result_dict = {}
|
||
cache_task_result_list = []
|
||
|
||
for task in tasks:
|
||
result = task.result()
|
||
cache_task_result_dict[result.urn] = result.model_dump_json()
|
||
cache_task_result_list.append(result)
|
||
|
||
modal_kv_cache.batch_update_cloudflare_kv(cache_task_result_dict)
|
||
return CacheResult(caches={media.urn: media for media in cache_task_result_list})
|
||
|
||
|
||
@router.delete("/",
|
||
tags=["缓存"],
|
||
summary="清除指定的所有缓存",
|
||
description="清除指定的所有缓存(包括KV记录和S3存储文件)",
|
||
dependencies=[Depends(verify_token)])
|
||
async def purge_media_kv_file(medias: MediaSources):
|
||
fn_id = current_function_call_id()
|
||
fn = modal.Function.from_name(config.modal_app_name, "cache_delete", environment_name=config.modal_environment)
|
||
|
||
@SentryUtils.sentry_tracker(name="清除媒体源缓存", op="cache.purge", fn_id=fn_id,
|
||
sentry_trace_id=None, sentry_baggage=None)
|
||
async def purge_handle(media: MediaSource):
|
||
cache_media = modal_kv_cache.pop(media.urn)
|
||
if cache_media:
|
||
deleted_cache: MediaSource = await fn.remote.aio(cache_media)
|
||
return deleted_cache.urn
|
||
return None
|
||
|
||
async with asyncio.TaskGroup() as group:
|
||
tasks = [group.create_task(purge_handle(media)) for media in medias.inputs]
|
||
|
||
keys = [task.result() for task in tasks]
|
||
modal_kv_cache.batch_remove_cloudflare_kv(keys)
|
||
return JSONResponse(content={"success": True, "keys": keys})
|
||
|
||
|
||
@router.post("/download",
|
||
tags=["缓存"],
|
||
summary="批量获取下载地址",
|
||
description="获取已缓存的视频下载地址",
|
||
dependencies=[Depends(verify_token)])
|
||
@sentry_sdk.trace
|
||
async def download_caches(medias: MediaSources) -> DownloadResult:
|
||
cdn_endpoint = config.S3_cdn_endpoint
|
||
urls = []
|
||
for media in medias.inputs:
|
||
urls.append(f"{cdn_endpoint}/{media.get_cdn_url()}")
|
||
return DownloadResult(urls=urls)
|
||
|
||
|
||
@router.get("/download",
|
||
tags=["缓存"],
|
||
summary="下载已缓存的视频",
|
||
description="通过CDN下载已缓存的视频文件")
|
||
@sentry_sdk.trace
|
||
async def download_cache(media: str) -> RedirectResponse:
|
||
cdn_endpoint = config.S3_cdn_endpoint
|
||
media = MediaSource.from_str(media)
|
||
return RedirectResponse(url=f"{cdn_endpoint}/{media.get_cdn_url()}", status_code=status.HTTP_302_FOUND)
|
||
|
||
|
||
@router.delete("/kv",
|
||
tags=["缓存"],
|
||
summary="清除KV记录",
|
||
description="清除当前环境下KV缓存过的所有数据(S3存储桶内的文件会保留)",
|
||
dependencies=[Depends(verify_token)])
|
||
async def purge_kv_all():
|
||
parent = sentry_sdk.get_current_span()
|
||
span = parent.start_child(name="清除缓存KV", op="purge.flush")
|
||
modal_kv_cache.clear()
|
||
span.set_data("cache.success", True)
|
||
span.finish()
|
||
return JSONResponse(content={"success": True})
|
||
|
||
|
||
@router.post("/kv",
|
||
tags=["缓存"],
|
||
summary="删除对应的KV记录",
|
||
description="删除请求中对应的视频缓存记录",
|
||
dependencies=[Depends(verify_token)])
|
||
async def purge_kv(medias: MediaSources):
|
||
try:
|
||
for media in medias.inputs:
|
||
modal_kv_cache.pop(media.urn)
|
||
keys = [media.urn for media in medias.inputs]
|
||
modal_kv_cache.batch_remove_cloudflare_kv(keys)
|
||
return JSONResponse(content={"success": True, "keys": keys})
|
||
except Exception as e:
|
||
return JSONResponse(content={"success": False, "error": str(e)})
|
||
|
||
|
||
@router.post("/media",
|
||
tags=["缓存"],
|
||
summary="清除指定的所有缓存",
|
||
description="清除指定的所有缓存(包括KV记录和S3存储文件), 将要被淘汰，使用DELETE /cache/替代",
|
||
deprecated=True,
|
||
dependencies=[Depends(verify_token)])
|
||
async def purge_media(medias: MediaSources):
|
||
fn_id = current_function_call_id()
|
||
fn = modal.Function.from_name(config.modal_app_name, "cache_delete", environment_name=config.modal_environment)
|
||
|
||
@SentryUtils.sentry_tracker(name="清除媒体源缓存", op="cache.purge", fn_id=fn_id,
|
||
sentry_trace_id=None, sentry_baggage=None)
|
||
async def purge_handle(media: MediaSource):
|
||
cache_media = modal_kv_cache.pop(media.urn)
|
||
if cache_media:
|
||
deleted_cache: MediaSource = await fn.remote.aio(cache_media)
|
||
return deleted_cache.urn
|
||
return None
|
||
|
||
async with asyncio.TaskGroup() as group:
|
||
tasks = [group.create_task(purge_handle(media)) for media in medias.inputs]
|
||
|
||
keys = [task.result() for task in tasks]
|
||
modal_kv_cache.batch_remove_cloudflare_kv(keys)
|
||
return JSONResponse(content={"success": True, "keys": keys})
|
||
|
||
|
||
@router.post("/upload-s3",
|
||
tags=['缓存'],
|
||
summary="上传文件到S3",
|
||
description="上传文件到S3的文件必须小于200M",
|
||
dependencies=[Depends(verify_token)])
|
||
async def s3_upload(file: Annotated[UploadFile, File(description="上传的文件")],
|
||
prefix: Annotated[Optional[str], Form()] = None) -> UploadResultResponse:
|
||
fn_id = current_function_call_id()
|
||
|
||
if file.size > 200 * 1024 * 1024: # 上传文件不大于200M
|
||
raise HTTPException(status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail="上传文件不可超过200MB")
|
||
|
||
key = f"upload/{prefix}/{file.filename}" if prefix else f"upload/{file.filename}"
|
||
local_path = f"{config.S3_mount_dir}/{key}"
|
||
logger.info(f"s3上传到{key}, size={file.size}")
|
||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||
with open(local_path, 'wb') as f:
|
||
f.write(file.file.read())
|
||
logger.info(f"{local_path} 保存成功")
|
||
media_source = MediaSource.from_str(f"s3://{config.S3_region}/{config.S3_bucket_name}/{key}")
|
||
media_source.status = MediaCacheStatus.ready
|
||
media_source.downloader_id = fn_id
|
||
return UploadResultResponse(media=media_source)
|
||
|
||
|
||
@router.post('/upload-s3-b64',
|
||
tags=['缓存'],
|
||
summary="基于Base64格式上传文件到S3",
|
||
description="上传文件到S3当文件必须小于200M",
|
||
dependencies=[Depends(verify_token)])
|
||
async def s3_upload_base64(body: UploadBase64Request) -> UploadResultResponse:
|
||
fn_id = current_function_call_id()
|
||
prefix = body.prefix
|
||
file = body.file
|
||
if file.size > 200 * 1024 * 1024: # 上传文件不大于200M
|
||
raise HTTPException(status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail="上传文件不可超过200MB")
|
||
key = f"upload/{prefix}/{file.filename}" if prefix else f"upload/{file.filename}"
|
||
local_path = f"{config.S3_mount_dir}/{key}"
|
||
logger.info(f"s3上传到{key}, size={file.size}")
|
||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||
with open(local_path, 'wb') as f:
|
||
f.write(file.raw_content)
|
||
logger.info(f"{local_path} 保存成功")
|
||
media_source = MediaSource.from_str(f"s3://{config.S3_region}/{config.S3_bucket_name}/{key}")
|
||
media_source.status = MediaCacheStatus.ready
|
||
media_source.downloader_id = fn_id
|
||
return UploadResultResponse(media=media_source)
|