转码成本降 71%，处理时长缩短一半

"""分片转码任务的中断处理。

关键设计：
    1. 分片是幂等的，重跑不会产生副作用
    2. Spot 中断通知触发主动上报，而不是等超时
    3. 重试计数在 DynamoDB 里，避免无限重试
    4. 连续失败的分片降级到按需实例，保证最终完成
"""

from __future__ import annotations

import json
import os
from datetime import datetime, timezone

import boto3

ddb = boto3.resource("dynamodb")
sqs = boto3.client("sqs")

TASK_TABLE = ddb.Table(os.environ["TASK_TABLE"])
SPOT_QUEUE = os.environ["SPOT_QUEUE_URL"]
ONDEMAND_QUEUE = os.environ["ONDEMAND_QUEUE_URL"]
MAX_SPOT_RETRY = 3


def requeue_shard(job_id: str, shard_id: str, reason: str) -> dict:
    """重新投递分片。超过 Spot 重试上限则转到按需队列。"""
    resp = TASK_TABLE.update_item(
        Key={"job_id": job_id, "shard_id": shard_id},
        UpdateExpression=(
            "SET #st = :pending, last_failure = :reason, updated_at = :now "
            "ADD retry_count :one"
        ),
        ExpressionAttributeNames={"#st": "status"},
        ExpressionAttributeValues={
            ":pending": "pending",
            ":reason": reason,
            ":now": datetime.now(timezone.utc).isoformat(),
            ":one": 1,
        },
        ReturnValues="ALL_NEW",
    )
    item = resp["Attributes"]
    retries = int(item.get("retry_count", 0))

    # 重试次数用完就上按需，宁可贵一点也要把任务做完
    queue = ONDEMAND_QUEUE if retries > MAX_SPOT_RETRY else SPOT_QUEUE

    sqs.send_message(
        QueueUrl=queue,
        MessageBody=json.dumps({
            "job_id": job_id,
            "shard_id": shard_id,
            "input_key": item["input_key"],
            "output_prefix": item["output_prefix"],
            "profile": item["profile"],
            "attempt": retries + 1,
        }),
        MessageGroupId=job_id,
    )

    return {
        "job_id": job_id,
        "shard_id": shard_id,
        "retry_count": retries,
        "queue": "on-demand" if retries > MAX_SPOT_RETRY else "spot",
    }


def lambda_handler(event, context):  # noqa: ARG001
    """处理两类事件：Spot 中断警告，以及 Batch 任务失败。"""
    source = event.get("source",")
    detail = event.get("detail", {})
    results = []

    if source == "aws.ec2" and event.get("detail-type") == "EC2 Spot Instance Interruption Warning":
        instance_id = detail["instance-id"]
        # 找出该实例上正在跑的所有分片，立即重新投递
        running = TASK_TABLE.query(
            IndexName="instance-index",
            KeyConditionExpression=boto3.dynamodb.conditions.Key("instance_id").eq(instance_id),
            FilterExpression=boto3.dynamodb.conditions.Attr("status").eq("running"),
        )["Items"]
        for shard in running:
            results.append(requeue_shard(
                shard["job_id"], shard["shard_id"], "spot-interruption"
            ))

    elif source == "aws.batch" and detail.get("status") == "FAILED":
        results.append(requeue_shard(
            detail["parameters"]["job_id"],
            detail["parameters"]["shard_id"],
            detail.get("statusReason", "batch-failed"),
        ))

    return {"requeued": len(results), "shards": results}