handler.py

6.9 kB · python · 179 lines

1import boto32import json3import os4import re5import time6from datetime import datetime, timedelta, timezone78BUCKET = os.environ["STATS_BUCKET"]9DISTRIBUTION = os.environ.get("STATS_DISTRIBUTION", "")10FUNCTIONS = [one.strip() for one in os.environ["STATS_FUNCTIONS"].split(",") if one.strip()]11STATS_KEY = os.environ["STATS_KEY"]12HOURS = 2413ERROR_LINES = 1014ERROR_PATTERN = '?ERROR ?"Task timed out" ?Traceback ?"Runtime exited"'15CLIP = 16016URL = re.compile(r"https?://\S+")17QUERY = re.compile(r"\?[^\s]*=[^\s]*")18SECRET = re.compile(r"(?i)(bearer\s+\S+|[\w.-]*(?:key|token|secret|password|auth|signature|credential)[\w.-]*\s*[=:]\s*(?:bearer\s+)?\S+)")1920cloudwatch = boto3.client("cloudwatch")21cloudwatch_global = boto3.client("cloudwatch", region_name="us-east-1")22logs = boto3.client("logs")23s3 = boto3.client("s3")2425# LISTS2627def rows(name: str) -> list[tuple[str, str, str]]:28    out = []29    for one in os.environ.get(name, "").split(","):30        label, _, rest = one.strip().partition("=")31        prefix, _, pattern = rest.partition("=")32        if label.strip(): out.append((label.strip(), prefix.strip(), pattern.strip()))33    return out3435SIZES = rows("STATS_SIZES")36COUNTS = rows("STATS_COUNTS")37FOLDERS = rows("STATS_FOLDERS")3839# METRICS4041def window() -> tuple[datetime, datetime]:42    end = datetime.now(timezone.utc).replace(second=0, microsecond=0)43    return end - timedelta(hours=HOURS), end4445def series(client, namespace: str, name: str, dimensions: list[dict], stat: str, period: int) -> list[tuple[int, float]]:46    start, end = window()47    query = {48        "Id": "q",49        "MetricStat": {"Metric": {"Namespace": namespace, "MetricName": name, "Dimensions": dimensions}, "Period": period, "Stat": stat},50        "ReturnData": True,51    }52    result = client.get_metric_data(MetricDataQueries=[query], StartTime=start, EndTime=end, ScanBy="TimestampAscending")["MetricDataResults"][0]53    return [(int(stamp.timestamp()), round(value, 3)) for stamp, value in zip(result["Timestamps"], result["Values"])]5455def total(points: list[tuple[int, float]]) -> float:56    return round(sum(value for _, value in points), 3)5758def mean(points: list[tuple[int, float]]) -> float:59    return round(sum(value for _, value in points) / len(points), 3) if points else 06061def cloudfront() -> dict:62    if not DISTRIBUTION:63        return {}64    dimensions = [{"Name": "DistributionId", "Value": DISTRIBUTION}, {"Name": "Region", "Value": "Global"}]65    requests = series(cloudwatch_global, "AWS/CloudFront", "Requests", dimensions, "Sum", 3600)66    bytes_out = series(cloudwatch_global, "AWS/CloudFront", "BytesDownloaded", dimensions, "Sum", 3600)67    errors_4xx = series(cloudwatch_global, "AWS/CloudFront", "4xxErrorRate", dimensions, "Average", 3600)68    errors_5xx = series(cloudwatch_global, "AWS/CloudFront", "5xxErrorRate", dimensions, "Average", 3600)69    return {70        "requests": total(requests),71        "bytes": total(bytes_out),72        "error_4xx": mean(errors_4xx),73        "error_5xx": mean(errors_5xx),74        "hourly": [{"at": stamp, "requests": value} for stamp, value in requests],75    }7677def lambdas() -> dict:78    out = {}79    for name in FUNCTIONS:80        dimensions = [{"Name": "FunctionName", "Value": name}]81        invocations = series(cloudwatch, "AWS/Lambda", "Invocations", dimensions, "Sum", 3600)82        errors = series(cloudwatch, "AWS/Lambda", "Errors", dimensions, "Sum", 3600)83        throttles = series(cloudwatch, "AWS/Lambda", "Throttles", dimensions, "Sum", 3600)84        duration = series(cloudwatch, "AWS/Lambda", "Duration", dimensions, "Average", 3600)85        out[name] = {86            "invocations": total(invocations),87            "errors": total(errors),88            "throttles": total(throttles),89            "duration_ms": mean(duration),90            "hourly": [{"at": stamp, "invocations": value} for stamp, value in invocations],91        }92    return out9394# LOGS9596def clean(message: str) -> str:97    text = URL.sub("<url>", str(message))98    text = SECRET.sub("<hidden>", QUERY.sub("", text))99    return " ".join(text.split())[:CLIP]100101def parse(message: str) -> dict:102    try:103        record = json.loads(message)104    except ValueError:105        return {"at": None, "level": "", "message": clean(message)}106    if "message" in record:107        return {"at": record.get("timestamp"), "level": record.get("level", ""), "message": clean(record["message"])}108    return {"at": record.get("time"), "level": record.get("type", ""), "message": clean(json.dumps(record.get("record", record)))}109110def recent_errors() -> dict:111    start, end = window()112    out = {}113    for name in FUNCTIONS:114        group = f"/aws/lambda/{name}"115        try:116            events = logs.filter_log_events(logGroupName=group, startTime=int(start.timestamp() * 1000), endTime=int(end.timestamp() * 1000), filterPattern=ERROR_PATTERN, limit=100)["events"]117        except logs.exceptions.ResourceNotFoundException:118            events = []119        out[name] = [parse(event["message"]) for event in events[-ERROR_LINES:]]120    return out121122# BUCKET123124def count_prefixes(prefix: str, pattern: re.Pattern = None) -> int:125    count = 0126    paginator = s3.get_paginator("list_objects_v2")127    for page in paginator.paginate(Bucket=BUCKET, Prefix=prefix, Delimiter="/"):128        for one in page.get("CommonPrefixes", []):129            tail = one["Prefix"][len(prefix):]130            if pattern is None or pattern.match(tail):131                count += 1132    return count133134def count_objects(prefix: str) -> tuple[int, int]:135    objects = 0136    size = 0137    paginator = s3.get_paginator("list_objects_v2")138    for page in paginator.paginate(Bucket=BUCKET, Prefix=prefix):139        for obj in page.get("Contents", []):140            objects += 1141            size += obj["Size"]142    return objects, size143144def bucket() -> dict:145    out = {}146    for label, prefix, _ in SIZES:147        objects, size = count_objects(prefix)148        out[f"{label}_objects"] = objects149        out[f"{label}_bytes"] = size150    for label, prefix, _ in COUNTS:151        out[label] = count_objects(prefix)[0]152    for label, prefix, pattern in FOLDERS:153        out[label] = count_prefixes(prefix, re.compile(pattern) if pattern else None)154    return out155156# RUN157158def build() -> dict:159    started = time.time()160    data = {161        "at": int(started),162        "hours": HOURS,163        "cdn": cloudfront(),164        "lambdas": lambdas(),165        "errors": recent_errors(),166        "bucket": bucket(),167    }168    data["seconds"] = round(time.time() - started, 1)169    return data170171def handler(event, context):172    data = build()173    s3.put_object(Bucket=BUCKET, Key=STATS_KEY, Body=json.dumps(data), ContentType="application/json", CacheControl="no-cache")174    summary = {name: (one["invocations"], one["errors"]) for name, one in data["lambdas"].items()}175    print(f"stats written in {data['seconds']} s: cdn {data['cdn'].get('requests')} requests, lambdas {summary}, bucket {data['bucket']}")176    return {"seconds": data["seconds"], "keys": len(data["bucket"])}177178if __name__ == "__main__":179    print(json.dumps(build(), indent=2))