Aurora DSQL Python Connector がリリースされたので、Lambda で性能を実測比較してみた (psycopg2 vs asyncpg)

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: Aurora DSQL Connector Benchmark on Lambda

Parameters:
  DsqlEndpoint:
    Type: String
  DsqlRegion:
    Type: String
    Default: us-west-2
  BenchIterations:
    Type: String
    Default: '30'

Globals:
  Function:
    Runtime: python3.12
    Timeout: 300
    MemorySize: 256
    Environment:
      Variables:
        DSQL_ENDPOINT: !Ref DsqlEndpoint
        DSQL_REGION: !Ref DsqlRegion
        BENCH_ITERATIONS: !Ref BenchIterations

Resources:
  Psycopg2Layer:
    Type: AWS::Serverless::LayerVersion
    Properties:
      LayerName: dsql-bench-psycopg2-x86
      ContentUri: layers/psycopg2-x86/
      CompatibleRuntimes: [python3.12]
      CompatibleArchitectures: [x86_64]
    Metadata:
      BuildMethod: python3.12
      BuildArchitecture: x86_64

  AsyncpgLayer:
    Type: AWS::Serverless::LayerVersion
    Properties:
      LayerName: dsql-bench-asyncpg-arm64
      ContentUri: layers/asyncpg-arm64/
      CompatibleRuntimes: [python3.12]
      CompatibleArchitectures: [arm64]
    Metadata:
      BuildMethod: python3.12
      BuildArchitecture: arm64

  BenchX86:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: dsql-bench-x86
      Handler: handler.lambda_handler
      CodeUri: src/
      Architectures: [x86_64]
      Layers: [!Ref Psycopg2Layer]
      Policies:
        - Statement:
            - Effect: Allow
              Action: dsql:DbConnectAdmin
              Resource: !Sub arn:aws:dsql:${DsqlRegion}:${AWS::AccountId}:cluster/*

  BenchArm64:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: dsql-bench-arm64
      Handler: handler.lambda_handler
      CodeUri: src/
      Architectures: [arm64]
      Layers: [!Ref AsyncpgLayer]
      Policies:
        - Statement:
            - Effect: Allow
              Action: dsql:DbConnectAdmin
              Resource: !Sub arn:aws:dsql:${DsqlRegion}:${AWS::AccountId}:cluster/*

import asyncio
import json
import os
import platform
import statistics
import time

DSQL_ENDPOINT = os.environ['DSQL_ENDPOINT']
DSQL_REGION = os.environ.get('DSQL_REGION', 'us-west-2')
ITERATIONS = int(os.environ.get('BENCH_ITERATIONS', '30'))
QUERY = 'SELECT 1'
ARCH = platform.machine()

def fmt(s):
    return round(s * 1000, 2)

def calc_stats(connect_times, query_times):
    n = len(connect_times)
    total = sum(connect_times) + sum(query_times)
    return {
        'connect_avg_ms': fmt(statistics.mean(connect_times)),
        'connect_p50_ms': fmt(statistics.median(connect_times)),
        'connect_p95_ms': fmt(sorted(connect_times)[int(n * 0.95)]),
        'query_avg_ms': fmt(statistics.mean(query_times)),
        'query_p50_ms': fmt(statistics.median(query_times)),
        'total_ms': fmt(total),
        'qps': round(n / total, 1),
    }

def bench_psycopg2():
    import aurora_dsql_psycopg2 as dsql
    ct, qt = [], []
    for _ in range(ITERATIONS):
        t0 = time.perf_counter()
        conn = dsql.connect(host=DSQL_ENDPOINT, region=DSQL_REGION, user='admin')
        t1 = time.perf_counter()
        ct.append(t1 - t0)
        cur = conn.cursor()
        cur.execute(QUERY)
        cur.fetchall()
        qt.append(time.perf_counter() - t1)
        cur.close()
        conn.close()
    return {'pattern': 'psycopg2 (every connect)', **calc_stats(ct, qt)}

async def bench_asyncpg_no_pool():
    import aurora_dsql_asyncpg as dsql
    ct, qt = [], []
    for _ in range(ITERATIONS):
        t0 = time.perf_counter()
        conn = await dsql.connect(host=DSQL_ENDPOINT, region=DSQL_REGION, user='admin')
        t1 = time.perf_counter()
        ct.append(t1 - t0)
        await conn.fetch(QUERY)
        qt.append(time.perf_counter() - t1)
        await conn.close()
    return {'pattern': 'asyncpg (every connect)', **calc_stats(ct, qt)}

async def bench_asyncpg_reuse():
    import aurora_dsql_asyncpg as dsql
    t0 = time.perf_counter()
    pool = await dsql.create_pool(
        host=DSQL_ENDPOINT, region=DSQL_REGION, user='admin',
        min_size=1, max_size=3)
    pool_create = time.perf_counter() - t0
    ct, qt = [], []
    for _ in range(ITERATIONS):
        t1 = time.perf_counter()
        async with pool.acquire() as conn:
            t2 = time.perf_counter()
            ct.append(t2 - t1)
            await conn.fetch(QUERY)
            qt.append(time.perf_counter() - t2)
    await pool.close()
    stats = calc_stats(ct, qt)
    stats['pool_create_ms'] = fmt(pool_create)
    return {'pattern': 'asyncpg (connection reuse)', **stats}

def lambda_handler(event, context):
    results = []
    results.append(bench_psycopg2())
    loop = asyncio.get_event_loop()
    results.append(loop.run_until_complete(bench_asyncpg_no_pool()))
    results.append(loop.run_until_complete(bench_asyncpg_reuse()))
    body = {
        'arch': ARCH,
        'runtime': f"python{platform.python_version()}",
        'iterations': ITERATIONS,
        'query': QUERY,
        'endpoint': DSQL_ENDPOINT,
        'results': results,
    }
    print(json.dumps(body, ensure_ascii=False, indent=2))
    return body