Measuring the Impact of Round-Trip Latency Between the App and the DB on Response Time with TiDB Cloud Starter

"""TiDB Cloud Starter で、通信遅延がアプリの 1 リクエストの応答時間に与える影響を計測する。

接続情報は環境変数 TIDB_HOST, TIDB_PORT, TIDB_USER, TIDB_PASSWORD から読む。
遅延は Toxiproxy (管理 API: TOXIPROXY_API、既定 http://127.0.0.1:8474) で加える。

    python latency_bench.py setup                # テーブル作成とデータ投入
    python latency_bench.py proxy                # Toxiproxy にプロキシ tidb を登録
    python latency_bench.py probe 0,20,80        # 接続時間と SELECT 1 の往復時間
    python latency_bench.py check 0,20,80        # 2 つの処理の結果一致と、遅延による増分の確認
    python latency_bench.py bench 0,20,80        # 本計測: 遅延 × 逐次取得・一括取得 (1 行 1 試行の JSON)
    python latency_bench.py conn 0,80            # 本計測: 接続を毎回作る場合と使い回す場合
    python latency_bench.py summarize bench.jsonl conn.jsonl  # 中央値と増分の集計
"""

import json
import os
import socket
import ssl
import statistics
import sys
import time
import urllib.request

import pymysql

DB = "latency_lab"
PROXY_NAME = "tidb"
PROXY_LISTEN = ("127.0.0.1", 13306)
API = os.environ.get("TOXIPROXY_API", "http://127.0.0.1:8474")

SETUP = [
    f"CREATE DATABASE IF NOT EXISTS {DB}",
    f"USE {DB}",
    "CREATE TABLE IF NOT EXISTS digit (n INT PRIMARY KEY)",
    "INSERT IGNORE INTO digit VALUES (0),(1),(2),(3),(4),(5),(6),(7),(8),(9)",
    """CREATE TABLE IF NOT EXISTS orders (
      id BIGINT PRIMARY KEY,
      customer_id INT NOT NULL,
      ordered_at DATETIME NOT NULL,
      KEY idx_customer (customer_id, ordered_at DESC)
    )""",
    """CREATE TABLE IF NOT EXISTS order_items (
      order_id BIGINT NOT NULL,
      line_no TINYINT NOT NULL,
      product_name VARCHAR(64) NOT NULL,
      qty TINYINT NOT NULL,
      unit_price INT NOT NULL,
      PRIMARY KEY (order_id, line_no)
    )""",
]

# 顧客 1,000 人 × 注文 10 件 = 注文 10,000 件、注文ごとに明細 5 件
LOAD = [
    """INSERT INTO orders
    SELECT x.id, 1 + MOD(x.id - 1, 1000),
           DATE_ADD('2026-09-01', INTERVAL x.id MINUTE)
    FROM (
      SELECT 1 + a.n + 10*b.n + 100*c.n + 1000*d.n AS id
      FROM digit a CROSS JOIN digit b CROSS JOIN digit c CROSS JOIN digit d
    ) AS x""",
    """INSERT INTO order_items
    SELECT o.id, e.n + 1, CONCAT('product-', MOD(o.id * 7 + e.n, 500)),
           1 + MOD(e.n, 3), 100 + MOD(o.id + e.n, 900)
    FROM orders AS o CROSS JOIN digit AS e WHERE e.n < 5""",
]

ORDERS_SQL = """SELECT id, ordered_at FROM orders
WHERE customer_id = %s ORDER BY ordered_at DESC LIMIT 5"""
ITEMS_ONE_SQL = """SELECT order_id, line_no, product_name, qty, unit_price
FROM order_items WHERE order_id = %s ORDER BY line_no"""
ITEMS_IN_SQL = """SELECT order_id, line_no, product_name, qty, unit_price
FROM order_items WHERE order_id IN ({}) ORDER BY order_id, line_no"""

def tls_context():
    ctx = ssl.create_default_context()
    ctx.check_hostname = True
    ctx.verify_mode = ssl.CERT_REQUIRED
    return ctx

def connect(via_proxy, database=DB):
    """TiDB に接続する。via_proxy なら TCP を Toxiproxy に向け、TLS のホスト名は TiDB のものを検証する。"""
    conn = pymysql.connect(
        host=os.environ["TIDB_HOST"],
        port=int(os.environ.get("TIDB_PORT", "4000")),
        user=os.environ["TIDB_USER"],
        password=os.environ["TIDB_PASSWORD"],
        database=database,
        ssl=tls_context(),
        autocommit=True,
        defer_connect=True,
    )
    sock = socket.create_connection(PROXY_LISTEN, timeout=30) if via_proxy else None
    conn.connect(sock=sock)
    return conn

def api(method, path, body=None):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(API + path, data=data, method=method,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as res:
        return res.read()

def set_delay(rtt_ms):
    """往復で rtt_ms だけ遅れるよう、上り・下りに半分ずつ latency を設定する。0 なら外す。"""
    for stream in ("upstream", "downstream"):
        name = f"latency_{stream}"
        try:
            api("DELETE", f"/proxies/{PROXY_NAME}/toxics/{name}")
        except urllib.error.HTTPError:
            pass
        if rtt_ms > 0:
            api("POST", f"/proxies/{PROXY_NAME}/toxics",
                {"name": name, "type": "latency", "stream": stream,
                 "attributes": {"latency": rtt_ms // 2, "jitter": 0}})

def setup():
    cur = connect(False, None).cursor()
    for sql in SETUP:
        cur.execute(sql)
    cur.execute("SELECT COUNT(*) FROM orders")
    if cur.fetchone()[0] == 0:
        for sql in LOAD:
            cur.execute(sql)
    cur.execute("SELECT (SELECT COUNT(*) FROM orders), (SELECT COUNT(*) FROM order_items)")
    print("orders, order_items:", cur.fetchone())

def proxy():
    try:
        api("DELETE", f"/proxies/{PROXY_NAME}")
    except urllib.error.HTTPError:
        pass
    api("POST", "/proxies", {"name": PROXY_NAME, "listen": "0.0.0.0:13306",
                             "upstream": f"{os.environ['TIDB_HOST']}:4000", "enabled": True})
    print(api("GET", f"/proxies/{PROXY_NAME}").decode())

def timed(cur, sql, args=None):
    start = time.perf_counter()
    cur.execute(sql, args)
    rows = cur.fetchall()
    return rows, time.perf_counter() - start

def order_page_sequential(cur, customer_id, sql_sec=None):
    """注文 5 件を取り、注文ごとに明細を取る (SQL 6 回)。"""
    sql_sec = [] if sql_sec is None else sql_sec
    orders, sec = timed(cur, ORDERS_SQL, (customer_id,))
    sql_sec.append(sec)
    page = []
    for order_id, ordered_at in orders:
        items, sec = timed(cur, ITEMS_ONE_SQL, (order_id,))
        sql_sec.append(sec)
        page.append({"id": order_id, "ordered_at": str(ordered_at), "items": [list(i[1:]) for i in items]})
    return page

def order_page_batched(cur, customer_id, sql_sec=None):
    """注文 5 件を取り、明細を IN 句で 1 回で取る (SQL 2 回)。"""
    sql_sec = [] if sql_sec is None else sql_sec
    orders, sec = timed(cur, ORDERS_SQL, (customer_id,))
    sql_sec.append(sec)
    ids = [o[0] for o in orders]
    items, sec = timed(cur, ITEMS_IN_SQL.format(",".join(["%s"] * len(ids))), ids)
    sql_sec.append(sec)
    by_order = {}
    for item in items:
        by_order.setdefault(item[0], []).append(list(item[1:]))
    return [{"id": oid, "ordered_at": str(at), "items": by_order.get(oid, [])} for oid, at in orders]

def summary(values_ms):
    q1, median, q3 = statistics.quantiles(values_ms, n=4, method="inclusive")
    return {"median": round(median, 2), "q1": round(q1, 2), "q3": round(q3, 2)}

def probe(delays, n=30):
    """接続確立の時間と SELECT 1 の往復時間を、プロキシなしと各遅延設定で測る。"""
    for label, via_proxy, delay in [("direct", False, 0)] + [(f"proxy+{d}", True, d) for d in delays]:
        if via_proxy:
            set_delay(delay)
        connect_ms, select1_ms = [], []
        for _ in range(5):
            start = time.perf_counter()
            conn = connect(via_proxy)
            connect_ms.append((time.perf_counter() - start) * 1000)
            conn.close()
        conn = connect(via_proxy)
        cur = conn.cursor()
        for _ in range(3):
            timed(cur, "SELECT 1")
        for _ in range(n):
            select1_ms.append(timed(cur, "SELECT 1")[1] * 1000)
        conn.close()
        print(json.dumps({"path": label, "connect_ms": summary(connect_ms),
                          "select1_ms": summary(select1_ms)}))
    set_delay(0)

def check(delays, n=5):
    """2 つの処理の結果一致と、各遅延設定での応答時間を数回だけ確かめる。"""
    for delay in delays:
        set_delay(delay)
        conn = connect(True)
        cur = conn.cursor()
        seq_ms, bat_ms = [], []
        for i in range(n):
            customer_id = 1 + (i * 37) % 1000
            start = time.perf_counter()
            a = order_page_sequential(cur, customer_id)
            seq_ms.append((time.perf_counter() - start) * 1000)
            start = time.perf_counter()
            b = order_page_batched(cur, customer_id)
            bat_ms.append((time.perf_counter() - start) * 1000)
            assert a == b and len(a) == 5 and all(len(o["items"]) == 5 for o in a), customer_id
        conn.close()
        print(json.dumps({"delay_ms": delay, "identical": True,
                          "sequential_ms": [round(x, 1) for x in seq_ms],
                          "batched_ms": [round(x, 1) for x in bat_ms]}))
    set_delay(0)

def select1_series(cur, n=30):
    return summary([timed(cur, "SELECT 1")[1] * 1000 for _ in range(n)])

def emit(row):
    print(json.dumps(row), flush=True)

def bench(delays, warmup=3, rounds=30, block=10):
    """遅延 × 逐次取得・一括取得。1 本の接続を使い回し、実装の実行順は周回ごと、遅延の順序はまとまりごとに入れ替える。"""
    conn = connect(True)
    cur = conn.cursor()
    pages = {"sequential": order_page_sequential, "batched": order_page_batched}
    for delay in delays:
        set_delay(delay)
        for i in range(warmup):
            for page in pages.values():
                page(cur, 1 + i)
        emit({"kind": "select1", "delay_ms": delay, "when": "before", **select1_series(cur)})
    for start in range(0, rounds, block):
        order = delays if (start // block) % 2 == 0 else list(reversed(delays))
        for delay in order:
            set_delay(delay)
            for i in range(start, start + block):
                customer_id = 1 + (i * 37) % 1000
                names = ["sequential", "batched"] if i % 2 == 0 else ["batched", "sequential"]
                for name in names:
                    sql_sec = []
                    t0 = time.perf_counter()
                    pages[name](cur, customer_id, sql_sec)
                    total = time.perf_counter() - t0
                    emit({"kind": "page", "delay_ms": delay, "round": i, "version": name,
                          "customer_id": customer_id, "total_ms": round(total * 1000, 2),
                          "sql_ms": [round(x * 1000, 2) for x in sql_sec]})
    for delay in delays:
        set_delay(delay)
        emit({"kind": "select1", "delay_ms": delay, "when": "after", **select1_series(cur)})
    set_delay(0)
    conn.close()

def conn_bench(delays, warmup=3, rounds=30):
    """一括取得の処理を、接続を毎回作る場合 (new) と使い回す場合 (reuse) で交互に実行する。"""
    for delay in delays:
        set_delay(delay)
        reuse = connect(True)
        reuse_cur = reuse.cursor()
        for i in range(warmup + rounds):
            customer_id = 1 + (i * 37) % 1000
            for mode in (["new", "reuse"] if i % 2 == 0 else ["reuse", "new"]):
                t0 = time.perf_counter()
                if mode == "new":
                    conn = connect(True)
                    connect_sec = time.perf_counter() - t0
                    order_page_batched(conn.cursor(), customer_id)
                    total = time.perf_counter() - t0
                    conn.close()
                else:
                    connect_sec = 0.0
                    order_page_batched(reuse_cur, customer_id)
                    total = time.perf_counter() - t0
                if i >= warmup:
                    emit({"kind": "conn", "delay_ms": delay, "round": i - warmup, "mode": mode,
                          "total_ms": round(total * 1000, 2), "connect_ms": round(connect_sec * 1000, 2)})
        reuse.close()
    set_delay(0)

def summarize(bench_path, conn_path):
    rows = [json.loads(line) for line in open(bench_path)]
    pages = [r for r in rows if r["kind"] == "page"]
    delays = sorted({r["delay_ms"] for r in pages})
    for version in ("sequential", "batched"):
        base = None
        for delay in delays:
            ms = summary([r["total_ms"] for r in pages if r["version"] == version and r["delay_ms"] == delay])
            base = ms["median"] if base is None else base
            print(f"{version:10} +{delay:>2} ms: median {ms['median']:7.1f} "
                  f"(q1 {ms['q1']:6.1f}, q3 {ms['q3']:6.1f})  increase {ms['median'] - base:6.1f}")
    rows = [json.loads(line) for line in open(conn_path)]
    for delay in sorted({r["delay_ms"] for r in rows}):
        for mode in ("new", "reuse"):
            target = [r for r in rows if r["delay_ms"] == delay and r["mode"] == mode]
            total = summary([r["total_ms"] for r in target])
            connect_ms = summary([r["connect_ms"] for r in target])
            print(f"conn {mode:5} +{delay:>2} ms: median {total['median']:7.1f} "
                  f"(q1 {total['q1']:6.1f}, q3 {total['q3']:6.1f})  connect {connect_ms['median']:6.1f}")

if __name__ == "__main__":
    command = sys.argv[1]
    if command == "summarize":
        summarize(sys.argv[2], sys.argv[3])
        sys.exit()
    if command in ("probe", "check", "bench", "conn"):
        {"probe": probe, "check": check, "bench": bench, "conn": conn_bench}[command](
            [int(d) for d in sys.argv[2].split(",")])
    else:
        {"setup": setup, "proxy": proxy}[command]()