Aurora PostgreSQL の CPU 高負荷時に自動で「インスタンス再起動」する仕組みを作った話

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "RdsReboot",
      "Effect": "Allow",
      "Action": [
        "rds:DescribeDBInstances",
        "rds:RebootDBInstance"
      ],
      "Resource": [
        "arn:aws:rds:ap-northeast-1:123456789012:db:myapp-prod-rdsinst",
        "arn:aws:rds:ap-northeast-1:123456789012:db:myapp-prod-rdsinst2"
      ]
    },
    {
      "Sid": "DynamoRebootHistory",
      "Effect": "Allow",
      "Action": [
        "dynamodb:GetItem",
        "dynamodb:PutItem"
      ],
      "Resource": "arn:aws:dynamodb:ap-northeast-1:123456789012:table/rds-auto-reboot-history"
    }
  ]
}

import os
import json
from datetime import datetime, timezone, timedelta

import boto3

rds = boto3.client("rds")
dynamo = boto3.client("dynamodb")

DDB_TABLE_NAME = os.environ["DDB_TABLE_NAME"]
MIN_REBOOT_INTERVAL_MINUTES = int(os.environ["MIN_REBOOT_INTERVAL_MINUTES"])


def lambda_handler(event, context):
    """
    CloudWatch アラームから呼び出されて RDS インスタンスを再起動する Lambda。

    対応イベント形式：
      1) 新しい CloudWatch アラーム → Lambda 形式
         {
           "source": "aws.cloudwatch",
           "alarmArn": "...",
           "alarmData": { "state": {...}, "configuration": {...} }
         }

      2) EventBridge (CloudWatch Alarm State Change) 経由
         { "detail-type": "...", "detail": { "state": {...}, "configuration": {...} } }

      3) 古い CloudWatch → Lambda 形式
         { "AlarmName": "...", "NewStateValue": "ALARM", "Trigger": { "Dimensions": [...] } }
    """
    print("Event:", json.dumps(event))

    db_instance_id = None

    # --- パターン1: 新しい CloudWatch アラーム → Lambda 形式 ---
    if event.get("source") == "aws.cloudwatch" and "alarmData" in event:
        alarm_data = event["alarmData"]
        state_value = alarm_data.get("state", {}).get("value")
        if state_value != "ALARM":
            print(f"[CloudWatch new] State is {state_value}. Do nothing.")
            return

        db_instance_id = extract_db_instance_id_from_metric_container(alarm_data)

    # --- パターン2: EventBridge (CloudWatch Alarm State Change) 経由 ---
    elif "detail-type" in event and "detail" in event:
        detail = event["detail"]
        state_value = detail.get("state", {}).get("value")
        if state_value != "ALARM":
            print(f"[EventBridge] State is {state_value}. Do nothing.")
            return

        db_instance_id = extract_db_instance_id_from_metric_container(detail)

    # --- パターン3: 古い CloudWatch → Lambda 形式 ---
    elif "AlarmName" in event:
        new_state = event.get("NewStateValue")
        if new_state != "ALARM":
            print(f"[LegacyAlarm] NewStateValue is {new_state}. Do nothing.")
            return

        db_instance_id = extract_db_instance_id_from_alarm_trigger(event)

    else:
        print(f"Unknown event format. keys={list(event.keys())}")
        return

    if not db_instance_id:
        print("DBInstanceIdentifier not found. Abort.")
        return

    print(f"Target DB instance: {db_instance_id}")

    # --- ループ防止: 直近で再起動していないか確認 ---
    if is_recently_rebooted(db_instance_id):
        print("Recently rebooted. Skip reboot.")
        return

    # --- RDS インスタンス状態確認 ---
    inst = rds.describe_db_instances(DBInstanceIdentifier=db_instance_id)["DBInstances"][0]
    status = inst["DBInstanceStatus"]
    print(f"Current status: {status}")
    if status != "available":
        print("Status is not 'available'. Skip reboot.")
        return

    # --- 再起動実行 (ForceFailover=False でフェイルオーバーはしない) ---
    resp = rds.reboot_db_instance(
        DBInstanceIdentifier=db_instance_id,
        ForceFailover=False
    )
    print("Reboot started:", resp["DBInstance"]["DBInstanceIdentifier"])

    # --- DynamoDB に再起動時刻を記録 ---
    now = datetime.now(timezone.utc).isoformat()
    dynamo.put_item(
        TableName=DDB_TABLE_NAME,
        Item={
            "DbInstanceIdentifier": {"S": db_instance_id},
            "LastRebootTime": {"S": now}
        }
    )

    return {"status": "reboot_started", "instance": db_instance_id}


def extract_db_instance_id_from_metric_container(container: dict) -> str | None:
    """
    alarmData / detail など、CloudWatch アラームの configuration.metrics[]
    を含むコンテナから DBInstanceIdentifier を取り出す。

    dimensions の形が
      - {"DBInstanceIdentifier": "..."} という dict
      - [{"name": "...", "value": "..."}] という list
    の両方をサポートする。
    """
    cfg = container.get("configuration", container)
    metrics = cfg.get("metrics", [])

    for m in metrics:
        metric_stat = m.get("metricStat", {})
        metric = metric_stat.get("metric", {})
        dims = metric.get("dimensions")

        if not dims:
            continue

        # パターンA: dict 形式 {"DBInstanceIdentifier": "..."}
        if isinstance(dims, dict):
            for k, v in dims.items():
                if k == "DBInstanceIdentifier":
                    return v

        # パターンB: list 形式 [{"name": "...", "value": "..."}]
        elif isinstance(dims, list):
            for d in dims:
                if not isinstance(d, dict):
                    continue
                name = d.get("name") or d.get("Name")
                value = d.get("value") or d.get("Value")
                if name == "DBInstanceIdentifier":
                    return value

    return None


def extract_db_instance_id_from_alarm_trigger(event: dict) -> str | None:
    """
    古い CloudWatch メトリクスアラームの Trigger.Dimensions[] から
    DBInstanceIdentifier を取得。
    """
    trigger = event.get("Trigger", {})
    dims = trigger.get("Dimensions", [])
    for d in dims:
        name = d.get("name") or d.get("Name")
        value = d.get("value") or d.get("Value")
        if name == "DBInstanceIdentifier":
            return value
    return None


def is_recently_rebooted(db_instance_id: str) -> bool:
    """
    DynamoDB に記録された最終再起動時刻と現在時刻を比較し、
    MIN_REBOOT_INTERVAL_MINUTES 以内なら True を返す。
    """
    res = dynamo.get_item(
        TableName=DDB_TABLE_NAME,
        Key={"DbInstanceIdentifier": {"S": db_instance_id}},
        ConsistentRead=True
    )
    item = res.get("Item")
    if not item:
        return False

    last_str = item.get("LastRebootTime", {}).get("S")
    if not last_str:
        return False

    last_dt = datetime.fromisoformat(last_str)
    now = datetime.now(timezone.utc)

    diff = now - last_dt
    print(f"Last reboot diff minutes: {diff.total_seconds() / 60:.1f}")
    return diff < timedelta(minutes=MIN_REBOOT_INTERVAL_MINUTES)

arn:aws:cloudwatch:ap-northeast-1:123456789012:alarm:rds-myapp-prod-rdsinst-high-cpu-auto-reboot
arn:aws:cloudwatch:ap-northeast-1:123456789012:alarm:rds-myapp-prod-rdsinst2-high-cpu-auto-reboot

# writer アラームからの Invoke を許可
aws lambda add-permission \
  --region ap-northeast-1 \
  --function-name rds-auto-reboot-on-high-cpu \
  --statement-id AllowCloudWatchInvokeWriterRdsCpuAlarm \
  --action lambda:InvokeFunction \
  --principal lambda.alarms.cloudwatch.amazonaws.com \
  --source-account 123456789012 \
  --source-arn arn:aws:cloudwatch:ap-northeast-1:123456789012:alarm:rds-myapp-prod-rdsinst-high-cpu-auto-reboot

# reader アラームからの Invoke を許可
aws lambda add-permission \
  --region ap-northeast-1 \
  --function-name rds-auto-reboot-on-high-cpu \
  --statement-id AllowCloudWatchInvokeReaderRdsCpuAlarm \
  --action lambda:InvokeFunction \
  --principal lambda.alarms.cloudwatch.amazonaws.com \
  --source-account 123456789012 \
  --source-arn arn:aws:cloudwatch:ap-northeast-1:123456789012:alarm:rds-myapp-prod-rdsinst2-high-cpu-auto-reboot

Event: {"source": "aws.cloudwatch", "alarmArn": "..."}
Target DB instance: myapp-prod-rdsinst
...

Reboot started: myapp-prod-rdsinst