# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Reader-writer lock against a plain mutex (and a snapshot pointer).

A shared table of 1,000 counters; N threads doing 99% reads and 1% writes
for a fixed time, protected by a mutex, by a pthread reader-writer lock, and
by an immutable snapshot read through one atomic pointer (writers copy and
swap). Total operations per second against N. The C program is compiled
and run by this script (needs cc).
"""
import argparse, json, os, subprocess, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--threads", type=lambda s: [int(v) for v in s.split(",")], default=[1, 2, 4, 8, 16], help="Thread counts N.")
p.add_argument("--seconds", type=float, default=1.0, help="Seconds per measurement.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

C = r"""
#include <pthread.h>
#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#define E 1000
#define RING 64
long table[E];
pthread_mutex_t mu = PTHREAD_MUTEX_INITIALIZER;
pthread_rwlock_t rw = PTHREAD_RWLOCK_INITIALIZER;
long *snaps[RING]; _Atomic(long *) current; _Atomic long next_snap = 1;
pthread_mutex_t wmu = PTHREAD_MUTEX_INITIALIZER;
volatile int stop = 0;
static double now(void) { struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec / 1e9; }
static void *worker(void *arg) {
    unsigned s = (unsigned)(long)arg * 2654435761u; long ops = 0, acc = 0;
    while (!stop) {
        s = s * 1103515245 + 12345; int i = (s >> 8) % E; int write = (s >> 4) % 100 == 0;
#if MODE == 0
        pthread_mutex_lock(&mu); if (write) table[i]++; else acc += table[i]; pthread_mutex_unlock(&mu);
#elif MODE == 1
        if (write) { pthread_rwlock_wrlock(&rw); table[i]++; } else { pthread_rwlock_rdlock(&rw); acc += table[i]; }
        pthread_rwlock_unlock(&rw);
#else
        if (write) {                            /* copy, modify, publish (old copies are reused from a ring) */
            pthread_mutex_lock(&wmu);
            long *cur = atomic_load(&current); long *nx = snaps[atomic_fetch_add(&next_snap, 1) % RING];
            memcpy(nx, cur, sizeof(long) * E); nx[i]++; atomic_store(&current, nx);
            pthread_mutex_unlock(&wmu);
        } else { acc += atomic_load_explicit(&current, memory_order_acquire)[i]; }
#endif
        ops++;
    }
    return (void *)(ops + (acc & 0));
}
int main(int argc, char **argv) {
    int n = atoi(argv[1]); double secs = atof(argv[2]);
    for (int r = 0; r < RING; r++) snaps[r] = calloc(E, sizeof(long));
    atomic_store(&current, snaps[0]);
    pthread_t th[64]; for (long i = 0; i < n; i++) pthread_create(&th[i], NULL, worker, (void *)(i + 1));
    struct timespec ts = { (time_t)secs, (long)((secs - (time_t)secs) * 1e9) };
    nanosleep(&ts, NULL); stop = 1;                      /* the main thread sleeps: it must not take a core */
    long total = 0; for (int i = 0; i < n; i++) { void *r; pthread_join(th[i], &r); total += (long)r; }
    printf("%.0f\n", total / secs);
}
"""
d = tempfile.mkdtemp(); src = os.path.join(d, "rw.c"); open(src, "w").write(C)
rows = []
for mode, label in ((0, "mutex"), (1, "reader-writer lock"), (2, "snapshot pointer")):
    exe = os.path.join(d, f"rw{mode}")
    subprocess.run(["cc", "-O2", "-pthread", src, "-o", exe, f"-DMODE={mode}"], check=True)
    for n in args.threads:
        ops = float(subprocess.run([exe, str(n), str(args.seconds)], capture_output=True, text=True, check=True).stdout)
        rows.append({"lock": label, "threads": n, "mops": ops / 1e6})
    print(f"{label:<20} " + "  ".join(f"N={r['threads']}: {r['mops']:6.1f}" for r in rows if r["lock"] == label) + "  million ops/s")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"throughput": rows}}, f)