# /// script
# requires-python = ">=3.10"
# dependencies = ["urllib3"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Stale pooled connections: find out what your pool does.

With urllib3's pool: (1) pool size 2 and 10 slow requests (the server sleeps
5 s): what happens to the 3rd? (2) A connection through a TCP proxy that is
killed and restarted: the first write on the pooled connection after its
silent death. (3) The same with a retry on that first failure. Prints the
library's defaults for exhaustion, max idle time and max lifetime.
"""
import argparse, json, socket, threading, time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import urllib3

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--slow-s", type=float, default=5.0, help="Seconds the slow endpoint sleeps.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

class H(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    def do_GET(self):
        if self.path == "/slow": time.sleep(args.slow_s)
        self.send_response(200); self.send_header("Content-Length", "2"); self.end_headers(); self.wfile.write(b"ok")
    def log_message(self, *a): pass
srv = ThreadingHTTPServer(("127.0.0.1", 0), H); threading.Thread(target=srv.serve_forever, daemon=True).start()
backend = srv.server_address[1]

rows = []
# (1) exhaustion: block=True waits for a free connection; block=False opens extra ones and discards them
for block in (False, True):
    pool = urllib3.HTTPConnectionPool("127.0.0.1", backend, maxsize=2, block=block, timeout=urllib3.Timeout(read=30))
    started = {}
    def call(i):
        t = time.perf_counter(); pool.request("GET", "/slow"); started[i] = time.perf_counter() - t
    ts = [threading.Thread(target=call, args=(i,)) for i in range(10)]
    t0 = time.perf_counter(); [t.start() for t in ts]; [t.join() for t in ts]
    total = time.perf_counter() - t0
    third = sorted(started.values())[2]
    rows.append({"experiment": f"10 slow requests, pool of 2, block={block}",
                 "result": f"3rd request finished after {third:.1f} s; all done after {total:.1f} s"})
    print(rows[-1]["experiment"], "->", rows[-1]["result"])

# (2) and (3): a proxy that dies silently between two requests on a pooled connection
class Proxy:
    """A TCP proxy. blackhole() is a silent death: it stops forwarding and stops listening,
    but sends no FIN or RST, as a NAT timeout or a crashed middlebox does."""
    def __init__(self, port): self.port, self.sock, self.conns, self.dead = port, None, [], False
    def start(self, listen_port=0):
        self.sock = socket.socket(); self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(("127.0.0.1", listen_port)); self.sock.listen(8)
        threading.Thread(target=self.loop, daemon=True).start(); return self.sock.getsockname()[1]
    def loop(self):
        while True:
            try: c, _ = self.sock.accept()
            except OSError: return
            u = socket.create_connection(("127.0.0.1", self.port)); self.conns += [c, u]
            for a, b in ((c, u), (u, c)): threading.Thread(target=self.pipe, args=(a, b), daemon=True).start()
    def pipe(self, a, b):
        try:
            while (d := a.recv(65536)):
                if self.dead: continue                                   # swallow everything, reply nothing
                b.sendall(d)
        except OSError: pass
    def blackhole(self):
        self.dead = True; self.sock.close()

for retries, label in ((0, "first request after the proxy died silently and restarted, no retry"), (1, "same, with one retry")):
    proxy = Proxy(backend); port = proxy.start()
    pool = urllib3.HTTPConnectionPool("127.0.0.1", port, maxsize=1, retries=retries, timeout=urllib3.Timeout(connect=2, read=2))
    pool.request("GET", "/")                                          # the connection goes back to the pool
    proxy.blackhole(); time.sleep(0.2); fresh = Proxy(backend); fresh.start(port)   # a new proxy on the same port
    t = time.perf_counter()
    try:
        r = pool.request("GET", "/"); outcome = f"status {r.status} after {time.perf_counter() - t:.1f} s"
    except Exception as e:
        outcome = f"failed after {time.perf_counter() - t:.1f} s: {type(e).__name__}"
    rows.append({"experiment": label, "result": outcome}); print(label, "->", outcome)
    fresh.blackhole()

defaults = urllib3.HTTPConnectionPool("x")
d = {"exhaustion (block)": f"block={defaults.block}: extra connections are opened and discarded, no waiting",
     "max idle time": "none (no idle eviction)", "max lifetime": "none",
     "retries on a stale connection": f"{urllib3.Retry.DEFAULT if hasattr(urllib3.Retry, 'DEFAULT') else 'Retry(3)'}"}
for k, v in d.items(): print(f"default {k}: {v}")
rows += [{"experiment": f"default: {k}", "result": v} for k, v in d.items()]

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