"""
02_select_for_update.py
-----------------------
演示两种"任务队列"消费模式：
  方案 A:  普通 FOR UPDATE       —— 多 worker 串行抢同一行
  方案 B:  FOR UPDATE SKIP LOCKED —— 多 worker 并行各自拿不同行

通过对比两种方案处理 100 个任务的耗时，直观感受 SKIP LOCKED 的威力。

依赖:  pip install psycopg[binary]>=3.1
运行:  python 02_select_for_update.py [--workers 5] [--mode skip|wait]
"""
import argparse
import time
import threading
import psycopg

DSN = "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres"


SQL_FETCH_SKIP = """
WITH job AS (
    SELECT id
      FROM ch11_task_queue
     WHERE status = 'pending'
     ORDER BY priority DESC, id
     FOR UPDATE SKIP LOCKED
     LIMIT 1
)
UPDATE ch11_task_queue t
   SET status = 'running',
       started_at = now(),
       locked_by = %s
  FROM job
 WHERE t.id = job.id
RETURNING t.id, t.payload;
"""

SQL_FETCH_WAIT = """
WITH job AS (
    SELECT id
      FROM ch11_task_queue
     WHERE status = 'pending'
     ORDER BY priority DESC, id
     FOR UPDATE
     LIMIT 1
)
UPDATE ch11_task_queue t
   SET status = 'running',
       started_at = now(),
       locked_by = %s
  FROM job
 WHERE t.id = job.id
RETURNING t.id, t.payload;
"""

SQL_DONE = """
UPDATE ch11_task_queue
   SET status = 'done', finished_at = now()
 WHERE id = %s
"""


def reset_queue():
    with psycopg.connect(DSN, autocommit=True) as conn:
        with conn.cursor() as cur:
            cur.execute(
                "UPDATE ch11_task_queue SET status='pending', "
                "started_at=NULL, finished_at=NULL, locked_by=NULL"
            )


def worker(name: str, sql: str, processed: list, lock: threading.Lock):
    with psycopg.connect(DSN) as conn:
        while True:
            with conn.transaction():
                with conn.cursor() as cur:
                    cur.execute(sql, (name,))
                    row = cur.fetchone()
                    if not row:
                        return
                    job_id = row[0]
            time.sleep(0.01)  # 模拟业务耗时
            with psycopg.connect(DSN, autocommit=True) as c2:
                with c2.cursor() as cur2:
                    cur2.execute(SQL_DONE, (job_id,))
            with lock:
                processed.append((name, job_id))


def run(workers: int, mode: str):
    sql = SQL_FETCH_SKIP if mode == "skip" else SQL_FETCH_WAIT
    reset_queue()
    processed: list = []
    lock = threading.Lock()

    threads = [
        threading.Thread(target=worker, args=(f"w{i}", sql, processed, lock))
        for i in range(workers)
    ]
    t0 = time.time()
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    dt = time.time() - t0
    by_w: dict = {}
    for w, _ in processed:
        by_w[w] = by_w.get(w, 0) + 1
    print(f"\n=== mode={mode}  workers={workers}  耗时={dt:.2f}s ===")
    print(f"总共消费 {len(processed)} 条，按 worker 分布：{by_w}")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--workers", type=int, default=5)
    p.add_argument("--mode", choices=["skip", "wait", "both"], default="both")
    args = p.parse_args()

    if args.mode in ("wait", "both"):
        run(args.workers, "wait")
    if args.mode in ("skip", "both"):
        run(args.workers, "skip")


if __name__ == "__main__":
    main()