源账号 source_account_id (cn-north-1)            目标账号 target_account_id (cn-north-1)
                    ┌──────────────────────────┐               ┌──────────────────────────┐
  模拟业务写入  ──► │ allinone_qa_LargeDataTable │               │ allinone_qa_LargeDataTable │
 (simulate_      │   ├─ PITR (全量基线来源)    │  Export S3    │   (Import from S3 新建)    │
  workload.py)   │   └─ Streams(增量来源)      │ ────────────► │                            │
                    └──────────┬───────────────┘   menglonq桶   └──────────▲───────────────┘
                               │ Stream                                     │ 跨账号写
                               ▼                                            │ (资源策略授权)
                    ┌──────────────────────────┐                          │
                    │ Lambda: ddb-xacct-        │ ─────────────────────────┘
                    │ replicator (事件源映射)    │
                    └──────────────────────────┘

"""
向 DynamoDB 表批量补写测试数据，schema 复刻表内现有"基础数据"。

用途：为 DynamoDB 跨账号迁移方案搭建测试基线数据。
表：allinone_qa_LargeDataTable（cn-north-1），主键 id(HASH,S) + timestamp(RANGE,S)。

显式锁定 default profile + cn-north-1，避免被环境变量 AWS_PROFILE=global / AWS_REGION 干扰。

用法：
    python load_data.py                # 默认补写 400000 条
    python load_data.py --count 100000 # 自定义条数
    python load_data.py --workers 16
"""

import argparse
import random
import string
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError

# ----------------------- 固定配置 -----------------------
PROFILE = "default"
REGION = "cn-north-1"
TABLE = "allinone_qa_LargeDataTable"

BATCH_SIZE = 25  # BatchWriteItem 单批上限
CATEGORIES = ["A", "B", "C", "D", "E"]
VERSIONS = ["v1", "v2", "v3"]
# 现有数据 random_string 长度约 2000+；用区间随机以贴近真实分布
RAND_MIN, RAND_MAX = 1800, 2200

_ALPHABET = string.ascii_letters + string.digits

# 进度计数（跨线程）
_lock = threading.Lock()
_written = 0
_failed = 0


def _make_client():
    """构造写表用的 DynamoDB client，连接池放大以支撑多线程并发。"""
    session = boto3.Session(profile_name=PROFILE, region_name=REGION)
    cfg = Config(
        max_pool_connections=64,
        retries={"max_attempts": 10, "mode": "adaptive"},
    )
    return session.client("dynamodb", config=cfg)


def _rand_string(n):
    return "".join(random.choices(_ALPHABET, k=n))


def _gen_item():
    """生成一条与现有基础数据同 schema 的 item（DynamoDB JSON 低层格式）。"""
    tag_count = random.randint(5, 10)
    return {
        "id": {"S": str(uuid.uuid4())},
        "timestamp": {"S": datetime.now(timezone.utc).isoformat()},
        "random_string": {"S": _rand_string(random.randint(RAND_MIN, RAND_MAX))},
        "category": {"S": random.choice(CATEGORIES)},
        "numeric_field": {"N": str(random.randint(1, 1_000_000))},
        "metadata": {
            "M": {
                "created_by": {"S": "load_script"},
                "version": {"S": random.choice(VERSIONS)},
                "tags": {"L": [{"S": f"tag_{i}"} for i in range(tag_count)]},
            }
        },
    }


def _write_batch(client, items):
    """写一批（≤25 条），对 UnprocessedItems 做指数退避重试。返回成功条数。"""
    request = {TABLE: [{"PutRequest": {"Item": it}} for it in items]}
    attempt = 0
    while request.get(TABLE):
        try:
            resp = client.batch_write_item(RequestItems=request)
        except ClientError as e:
            # 限流等可重试错误：退避后整批重试（已配 adaptive retries 兜底）
            attempt += 1
            if attempt > 10:
                raise
            time.sleep(min(2 ** attempt * 0.05, 5))
            continue
        unprocessed = resp.get("UnprocessedItems", {})
        if not unprocessed.get(TABLE):
            break
        request = unprocessed
        attempt += 1
        time.sleep(min(2 ** attempt * 0.05, 5))
    return len(items)


def _worker(client, n_items, report_every):
    """单线程任务：生成并写入 n_items 条。"""
    global _written, _failed
    buf = []
    done_local = 0
    for _ in range(n_items):
        buf.append(_gen_item())
        if len(buf) == BATCH_SIZE:
            done_local += _flush(client, buf)
            buf = []
            if done_local >= report_every:
                _bump(done_local, 0)
                done_local = 0
    if buf:
        done_local += _flush(client, buf)
    if done_local:
        _bump(done_local, 0)


def _flush(client, buf):
    global _failed
    try:
        return _write_batch(client, buf)
    except Exception as e:  # noqa: BLE001 - 单批失败不应中断整体载入
        with _lock:
            _failed += len(buf)
        print(f"[WARN] batch failed ({len(buf)} items): {e}")
        return 0


def _bump(ok, bad):
    global _written, _failed
    with _lock:
        _written += ok
        _failed += bad


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--count", type=int, default=400_000, help="补写条数")
    parser.add_argument("--workers", type=int, default=16, help="并发线程数")
    args = parser.parse_args()

    total = args.count
    workers = args.workers
    client = _make_client()

    # 平均分配到各线程
    per = total // workers
    remainder = total % workers
    chunks = [per + (1 if i < remainder else 0) for i in range(workers)]

    print(f"目标表    : {TABLE} ({REGION}, profile={PROFILE})")
    print(f"补写条数  : {total:,}  并发线程: {workers}  每批: {BATCH_SIZE}")
    print(f"线程分配  : {chunks}")
    print("-" * 60)

    start = time.time()

    # 后台进度打印
    stop_flag = {"stop": False}

    def _progress():
        while not stop_flag["stop"]:
            time.sleep(5)
            with _lock:
                w, f = _written, _failed
            elapsed = time.time() - start
            rate = w / elapsed if elapsed > 0 else 0
            pct = w / total * 100 if total else 100
            eta = (total - w) / rate if rate > 0 else 0
            print(f"  进度 {w:>9,}/{total:,} ({pct:5.1f}%)  "
                  f"失败 {f}  速率 {rate:8.0f} 条/s  ETA {eta:6.0f}s")

    monitor = threading.Thread(target=_progress, daemon=True)
    monitor.start()

    report_every = max(1000, total // workers // 20)
    with ThreadPoolExecutor(max_workers=workers) as ex:
        futures = [ex.submit(_worker, client, c, report_every) for c in chunks]
        for fut in as_completed(futures):
            fut.result()  # 抛出未捕获异常（理论上 _worker 内已吞批级错误）

    stop_flag["stop"] = True
    elapsed = time.time() - start
    print("-" * 60)
    print(f"完成：成功 {_written:,} 条，失败 {_failed} 条，耗时 {elapsed:.1f}s，"
          f"平均 {_written/elapsed if elapsed else 0:.0f} 条/s")


if __name__ == "__main__":
    main()

PROFILE = "default"; REGION = "cn-north-1"; TABLE = "allinone_qa_LargeDataTable"
BATCH_SIZE = 25                    # BatchWriteItem 单批硬上限就是 25
RAND_MIN, RAND_MAX = 1800, 2200    # random_string 长度区间，贴近现有数据的 ~2048

def _write_batch(client, items):
    request = {TABLE: [{"PutRequest": {"Item": it}} for it in items]}
    attempt = 0
    while request.get(TABLE):
        resp = client.batch_write_item(RequestItems=request)
        unprocessed = resp.get("UnprocessedItems", {})
        if not unprocessed.get(TABLE):
            break
        request = unprocessed
        attempt += 1
        time.sleep(min(2 ** attempt * 0.05, 5))   # 指数退避
    return len(items)

# 先冒烟 100 条
python load_data.py --count 100 --workers 4      # 100 条 0 失败，132 条/s

# 全量 40w（后台运行，PYTHONIOENCODING=utf-8 避免 Windows GBK 控制台中文乱码）
PYTHONIOENCODING=utf-8 python load_data.py --count 400000 --workers 16

"""
模拟真实业务应用对 DynamoDB 表的持续增删改（INSERT / UPDATE / DELETE）。

用途：DynamoDB 跨账号迁移方案验证。在执行 Export -> Import -> Streams+Lambda
增量同步的全过程中，让源表持续产生变更，以检验：
  - 全量基线导出期间的写入是否被增量通道捕获
  - cutover 前增量能否追平、数据是否最终一致

设计：
  - 维护一个进程内"已知主键池"(known keys)：启动时从表抽样捞一批，运行中新 INSERT 的
    key 也加入池；UPDATE/DELETE 从池中随机取，保证操作命中真实存在的 item。
  - 三类操作按 --insert/--update/--delete 权重随机混合。
  - 通过 --tps 控制目标吞吐（令牌桶节流），--duration 控制运行时长，--workers 控制并发。
  - 每条变更都在 metadata.last_op / metadata.op_ts 留痕，便于迁移后做数据比对核验。

显式锁定 default profile + cn-north-1，避免被环境变量 AWS_PROFILE=global 干扰。

用法：
    python simulate_workload.py --tps 50 --duration 600
    python simulate_workload.py --tps 100 --duration 0   # duration=0 表示一直跑到 Ctrl+C
    python simulate_workload.py --insert 3 --update 6 --delete 1 --tps 80
"""

import argparse
import random
import signal
import string
import threading
import time
import uuid
from collections import deque
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError

# ----------------------- 固定配置 -----------------------
PROFILE = "default"
REGION = "cn-north-1"
TABLE = "allinone_qa_LargeDataTable"

CATEGORIES = ["A", "B", "C", "D", "E"]
VERSIONS = ["v1", "v2", "v3"]
RAND_MIN, RAND_MAX = 1800, 2200
_ALPHABET = string.ascii_letters + string.digits

# 启动时从表抽样捞多少 key 进池（供 UPDATE/DELETE 命中已有数据）
SEED_KEY_TARGET = 5000

_stop = threading.Event()


# ----------------------- 主键池 -----------------------
class KeyPool:
    """线程安全的主键池，存 (id, timestamp) 元组。

    用 deque 限制上限，避免长时间运行内存膨胀；满了丢最旧的。
    """

    def __init__(self, maxlen=200_000):
        self._dq = deque(maxlen=maxlen)
        self._lock = threading.Lock()

    def add(self, key):
        with self._lock:
            self._dq.append(key)

    def random(self):
        with self._lock:
            if not self._dq:
                return None
            return random.choice(self._dq)

    def discard(self, key):
        # deque 无高效随机删除；DELETE 后不强制移除，靠条件写幂等兜底即可。
        # 简化处理：不移除，重复 DELETE 同 key 时第二次条件失败计为 noop。
        pass

    def __len__(self):
        with self._lock:
            return len(self._dq)


# ----------------------- 统计 -----------------------
class Stats:
    def __init__(self):
        self.lock = threading.Lock()
        self.insert = 0
        self.update = 0
        self.delete = 0
        self.noop = 0      # 条件写未命中（如 DELETE 已删的 key）
        self.error = 0

    def bump(self, field, n=1):
        with self.lock:
            setattr(self, field, getattr(self, field) + n)

    def snapshot(self):
        with self.lock:
            return (self.insert, self.update, self.delete, self.noop, self.error)


def _make_client():
    session = boto3.Session(profile_name=PROFILE, region_name=REGION)
    cfg = Config(max_pool_connections=64,
                 retries={"max_attempts": 10, "mode": "adaptive"})
    return session.client("dynamodb", config=cfg)


def _rand_string(n):
    return "".join(random.choices(_ALPHABET, k=n))


def _now_iso():
    return datetime.now(timezone.utc).isoformat()


def _seed_pool(client, pool, target):
    """启动时分页 scan 抽样，捞 target 个 key 进池。只取主键属性，省流量。"""
    scanned = 0
    kwargs = {
        "TableName": TABLE,
        "ProjectionExpression": "id, #ts",
        "ExpressionAttributeNames": {"#ts": "timestamp"},
        "Limit": 1000,
    }
    while len(pool) < target:
        resp = client.scan(**kwargs)
        for it in resp.get("Items", []):
            pool.add((it["id"]["S"], it["timestamp"]["S"]))
        scanned += resp.get("ScannedCount", 0)
        lek = resp.get("LastEvaluatedKey")
        if not lek:
            break
        kwargs["ExclusiveStartKey"] = lek
    print(f"[seed] 已抽样 {len(pool)} 个主键进池（scanned={scanned}）")


# ----------------------- 三类操作 -----------------------
def _do_insert(client, pool, stats):
    key_id = str(uuid.uuid4())
    ts = _now_iso()
    item = {
        "id": {"S": key_id},
        "timestamp": {"S": ts},
        "random_string": {"S": _rand_string(random.randint(RAND_MIN, RAND_MAX))},
        "category": {"S": random.choice(CATEGORIES)},
        "numeric_field": {"N": str(random.randint(1, 1_000_000))},
        "metadata": {"M": {
            "created_by": {"S": "workload_sim"},
            "version": {"S": random.choice(VERSIONS)},
            "last_op": {"S": "INSERT"},
            "op_ts": {"S": ts},
            "tags": {"L": [{"S": f"tag_{i}"} for i in range(random.randint(5, 10))]},
        }},
    }
    client.put_item(TableName=TABLE, Item=item)
    pool.add((key_id, ts))
    stats.bump("insert")


def _do_update(client, pool, stats):
    key = pool.random()
    if key is None:
        stats.bump("noop")
        return
    key_id, ts = key
    now = _now_iso()
    try:
        client.update_item(
            TableName=TABLE,
            Key={"id": {"S": key_id}, "timestamp": {"S": ts}},
            UpdateExpression=("SET numeric_field = :n, category = :c, "
                              "metadata.last_op = :op, metadata.op_ts = :t"),
            ExpressionAttributeValues={
                ":n": {"N": str(random.randint(1, 1_000_000))},
                ":c": {"S": random.choice(CATEGORIES)},
                ":op": {"S": "UPDATE"},
                ":t": {"S": now},
            },
            # 仅当该 item 仍存在时才更新（已被 DELETE 的不重建）
            ConditionExpression="attribute_exists(id)",
        )
        stats.bump("update")
    except ClientError as e:
        if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
            stats.bump("noop")
        else:
            raise


def _do_delete(client, pool, stats):
    key = pool.random()
    if key is None:
        stats.bump("noop")
        return
    key_id, ts = key
    try:
        client.delete_item(
            TableName=TABLE,
            Key={"id": {"S": key_id}, "timestamp": {"S": ts}},
            ConditionExpression="attribute_exists(id)",
        )
        stats.bump("delete")
    except ClientError as e:
        if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
            stats.bump("noop")  # 已被删过，幂等 noop
        else:
            raise


def _pick_op(weights):
    r = random.random() * weights["total"]
    if r < weights["insert"]:
        return _do_insert
    if r < weights["insert"] + weights["update"]:
        return _do_update
    return _do_delete


# ----------------------- 令牌桶节流 -----------------------
class RateLimiter:
    """简单令牌桶：每秒补充 tps 个令牌。tps<=0 表示不限速。"""

    def __init__(self, tps):
        self.tps = tps
        self.lock = threading.Lock()
        self.allowance = float(tps)
        self.last = time.monotonic()

    def acquire(self):
        if self.tps <= 0:
            return
        while True:
            with self.lock:
                now = time.monotonic()
                self.allowance += (now - self.last) * self.tps
                self.last = now
                if self.allowance > self.tps:
                    self.allowance = self.tps
                if self.allowance >= 1:
                    self.allowance -= 1
                    return
            time.sleep(0.002)


def _worker(client, pool, stats, weights, limiter):
    while not _stop.is_set():
        limiter.acquire()
        if _stop.is_set():
            break
        op = _pick_op(weights)
        try:
            op(client, pool, stats)
        except Exception as e:  # noqa: BLE001 - 单次操作失败不中断模拟
            stats.bump("error")
            print(f"[ERR] {op.__name__}: {e}")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--tps", type=int, default=50, help="目标总吞吐(操作/秒)，<=0 不限速")
    p.add_argument("--duration", type=int, default=300, help="运行秒数，0=直到 Ctrl+C")
    p.add_argument("--workers", type=int, default=8, help="并发线程数")
    p.add_argument("--insert", type=int, default=3, help="INSERT 权重")
    p.add_argument("--update", type=int, default=5, help="UPDATE 权重")
    p.add_argument("--delete", type=int, default=2, help="DELETE 权重")
    p.add_argument("--seed-keys", type=int, default=SEED_KEY_TARGET,
                   help="启动时抽样进池的 key 数")
    args = p.parse_args()

    weights = {
        "insert": args.insert,
        "update": args.update,
        "delete": args.delete,
        "total": args.insert + args.update + args.delete,
    }

    client = _make_client()
    pool = KeyPool()

    print(f"目标表    : {TABLE} ({REGION}, profile={PROFILE})")
    print(f"操作权重  : INSERT={args.insert} UPDATE={args.update} DELETE={args.delete}")
    print(f"目标 TPS  : {args.tps}  时长: {args.duration or '∞'}s  线程: {args.workers}")
    print("-" * 64)

    _seed_pool(client, pool, args.seed_keys)
    if len(pool) == 0:
        print("[警告] 主键池为空，UPDATE/DELETE 将全部 noop，仅 INSERT 生效。")

    # Ctrl+C 优雅停止
    signal.signal(signal.SIGINT, lambda *_: _stop.set())

    limiter = RateLimiter(args.tps)
    stats = Stats()
    start = time.time()

    def _progress():
        last = stats.snapshot()
        last_t = start
        while not _stop.is_set():
            time.sleep(5)
            now = time.time()
            cur = stats.snapshot()
            di = sum(cur[:3]) - sum(last[:3])
            cur_tps = di / (now - last_t) if now > last_t else 0
            i, u, d, n, e = cur
            print(f"  t+{now-start:5.0f}s | INS {i:>7,} UPD {u:>7,} DEL {d:>7,} "
                  f"| noop {n:>6,} err {e} | 池 {len(pool):>6,} | ~{cur_tps:5.0f} ops/s")
            last, last_t = cur, now

    threading.Thread(target=_progress, daemon=True).start()

    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        for _ in range(args.workers):
            ex.submit(_worker, client, pool, stats, weights, limiter)
        # 主线程负责计时
        try:
            while not _stop.is_set():
                if args.duration and (time.time() - start) >= args.duration:
                    _stop.set()
                    break
                time.sleep(0.2)
        except KeyboardInterrupt:
            _stop.set()

    i, u, d, n, e = stats.snapshot()
    elapsed = time.time() - start
    total_ops = i + u + d
    print("-" * 64)
    print(f"运行结束：耗时 {elapsed:.1f}s")
    print(f"  INSERT={i:,}  UPDATE={u:,}  DELETE={d:,}  noop={n:,}  error={e}")
    print(f"  有效变更 {total_ops:,} 次，平均 {total_ops/elapsed if elapsed else 0:.1f} ops/s")
    print(f"  净增条数 ≈ INSERT - DELETE = {i - d:,}")


if __name__ == "__main__":
    main()

self.allowance += (now - self.last) * self.tps   # 按时间补充令牌
if self.allowance >= 1: self.allowance -= 1; return

PYTHONIOENCODING=utf-8 python simulate_workload.py --tps 40 --duration 60 --workers 8 --seed-keys 3000

[seed] 已抽样 3289 个主键进池
INSERT=735  UPDATE=1,150  DELETE=447  noop=97  error=0
有效变更 2,332 次，平均 38.7 ops/s，净增 288 条

tmp/ddb/backup/AWSDynamoDB/01782130127268-c860c940/
    ├── manifest-summary.json
    ├── manifest-files.json
    └── data/                          # 24 个 .json.gz 文件，每个约 100 MB
        ├── 3goqtddyim5ynm75oghnyy6s7u.json.gz   (101,879,132 bytes)
        ├── 4aouzthtju5wpcsb2luvcnsgk4.json.gz   (101,466,287 bytes)
        └── ... (共 24 个)

{"Version":"2012-10-17","Statement":[{"Effect":"Allow",
  "Principal":{"Service":"lambda.amazonaws.com"},"Action":"sts:AssumeRole"}]}

{
  "Version":"2012-10-17",
  "Statement":[
    {"Sid":"Logs","Effect":"Allow",
     "Action":["logs:CreateLogGroup","logs:CreateLogStream","logs:PutLogEvents"],
     "Resource":"arn:aws-cn:logs:*:*:*"},
    {"Sid":"ReadSourceStream","Effect":"Allow",
     "Action":["dynamodb:GetRecords","dynamodb:GetShardIterator",
               "dynamodb:DescribeStream","dynamodb:ListStreams"],
     "Resource":"arn:aws-cn:dynamodb:cn-north-1:source_account_id:table/allinone_qa_LargeDataTable/stream/*"},
    {"Sid":"WriteTargetTableCrossAccount","Effect":"Allow",
     "Action":["dynamodb:PutItem","dynamodb:UpdateItem","dynamodb:DeleteItem","dynamodb:BatchWriteItem"],
     "Resource":"arn:aws-cn:dynamodb:cn-north-1:target_account_id:table/allinone_qa_LargeDataTable"}
  ]
}

TARGET_TABLE_ARN = arn:aws-cn:dynamodb:cn-north-1:target_account_id:table/allinone_qa_LargeDataTable
TARGET_REGION    = cn-north-1
EXPORT_TS        = 1782130007
REPL_TS_ATTR     = _replication_ts

lam.create_event_source_mapping(
    EventSourceArn=SRC_STREAM_ARN,
    FunctionName="ddb-xacct-replicator",
    Enabled=False,                          # 先禁用！
    StartingPosition="TRIM_HORIZON",
    BatchSize=100,
    MaximumRetryAttempts=-1,
    BisectBatchOnFunctionError=True,
    FunctionResponseTypes=["ReportBatchItemFailures"],
)

[role] created ddb-xacct-replicator-role
[role] inline policy attached
[fn]   created ddb-xacct-replicator
[esm]  created 5706f37f-bcb7-43c0-85c9-6696cb5fbee7 (Enabled=False)

{
  "Version":"2012-10-17",
  "Statement":[{
    "Sid":"AllowCrossAccountReplicatorWrites",
    "Effect":"Allow",
    "Principal":{"AWS":"arn:aws-cn:iam::source_account_id:role/ddb-xacct-replicator-role"},
    "Action":["dynamodb:PutItem","dynamodb:UpdateItem","dynamodb:DeleteItem","dynamodb:BatchWriteItem"],
    "Resource":"arn:aws-cn:dynamodb:cn-north-1:target_account_id:table/allinone_qa_LargeDataTable"
  }]
}

lam.update_event_source_mapping(UUID="5706f37f-...", Enabled=True)
# State 秒切 Enabled（小表），LastProcessingResult: No records processed → 随即 OK

时刻      IteratorAge(Max)     说明
20:55      3037.9 s          ┐ 阶段2：积压下降中（模拟客户端先写了~50分钟才enable，
20:56      2769.1 s          │         TRIM_HORIZON 从最老记录读起，故起跳≈积压时长）
20:57      2406.2 s          │
20:58      1706.8 s          │
20:59       821.0 s          ┘
21:00         1.5 s          ← 阶段3：贴零，管道排空，目标表实时跟随源表

def seg(i):
    cnt=0; lek=None
    while True:
        kw=dict(TableName=TABLE, Select='COUNT', Segment=i, TotalSegments=8)
        if lek: kw['ExclusiveStartKey']=lek
        r=ddb.scan(**kw); cnt+=r['Count']; lek=r.get('LastEvaluatedKey')
        if not lek: break
    return cnt
# 源账号与目标账号各跑一次