CloudFront Origin GroupのMembers入れ替えをStep FunctionsとLambdaで試してみた

{
  "Comment": "CloudFront Origin Group - Primary側3回失敗でMembers入れ替え",
  "QueryLanguage": "JSONata",
  "StartAt": "Initialize",
  "States": {
    "Initialize": {
      "Type": "Pass",
      "Comment": "入力値を変数に格納 & カウンター初期化",
      "Assign": {
        "primaryUrl": "{% $states.input.primaryUrl %}",
        "secondaryUrl": "{% $states.input.secondaryUrl %}",
        "connectionArn": "{% $states.input.connectionArn %}",
        "distributionId": "{% $states.input.distributionId %}",
        "counter": 0,
        "primaryFailures": 0,
        "secondaryOK": 0
      },
      "Next": "HealthCheck"
    },
    "HealthCheck": {
      "Type": "Parallel",
      "Comment": "Primary/Secondaryに並行ヘルスチェック",
      "Branches": [
        {
          "StartAt": "CheckPrimary",
          "States": {
            "CheckPrimary": {
              "Type": "Task",
              "Resource": "arn:aws:states:::http:invoke",
              "Arguments": {
                "ApiEndpoint": "{% $primaryUrl %}",
                "Method": "GET",
                "Authentication": {
                  "ConnectionArn": "{% $connectionArn %}"
                }
              },
              "Output": {
                "statusCode": "{% $states.result.StatusCode %}",
                "healthy": true
              },
              "Catch": [
                {
                  "ErrorEquals": ["States.ALL"],
                  "Next": "PrimaryFailed"
                }
              ],
              "End": true
            },
            "PrimaryFailed": {
              "Type": "Pass",
              "Output": {
                "statusCode": 0,
                "healthy": false
              },
              "End": true
            }
          }
        },
        {
          "StartAt": "CheckSecondary",
          "States": {
            "CheckSecondary": {
              "Type": "Task",
              "Resource": "arn:aws:states:::http:invoke",
              "Arguments": {
                "ApiEndpoint": "{% $secondaryUrl %}",
                "Method": "GET",
                "Authentication": {
                  "ConnectionArn": "{% $connectionArn %}"
                }
              },
              "Output": {
                "statusCode": "{% $states.result.StatusCode %}",
                "healthy": true
              },
              "Catch": [
                {
                  "ErrorEquals": ["States.ALL"],
                  "Next": "SecondaryFailed"
                }
              ],
              "End": true
            },
            "SecondaryFailed": {
              "Type": "Pass",
              "Output": {
                "statusCode": 0,
                "healthy": false
              },
              "End": true
            }
          }
        }
      ],
      "Assign": {
        "counter": "{% $counter + 1 %}",
        "primaryFailures": "{% $states.result[0].healthy = false ? $primaryFailures + 1 : $primaryFailures %}",
        "secondaryOK": "{% $states.result[1].healthy = true ? $secondaryOK + 1 : $secondaryOK %}"
      },
      "Next": "CheckLoopComplete"
    },
    "CheckLoopComplete": {
      "Type": "Choice",
      "Comment": "3回チェック完了か判定",
      "Choices": [
        {
          "Condition": "{% $counter >= 3 %}",
          "Next": "EvaluateResults"
        }
      ],
      "Default": "WaitBetweenChecks"
    },
    "WaitBetweenChecks": {
      "Type": "Wait",
      "Seconds": 20,
      "Next": "HealthCheck"
    },
    "EvaluateResults": {
      "Type": "Choice",
      "Comment": "Primary全滅 & Secondary全正常ならMembers入れ替え",
      "Choices": [
        {
          "Condition": "{% $primaryFailures = 3 and $secondaryOK = 3 %}",
          "Next": "GetDistributionConfig"
        }
      ],
      "Default": "NoActionNeeded"
    },
    "NoActionNeeded": {
      "Type": "Succeed",
      "Comment": "入れ替え条件未達、正常終了"
    },
    "GetDistributionConfig": {
      "Type": "Task",
      "Resource": "arn:aws:states:::aws-sdk:cloudfront:getDistributionConfig",
      "Arguments": {
        "Id": "{% $distributionId %}"
      },
      "Assign": {
        "etag": "{% $states.result.ETag %}",
        "distributionConfig": "{% $states.result.DistributionConfig %}"
      },
      "Next": "SwapOriginGroupMembers"
    },
    "SwapOriginGroupMembers": {
      "Type": "Pass",
      "Comment": "Origin GroupのMembers順序を入れ替え（Primary⇔Secondary）",
      "Assign": {
        "updatedConfig": "{% ( $swapped := $distributionConfig.OriginGroups.Items.{ 'Id': Id, 'FailoverCriteria': FailoverCriteria, 'Members': { 'Quantity': Members.Quantity, 'Items': $reverse(Members.Items) } }; $merge([$distributionConfig, { 'OriginGroups': { 'Quantity': $distributionConfig.OriginGroups.Quantity, 'Items': $type($swapped) = 'array' ? $swapped : [$swapped] } }]) ) %}"
      },
      "Next": "UpdateDistribution"
    },
    "UpdateDistribution": {
      "Type": "Task",
      "Resource": "arn:aws:states:::aws-sdk:cloudfront:updateDistribution",
      "Arguments": {
        "Id": "{% $distributionId %}",
        "IfMatch": "{% $etag %}",
        "DistributionConfig": "{% $updatedConfig %}"
      },
      "Output": {
        "status": "UPDATE_ACCEPTED",
        "distributionId": "{% $distributionId %}",
        "etag": "{% $states.result.ETag %}"
      },
      "End": true
    }
  }
}

import json
import os
import time
import urllib.request
import urllib.error
import ssl

import boto3

def get_config():
    return {
        "distribution_id": os.environ["DISTRIBUTION_ID"],
        "origin_group_id": os.environ.get("ORIGIN_GROUP_ID", "failover-group"),
        "check_count": int(os.environ.get("CHECK_COUNT", "5")),
        "check_interval": int(os.environ.get("CHECK_INTERVAL", "10")),
        "max_latency": float(os.environ.get("MAX_LATENCY", "1.0")),
        "max_error_rate": float(os.environ.get("MAX_ERROR_RATE", "0.0")),
        "health_check_path": os.environ.get("HEALTH_CHECK_PATH", "/"),
    }

def get_distribution_config(client, distribution_id):
    resp = client.get_distribution_config(Id=distribution_id)
    return resp["DistributionConfig"], resp["ETag"]

def find_origin_group(distribution_config, origin_group_id):
    origin_groups = distribution_config.get("OriginGroups", {})
    for group in origin_groups.get("Items", []):
        if group["Id"] == origin_group_id:
            return group
    raise ValueError(f"Origin Group '{origin_group_id}' not found")

def get_origin_domain(distribution_config, origin_id):
    for origin in distribution_config["Origins"]["Items"]:
        if origin["Id"] == origin_id:
            return origin["DomainName"]
    raise ValueError(f"Origin '{origin_id}' not found")

def health_check_once(domain, path, timeout=10):
    url = f"https://{domain}{path}"
    ctx = ssl.create_default_context()
    start = time.time()
    try:
        req = urllib.request.Request(url, method="GET")
        req.add_header("User-Agent", "CloudFront-Failover-HealthCheck/1.0")
        with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
            latency = time.time() - start
            return {"success": 200 <= resp.status < 400, "status_code": resp.status, "latency": round(latency, 3)}
    except urllib.error.HTTPError as e:
        latency = time.time() - start
        return {"success": False, "status_code": e.code, "latency": round(latency, 3)}
    except Exception as e:
        latency = time.time() - start
        return {"success": False, "status_code": 0, "latency": round(latency, 3), "error": str(e)}

def run_health_checks(domain, path, count, interval):
    results = []
    for i in range(count):
        if i > 0:
            time.sleep(interval)
        result = health_check_once(domain, path)
        results.append(result)
        print(f"  Check {i+1}/{count}: status={result.get('status_code')} latency={result['latency']}s success={result['success']}")
    return results

def evaluate_health(results, max_latency, max_error_rate):
    if not results:
        return False, "No results"
    total = len(results)
    errors = sum(1 for r in results if not r["success"])
    error_rate = errors / total
    avg_latency = sum(r["latency"] for r in results) / total
    summary = f"avg_latency={avg_latency:.3f}s (max={max_latency}s), error_rate={error_rate:.0%} (max={max_error_rate:.0%})"
    healthy = avg_latency <= max_latency and error_rate <= max_error_rate
    return healthy, summary

def swap_origin_group_members(distribution_config, origin_group_id):
    origin_groups = distribution_config.get("OriginGroups", {})
    for group in origin_groups.get("Items", []):
        if group["Id"] == origin_group_id:
            members = group["Members"]["Items"]
            if len(members) != 2:
                raise ValueError(f"Origin Group '{origin_group_id}' has {len(members)} members, expected 2")
            group["Members"]["Items"] = [members[1], members[0]]
            return members[0]["OriginId"], members[1]["OriginId"]
    raise ValueError(f"Origin Group '{origin_group_id}' not found")

def handler(event, context):
    config = get_config()
    print(f"Config: distribution_id={config['distribution_id']}, origin_group_id={config['origin_group_id']}")

    cf_client = boto3.client("cloudfront")
    distribution_config, etag = get_distribution_config(cf_client, config["distribution_id"])

    origin_group = find_origin_group(distribution_config, config["origin_group_id"])
    members = origin_group["Members"]["Items"]
    primary_id = members[0]["OriginId"]
    secondary_id = members[1]["OriginId"]

    primary_domain = get_origin_domain(distribution_config, primary_id)
    secondary_domain = get_origin_domain(distribution_config, secondary_id)

    print(f"Current order:")
    print(f"  Primary:   {primary_id} ({primary_domain})")
    print(f"  Secondary: {secondary_id} ({secondary_domain})")

    print(f"\nHealth checking secondary: {secondary_domain}")
    results = run_health_checks(
        domain=secondary_domain,
        path=config["health_check_path"],
        count=config["check_count"],
        interval=config["check_interval"],
    )

    healthy, summary = evaluate_health(results, config["max_latency"], config["max_error_rate"])
    print(f"\nSecondary direct check: {'PASSED' if healthy else 'FAILED'} - {summary}")

    if not healthy:
        msg = f"Secondary direct check failed. No member swap requested. {summary}"
        print(msg)
        return {
            "action": "none",
            "reason": "secondary_direct_check_failed",
            "summary": summary,
            "primary": primary_id,
            "secondary": secondary_id,
        }

    print(f"\nSwapping origin group members: {primary_id} <-> {secondary_id}")
    old_primary, old_secondary = swap_origin_group_members(distribution_config, config["origin_group_id"])

    new_etag = cf_client.update_distribution(
        Id=config["distribution_id"],
        IfMatch=etag,
        DistributionConfig=distribution_config,
    )["ETag"]

    print(f"Distribution update accepted. New ETag: {new_etag}")
    print(f"Requested order: Primary={old_secondary}, Secondary={old_primary}")

    return {
        "action": "swapped",
        "old_primary": old_primary,
        "old_secondary": old_secondary,
        "new_primary": old_secondary,
        "new_secondary": old_primary,
        "health_summary": summary,
    }

AWSTemplateFormatVersion: '2010-09-09'
Description: >-
  Test environment for CloudFront Origin Group Failover Lambda.
  Creates a CloudFront distribution with origin group (failover-group)
  using Lambda Function URLs as primary/secondary origins.

Parameters:
  ProjectName:
    Type: String
    Default: cf-failover-test

  PrimaryDelaySeconds:
    Type: Number
    Default: 0
    Description: Primary Lambda response delay in seconds

  SecondaryDelaySeconds:
    Type: Number
    Default: 0
    Description: Secondary Lambda response delay in seconds (keep 0 for healthy secondary)

  HealthCheckCount:
    Type: Number
    Default: 5
    Description: Number of health checks the failover Lambda performs

  HealthCheckInterval:
    Type: Number
    Default: 10
    Description: Interval between health checks in seconds

Resources:
  OriginLambdaRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub ${ProjectName}-origin-role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  PrimaryFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub ${ProjectName}-primary
      Runtime: python3.12
      Handler: index.handler
      Role: !GetAtt OriginLambdaRole.Arn
      Timeout: 30
      MemorySize: 128
      Environment:
        Variables:
          DELAY_SECONDS: !Ref PrimaryDelaySeconds
      Code:
        ZipFile: |
          import json, os, time
          def handler(event, context):
              delay = int(os.environ.get("DELAY_SECONDS", "0"))
              if delay > 0:
                  time.sleep(delay)
              return {
                  "statusCode": 200,
                  "headers": {"Content-Type": "application/json"},
                  "body": json.dumps({"origin": "primary", "delay": delay})
              }

  PrimaryFunctionUrl:
    Type: AWS::Lambda::Url
    Properties:
      AuthType: NONE
      TargetFunctionArn: !GetAtt PrimaryFunction.Arn

  PrimaryFunctionUrlPermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt PrimaryFunction.Arn
      Action: lambda:InvokeFunctionUrl
      Principal: '*'
      FunctionUrlAuthType: NONE

  PrimaryFunctionInvokePermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt PrimaryFunction.Arn
      Action: lambda:InvokeFunction
      Principal: '*'
      InvokedViaFunctionUrl: true

  SecondaryFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub ${ProjectName}-secondary
      Runtime: python3.12
      Handler: index.handler
      Role: !GetAtt OriginLambdaRole.Arn
      Timeout: 30
      MemorySize: 128
      Environment:
        Variables:
          DELAY_SECONDS: !Ref SecondaryDelaySeconds
      Code:
        ZipFile: |
          import json, os, time
          def handler(event, context):
              delay = int(os.environ.get("DELAY_SECONDS", "0"))
              if delay > 0:
                  time.sleep(delay)
              return {
                  "statusCode": 200,
                  "headers": {"Content-Type": "application/json"},
                  "body": json.dumps({"origin": "secondary", "delay": delay})
              }

  SecondaryFunctionUrl:
    Type: AWS::Lambda::Url
    Properties:
      AuthType: NONE
      TargetFunctionArn: !GetAtt SecondaryFunction.Arn

  SecondaryFunctionUrlPermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt SecondaryFunction.Arn
      Action: lambda:InvokeFunctionUrl
      Principal: '*'
      FunctionUrlAuthType: NONE

  SecondaryFunctionInvokePermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt SecondaryFunction.Arn
      Action: lambda:InvokeFunction
      Principal: '*'
      InvokedViaFunctionUrl: true

  Distribution:
    Type: AWS::CloudFront::Distribution
    Properties:
      DistributionConfig:
        Enabled: true
        Comment: !Sub ${ProjectName} - Origin Group Failover Test
        DefaultCacheBehavior:
          TargetOriginId: failover-group
          ViewerProtocolPolicy: redirect-to-https
          CachePolicyId: 4135ea2d-6df8-44a3-9df3-4b5a84be39ad
          OriginRequestPolicyId: b689b0a8-53d0-40ab-baf2-68738e2966ac
          AllowedMethods: [GET, HEAD]
          Compress: true
        Origins:
          - Id: primary-origin
            DomainName: !Select [2, !Split ['/', !GetAtt PrimaryFunctionUrl.FunctionUrl]]
            CustomOriginConfig:
              HTTPSPort: 443
              OriginProtocolPolicy: https-only
              OriginSSLProtocols: [TLSv1.2]
          - Id: secondary-origin
            DomainName: !Select [2, !Split ['/', !GetAtt SecondaryFunctionUrl.FunctionUrl]]
            CustomOriginConfig:
              HTTPSPort: 443
              OriginProtocolPolicy: https-only
              OriginSSLProtocols: [TLSv1.2]
        OriginGroups:
          Quantity: 1
          Items:
            - Id: failover-group
              FailoverCriteria:
                StatusCodes:
                  Quantity: 4
                  Items: [500, 502, 503, 504]
              Members:
                Quantity: 2
                Items:
                  - OriginId: primary-origin
                  - OriginId: secondary-origin
        HttpVersion: http2and3
        PriceClass: PriceClass_100

  FailoverLambdaRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub ${ProjectName}-failover-role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: CloudFrontAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - cloudfront:GetDistribution
                  - cloudfront:GetDistributionConfig
                  - cloudfront:UpdateDistribution
                Resource: !Sub arn:aws:cloudfront::${AWS::AccountId}:distribution/${Distribution}

  FailoverFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub ${ProjectName}-failover
      Runtime: python3.12
      Handler: index.handler
      Role: !GetAtt FailoverLambdaRole.Arn
      Timeout: 120
      MemorySize: 128
      Environment:
        Variables:
          DISTRIBUTION_ID: !Ref Distribution
          ORIGIN_GROUP_ID: failover-group
          CHECK_COUNT: !Ref HealthCheckCount
          CHECK_INTERVAL: !Ref HealthCheckInterval
          MAX_LATENCY: '1.0'
          MAX_ERROR_RATE: '0.0'
          HEALTH_CHECK_PATH: /
      Code:
        ZipFile: |
          import json, os, time, urllib.request, urllib.error, ssl
          import boto3

          def get_config():
              return {
                  "distribution_id": os.environ["DISTRIBUTION_ID"],
                  "origin_group_id": os.environ.get("ORIGIN_GROUP_ID", "failover-group"),
                  "check_count": int(os.environ.get("CHECK_COUNT", "5")),
                  "check_interval": int(os.environ.get("CHECK_INTERVAL", "10")),
                  "max_latency": float(os.environ.get("MAX_LATENCY", "1.0")),
                  "max_error_rate": float(os.environ.get("MAX_ERROR_RATE", "0.0")),
                  "health_check_path": os.environ.get("HEALTH_CHECK_PATH", "/"),
              }

          def get_distribution_config(client, distribution_id):
              resp = client.get_distribution_config(Id=distribution_id)
              return resp["DistributionConfig"], resp["ETag"]

          def find_origin_group(distribution_config, origin_group_id):
              for group in distribution_config.get("OriginGroups", {}).get("Items", []):
                  if group["Id"] == origin_group_id:
                      return group
              raise ValueError(f"Origin Group '{origin_group_id}' not found")

          def get_origin_domain(distribution_config, origin_id):
              for origin in distribution_config["Origins"]["Items"]:
                  if origin["Id"] == origin_id:
                      return origin["DomainName"]
              raise ValueError(f"Origin '{origin_id}' not found")

          def health_check_once(domain, path, timeout=10):
              url = f"https://{domain}{path}"
              ctx = ssl.create_default_context()
              start = time.time()
              try:
                  req = urllib.request.Request(url, method="GET")
                  req.add_header("User-Agent", "CloudFront-Failover-HealthCheck/1.0")
                  with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
                      latency = time.time() - start
                      return {"success": 200 <= resp.status < 400, "status_code": resp.status, "latency": round(latency, 3)}
              except urllib.error.HTTPError as e:
                  latency = time.time() - start
                  return {"success": False, "status_code": e.code, "latency": round(latency, 3)}
              except Exception as e:
                  latency = time.time() - start
                  return {"success": False, "status_code": 0, "latency": round(latency, 3), "error": str(e)}

          def run_health_checks(domain, path, count, interval):
              results = []
              for i in range(count):
                  if i > 0:
                      time.sleep(interval)
                  result = health_check_once(domain, path)
                  results.append(result)
                  print(f"  Check {i+1}/{count}: status={result.get('status_code')} latency={result['latency']}s success={result['success']}")
              return results

          def evaluate_health(results, max_latency, max_error_rate):
              if not results:
                  return False, "No results"
              total = len(results)
              errors = sum(1 for r in results if not r["success"])
              error_rate = errors / total
              avg_latency = sum(r["latency"] for r in results) / total
              summary = f"avg_latency={avg_latency:.3f}s (max={max_latency}s), error_rate={error_rate:.0%} (max={max_error_rate:.0%})"
              healthy = avg_latency <= max_latency and error_rate <= max_error_rate
              return healthy, summary

          def swap_origin_group_members(distribution_config, origin_group_id):
              for group in distribution_config.get("OriginGroups", {}).get("Items", []):
                  if group["Id"] == origin_group_id:
                      members = group["Members"]["Items"]
                      if len(members) != 2:
                          raise ValueError(f"Expected 2 members, got {len(members)}")
                      group["Members"]["Items"] = [members[1], members[0]]
                      return members[0]["OriginId"], members[1]["OriginId"]
              raise ValueError(f"Origin Group '{origin_group_id}' not found")

          def handler(event, context):
              config = get_config()
              print(f"Config: distribution_id={config['distribution_id']}, origin_group_id={config['origin_group_id']}")
              cf_client = boto3.client("cloudfront")
              distribution_config, etag = get_distribution_config(cf_client, config["distribution_id"])
              origin_group = find_origin_group(distribution_config, config["origin_group_id"])
              members = origin_group["Members"]["Items"]
              primary_id = members[0]["OriginId"]
              secondary_id = members[1]["OriginId"]
              primary_domain = get_origin_domain(distribution_config, primary_id)
              secondary_domain = get_origin_domain(distribution_config, secondary_id)
              print(f"Current order:\n  Primary:   {primary_id} ({primary_domain})\n  Secondary: {secondary_id} ({secondary_domain})")
              print(f"\nHealth checking secondary: {secondary_domain}")
              results = run_health_checks(secondary_domain, config["health_check_path"], config["check_count"], config["check_interval"])
              healthy, summary = evaluate_health(results, config["max_latency"], config["max_error_rate"])
              print(f"\nSecondary direct check: {'PASSED' if healthy else 'FAILED'} - {summary}")
              if not healthy:
                  print(f"Secondary direct check failed. No member swap requested. {summary}")
                  return {"action": "none", "reason": "secondary_direct_check_failed", "summary": summary}
              print(f"\nSwapping origin group members: {primary_id} <-> {secondary_id}")
              old_primary, old_secondary = swap_origin_group_members(distribution_config, config["origin_group_id"])
              new_etag = cf_client.update_distribution(Id=config["distribution_id"], IfMatch=etag, DistributionConfig=distribution_config)["ETag"]
              print(f"Distribution update accepted. New ETag: {new_etag}")
              print(f"Requested order: Primary={old_secondary}, Secondary={old_primary}")
              return {"action": "swapped", "old_primary": old_primary, "new_primary": old_secondary, "health_summary": summary}