# /// script
# requires-python = ">=3.10"
# dependencies = ["boto3"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Find the small-file crossover on object storage.

1,000 objects of 10 KB and one object of 10 MB in an S3-compatible store.
Read all the small ones with one thread and with 32, and the large one once:
concurrency changes the time, never the request count, and the fee follows
the request count.
Needs the S3 service from the compose file (SeaweedFS, port 59000), or set --endpoint.
"""
import argparse, json, os, time
from concurrent.futures import ThreadPoolExecutor
import boto3
from botocore.config import Config

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--endpoint", default="http://localhost:59000", help="S3 endpoint (empty for AWS).")
p.add_argument("--bucket", default="lab", help="Bucket.")
p.add_argument("--small", type=int, default=1000, help="Number of small objects.")
p.add_argument("--small-kb", type=int, default=10, help="Size of each small object, KB.")
p.add_argument("--threads", type=lambda s: [int(v) for v in s.split(",")], default=[1, 8, 32], help="Thread counts to try.")
p.add_argument("--get-price", type=float, default=0.0004, help="Price per 1,000 GET requests (S3 Standard list price).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

s3 = boto3.client("s3", endpoint_url=args.endpoint or None, aws_access_key_id="lab", aws_secret_access_key="lablablab",
                  region_name="us-east-1", config=Config(max_pool_connections=64))
try: s3.create_bucket(Bucket=args.bucket)
except s3.exceptions.BucketAlreadyOwnedByYou: pass
blob = os.urandom(args.small_kb * 1024)
keys = [f"small/{i}.bin" for i in range(args.small)]
for k in keys: s3.put_object(Bucket=args.bucket, Key=k, Body=blob)
s3.put_object(Bucket=args.bucket, Key="large.bin", Body=os.urandom(args.small * args.small_kb * 1024))

def get(k): return len(s3.get_object(Bucket=args.bucket, Key=k)["Body"].read())

rows = []
for n in args.threads:
    t = time.perf_counter()
    with ThreadPoolExecutor(n) as ex: total = sum(ex.map(get, keys))
    dt = time.perf_counter() - t
    rows.append({"read": f"{args.small:,} small objects, {n} thread{'s' if n > 1 else ''}", "threads": n, "seconds": dt,
                 "requests": args.small, "fee": args.small / 1000 * args.get_price})
t = time.perf_counter(); get("large.bin"); dt = time.perf_counter() - t
rows.append({"read": "one large object", "threads": 1, "seconds": dt, "requests": 1, "fee": 1 / 1000 * args.get_price})
for r in rows:
    print(f"{r['read']:<34} {r['seconds']:7.3f} s   {r['requests']:>5,} GETs   fee {r['fee']:.6f}")
small1 = rows[0]["seconds"]; small32 = next(r["seconds"] for r in rows if r["threads"] == max(args.threads))
print(f"\n{max(args.threads)} threads: {small1 / small32:.1f}x faster than 1 thread, same {args.small:,} requests; "
      f"the large object costs 1/{args.small:,} of the fee for the same bytes")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"reads": rows, "summary": [{"speedup": small1 / small32, "requests": args.small,
                   "fee_ratio": args.small, "large_s": rows[-1]["seconds"]}]}}, f)