371 lines
18 KiB
Python
371 lines
18 KiB
Python
import asyncio
|
||
import datetime
|
||
import os
|
||
from typing import Annotated, Optional, List, Tuple
|
||
|
||
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
|
||
import boto3
|
||
from botocore.config import Config
|
||
|
||
from ..config import WorkerConfig
|
||
from ..middleware.authorization import verify_token
|
||
from ..models.media_model import (MediaSources,
|
||
CacheResult,
|
||
MediaSource,
|
||
MediaCacheStatus,
|
||
DownloadResult,
|
||
UploadResultResponse,
|
||
UploadBase64Request, UploadPresignRequest, UploadPresignResponse,
|
||
UploadMultipartPresignRequest, UploadMultipartPresignResponse, MediaProtocol
|
||
)
|
||
from ..models.web_model import SentryTransactionInfo, MonitorLiveRoomProductRequest, ModalTaskResponse, \
|
||
LiveRoomProductCachesResponse, CacheDeleteTaskResponse
|
||
from ..utils.KVCache import MediaSourceKVCache, LiveProductKVCache
|
||
from ..utils.SentryUtils import SentryUtils
|
||
|
||
config = WorkerConfig()
|
||
|
||
client = boto3.client("s3",
|
||
aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
|
||
aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
|
||
region_name=config.S3_region,
|
||
endpoint_url="https://s3-accelerate.amazonaws.com",
|
||
config=Config(
|
||
s3={'addressing_style': 'virtual'},
|
||
signature_version='s3v4', )
|
||
)
|
||
|
||
router = APIRouter(prefix="/cache", tags=['缓存'], )
|
||
|
||
modal_kv_cache = MediaSourceKVCache(kv_name=config.modal_kv_name,
|
||
environment=config.modal_environment)
|
||
|
||
modal_kv_product_cache = LiveProductKVCache(kv_name=config.modal_product_kv_name,
|
||
environment=config.modal_environment)
|
||
|
||
|
||
@router.post("/",
|
||
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("/",
|
||
summary="清除指定的所有缓存",
|
||
description="清除指定的所有缓存(包括KV记录和S3存储文件)",
|
||
dependencies=[Depends(verify_token)])
|
||
async def purge_media_kv_file(medias: MediaSources) -> CacheDeleteTaskResponse:
|
||
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) -> Tuple[Optional[str], int]:
|
||
try:
|
||
cache_media = modal_kv_cache.pop(media.urn)
|
||
if cache_media:
|
||
deleted_cache: MediaSource = await fn.remote.aio(cache_media)
|
||
return deleted_cache.urn, 1
|
||
except KeyError as e:
|
||
logger.exception(e)
|
||
if media.local_exists:
|
||
deleted_cache: MediaSource = await fn.remote.aio(media)
|
||
if deleted_cache.status == MediaCacheStatus.missing:
|
||
logger.warning(f"不存在s3挂载文件 {deleted_cache.urn}")
|
||
return None, 0
|
||
else:
|
||
logger.warning(f"s3挂载文件 {deleted_cache.urn} 已删除")
|
||
return deleted_cache.urn, 0
|
||
return media.urn, -1
|
||
|
||
async with asyncio.TaskGroup() as group:
|
||
tasks = [group.create_task(purge_handle(media)) for media in medias.inputs]
|
||
|
||
keys: List[str] = []
|
||
non_kv_keys: List[str] = []
|
||
error_keys: List[str] = []
|
||
|
||
for task in tasks:
|
||
urn, task_status = task.result()
|
||
if urn:
|
||
if task_status == 1: # 成功从kv和s3删除
|
||
keys.append(urn)
|
||
elif task_status == 2: # 只从s3删除
|
||
non_kv_keys.append(urn)
|
||
else:
|
||
error_keys.append(urn)
|
||
# keys = [task.result() for task in tasks]
|
||
modal_kv_cache.batch_remove_cloudflare_kv(keys)
|
||
return CacheDeleteTaskResponse(success=True, keys=keys, nonKVKeys=non_kv_keys, notFoundKeys=error_keys)
|
||
# return JSONResponse(content={"success": True, "keys": keys, "nonKVKeys": non_kv_keys, "errorKeys": error_keys})
|
||
|
||
|
||
@router.post("/download",
|
||
summary="批量获取下载地址, 返回的MediaSource类自带CDN访问URL, 不需要另外请求获取",
|
||
deprecated=True,
|
||
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",
|
||
summary="下载已缓存的视频",
|
||
deprecated=True,
|
||
description="通过CDN下载已缓存的视频文件, 不在提供通过此接口下载，请使用CDN URL直接下载文件")
|
||
@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",
|
||
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",
|
||
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("/upload-s3",
|
||
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',
|
||
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)
|
||
|
||
|
||
@router.post("/monitor_live_room_product_trigger",
|
||
summary="触发监控直播间商品信息并缓存",
|
||
description="触发监控直播间商品信息并缓存, 如果直播结束清除缓存, 触发间隔请控制在60s以上",
|
||
dependencies=[Depends(verify_token)], deprecated=True)
|
||
async def monitor_live_room_product(body: MonitorLiveRoomProductRequest) -> LiveRoomProductCachesResponse:
|
||
fn = modal.Function.from_name(config.modal_app_name, "monitor_live_room_product_trigger",
|
||
environment_name=config.modal_environment)
|
||
status = await fn.remote.aio(body.cookie, body.room_id, body.author_id)
|
||
if status == 0:
|
||
product_list = modal_kv_product_cache.get_cache(body.room_id)
|
||
return LiveRoomProductCachesResponse(status=status, cache_json=product_list.model_dump_json())
|
||
elif status == 1:
|
||
return LiveRoomProductCachesResponse(status=status, message="直播已结束")
|
||
elif status == 2:
|
||
return LiveRoomProductCachesResponse(status=status, message="分配到风控IP, 请稍后重试")
|
||
elif status == 3:
|
||
return LiveRoomProductCachesResponse(status=status, message="请求Tikhub API出现错误")
|
||
else:
|
||
return LiveRoomProductCachesResponse(status=4, message="内部错误")
|
||
|
||
|
||
@router.post('/upload-s3/simple/presign',
|
||
summary="S3简单上传预签名",
|
||
description="利用S3就近接入点上传",
|
||
dependencies=[Depends(verify_token)])
|
||
async def s3_presign_upload(body: UploadPresignRequest) -> UploadPresignResponse:
|
||
expires_in = 3600
|
||
expired_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in)
|
||
signed_url = client.generate_presigned_url("put_object",
|
||
Params={
|
||
'Bucket': config.S3_bucket_name,
|
||
'Key': f"upload/{body.key}",
|
||
"ContentType": body.content_type,
|
||
}, ExpiresIn=expires_in, )
|
||
return UploadPresignResponse(url=signed_url,
|
||
urn=f"s3://{config.S3_region}/{config.S3_bucket_name}/upload/{body.key}",
|
||
expired_at=expired_at)
|
||
|
||
|
||
@router.post("/upload-s3/multipart/presign",
|
||
summary="S3分片上传预签名",
|
||
description="""
|
||
1. 本地按文件总大小分Chunk大小按urls链接内的顺序通过HTTP PUT请求上传文件分片; 并将上传完成后获得的返回头ETag值记录与PartNumber对应, PartNumber对应使用url在urls内的顺位, 以1开始
|
||
\n2. 所有分片上传完成后使用XML格式拼装出用于确认上传的body; 并通过HTTP POST complete_url确认上传, ContentType 需要确保为application/xml\n\n
|
||
<CompleteMultipartUpload>
|
||
<Part>
|
||
<PartNumber>1</PartNumber>
|
||
<ETag>"60575364b098a1a48765a28c3a48e0ef"</ETag>
|
||
</Part>
|
||
<Part>
|
||
<PartNumber>2</PartNumber>
|
||
<ETag>"a38691c31fd242faee5533c65b4501d7"</ETag>
|
||
</Part>
|
||
<Part>
|
||
<PartNumber>3</PartNumber>
|
||
<ETag>"7a7510cc83f98feea28f319567e4cf66"</ETag>
|
||
</Part>
|
||
</CompleteMultipartUpload>
|
||
\n3. 如无法确认上传，可使用HTTP GET list_url debug当前分片上传状态，确认成功后list_url无法返回有效数据
|
||
""", dependencies=[Depends(verify_token)])
|
||
async def s3_presign_upload_multipart(body: UploadMultipartPresignRequest) -> UploadMultipartPresignResponse:
|
||
chunk_count = body.parts_count
|
||
multipart_upload_response = client.create_multipart_upload(Bucket=config.S3_bucket_name, Key=body.key,
|
||
ContentType=body.content_type, )
|
||
upload_id = multipart_upload_response.get("UploadId")
|
||
signed_urls = []
|
||
expires_in = 3600
|
||
expired_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in)
|
||
for i in range(chunk_count):
|
||
signed_url = client.generate_presigned_url("upload_part",
|
||
Params={
|
||
'Bucket': config.S3_bucket_name,
|
||
'Key': body.key,
|
||
'PartNumber': i + 1,
|
||
'UploadId': upload_id,
|
||
}, ExpiresIn=expires_in)
|
||
signed_urls.append(signed_url)
|
||
|
||
signed_completed_url = client.generate_presigned_url("complete_multipart_upload",
|
||
Params={
|
||
'Bucket': config.S3_bucket_name,
|
||
'Key': body.key,
|
||
'UploadId': upload_id,
|
||
}, ExpiresIn=expires_in)
|
||
signed_list_url = client.generate_presigned_url("list_parts",
|
||
Params={
|
||
'Bucket': config.S3_bucket_name,
|
||
'Key': body.key,
|
||
'UploadId': upload_id,
|
||
}, ExpiresIn=expires_in)
|
||
return UploadMultipartPresignResponse(urls=signed_urls,
|
||
urn=f"s3://{config.S3_region}/{config.S3_bucket_name}/upload/{body.key}",
|
||
expired_at=expired_at,
|
||
complete_url=signed_completed_url,
|
||
list_url=signed_list_url)
|