handler.py

import base64
import gzip
import json
import os

import boto3

ecs = boto3.client("ecs")
iot_data = boto3.client("iot-data", endpoint_url=f"https://{os.environ['IOT_ENDPOINT']}")

ECS_CLUSTER = os.environ["ECS_CLUSTER"]

# Per-warm-instance cache: ecsTaskId -> (testId, taskIndex). Avoids one
# DescribeTasks call per log line for tasks already seen by this instance.
_task_cache = {}


def _resolve_task(ecs_task_id):
    if ecs_task_id in _task_cache:
        return _task_cache[ecs_task_id]

    resp = ecs.describe_tasks(cluster=ECS_CLUSTER, tasks=[ecs_task_id], include=["TAGS"])
    tasks = resp.get("tasks", [])
    if not tasks:
        return None

    tags = {t["key"]: t["value"] for t in tasks[0].get("tags", [])}
    test_id = tags.get("TestId")
    task_index = tags.get("TaskIndex")
    if test_id is None or task_index is None:
        return None

    result = (test_id, task_index)
    _task_cache[ecs_task_id] = result
    return result


def lambda_handler(event, context):
    payload = json.loads(gzip.decompress(base64.b64decode(event["awslogs"]["data"])))

    for log_event in payload.get("logEvents", []):
        # logStream format: runner/<containerName>/<ecsTaskId>
        parts = payload.get("logStream", "").split("/")
        if len(parts) != 3:
            continue
        _, container, ecs_task_id = parts

        resolved = _resolve_task(ecs_task_id)
        if resolved is None:
            continue
        test_id, task_index = resolved

        iot_data.publish(
            topic=f"loadtest/{test_id}/{task_index}/logs",
            qos=0,
            payload=json.dumps(
                {
                    "testId": test_id,
                    "taskIndex": task_index,
                    "container": container,
                    "timestamp": log_event["timestamp"],
                    "message": log_event["message"],
                }
            ),
        )