Sizing, Bounding, Backpressure and Rate Limits

import math


def workers_for(rate_per_s: float, p95_latency_s: float, headroom: float = 1.2, cap: int = 64) -> int:
    return max(1, min(cap, math.ceil(rate_per_s * p95_latency_s * headroom)))


print(workers_for(500 / 60, 4.0))      # 40
print(workers_for(3000 / 60, 0.4))     # 24
print(workers_for(200, 0.05))          # 12

import tracemalloc
from concurrent.futures import ThreadPoolExecutor


def lines(n: int):
    for i in range(n):
        yield "x" * 200 + str(i)                     # stands in for a JSONL record


tracemalloc.start()
with ThreadPoolExecutor(max_workers=4) as ex:
    it = ex.map(len, lines(50_000))                  # noqa: TPE004 - all 50,000 submitted right here
    peak_after_map = tracemalloc.get_traced_memory()[1]
    total = sum(it)
tracemalloc.stop()
print(total, peak_after_map > 10_000_000)            # 10238890 True  (>10 MB for 50k small items)

import itertools
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from typing import Callable, Iterable, Iterator, Tuple, TypeVar

T = TypeVar("T")
R = TypeVar("R")


def window_map(fn: Callable[[T], R], items: Iterable[T], ex: ThreadPoolExecutor,
               max_in_flight: int = 64) -> Iterator[Tuple[int, R]]:
    """Yield (index, result) as tasks finish; never more than max_in_flight pending."""
    source = enumerate(items)
    pending = {ex.submit(fn, item): i for i, item in itertools.islice(source, max_in_flight)}
    try:
        while pending:
            done, _ = wait(pending, return_when=FIRST_COMPLETED)
            for fut in done:
                yield pending.pop(fut), fut.result()
            for i, item in itertools.islice(source, len(done)):
                pending[ex.submit(fn, item)] = i
    finally:
        for fut in pending:
            fut.cancel()


with ThreadPoolExecutor(max_workers=8) as ex:
    out = dict(window_map(lambda s: s.upper(), (f"doc{i}" for i in range(1000)), ex, max_in_flight=16))
print(len(out), out[999])                            # 1000 DOC999

import collections
import itertools
import time
from concurrent.futures import ThreadPoolExecutor


def ordered_window_map(fn, items, ex, max_in_flight=32):
    source = iter(items)
    window = collections.deque(ex.submit(fn, x) for x in itertools.islice(source, max_in_flight))
    try:
        while window:
            result = window.popleft().result()      # oldest first: preserves order
            for x in itertools.islice(source, 1):
                window.append(ex.submit(fn, x))
            yield result
    finally:
        for fut in window:
            fut.cancel()


def slow_even(i: int) -> int:
    time.sleep(0.01 if i % 2 == 0 else 0.001)
    return i * i


with ThreadPoolExecutor(max_workers=4) as ex:
    print(list(ordered_window_map(slow_even, range(10), ex, max_in_flight=4)))
# [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

import threading
from concurrent.futures import Future, ThreadPoolExecutor
from typing import Any, Callable, Optional


class BoundedExecutor:
    """submit() blocks while `bound` tasks are pending (queued + running)."""

    def __init__(self, bound: int, max_workers: int, **kwargs: Any) -> None:
        self._slots = threading.BoundedSemaphore(bound)
        self._ex = ThreadPoolExecutor(max_workers=max_workers, **kwargs)

    def submit(self, fn: Callable[..., Any], *args: Any, timeout: Optional[float] = None, **kwargs: Any) -> Future:
        if not self._slots.acquire(timeout=timeout):
            raise TimeoutError("no free slot")
        try:
            fut = self._ex.submit(fn, *args, **kwargs)
        except BaseException:
            self._slots.release()
            raise
        fut.add_done_callback(lambda _f: self._slots.release())
        return fut

    def shutdown(self, wait: bool = True, cancel_futures: bool = False) -> None:
        self._ex.shutdown(wait=wait, cancel_futures=cancel_futures)

    def __enter__(self) -> "BoundedExecutor":
        return self

    def __exit__(self, *exc: Any) -> None:
        self.shutdown()


peak = 0
lock = threading.Lock()
active = 0


def task(i: int) -> int:
    global peak, active
    with lock:
        active += 1
        peak = max(peak, active)
    with lock:
        active -= 1
    return i


with BoundedExecutor(bound=8, max_workers=4) as bx:
    futures = [bx.submit(task, i) for i in range(200)]
print(sum(f.result() for f in futures), peak <= 4)   # 19900 True

import queue
import threading

STOP = object()


def stage(inbox: queue.Queue, outbox: queue.Queue, fn) -> None:
    while True:
        item = inbox.get()
        try:
            if item is STOP:
                return
            outbox.put(fn(item))
        finally:
            inbox.task_done()


urls: queue.Queue = queue.Queue(maxsize=32)          # producer blocks when parsers fall behind
pages: queue.Queue = queue.Queue(maxsize=32)
workers = [threading.Thread(target=stage, args=(urls, pages, str.upper), name=f"parse-{i}") for i in range(4)]
for t in workers:
    t.start()


def produce() -> None:
    for i in range(100):
        urls.put(f"https://site/{i}")
    for _ in workers:
        urls.put(STOP)                               # one sentinel per worker


producer = threading.Thread(target=produce, name="producer")
producer.start()
results = [pages.get() for _ in range(100)]          # consumer drains so parsers never block forever
producer.join()
for t in workers:
    t.join()
print(len(results), results[0].startswith("HTTPS://"))   # 100 True

import sys
from concurrent.futures import ThreadPoolExecutor


def records():
    for i in range(100_000):
        yield i


with ThreadPoolExecutor(max_workers=8) as ex:
    if sys.version_info >= (3, 14):
        total = sum(ex.map(abs, records(), buffersize=64))   # at most 64 submitted ahead
    else:
        total = sum(abs(r) for r in records())               # use a window on older versions
print(total)                                                 # 4999950000

import threading
import time
from concurrent.futures import ThreadPoolExecutor


class TokenBucket:
    def __init__(self, rate_per_minute: float, burst: int,
                 clock=time.monotonic, sleep=time.sleep) -> None:
        self.rate = rate_per_minute / 60.0
        self.capacity = float(burst)
        self.tokens = float(burst)
        self.clock, self.sleep = clock, sleep
        self.updated = clock()
        self.lock = threading.Lock()

    def acquire(self, n: float = 1.0) -> None:
        while True:
            with self.lock:
                now = self.clock()
                self.tokens = min(self.capacity, self.tokens + (now - self.updated) * self.rate)
                self.updated = now
                if self.tokens >= n:
                    self.tokens -= n
                    return
                wait_s = (n - self.tokens) / self.rate
            self.sleep(wait_s)


bucket = TokenBucket(rate_per_minute=600, burst=3)    # 10/s
stamps = []


def call(i: int) -> None:
    bucket.acquire()
    stamps.append(time.monotonic())


t0 = time.monotonic()
with ThreadPoolExecutor(max_workers=8) as ex:
    list(ex.map(call, range(13)))
print(f"{time.monotonic() - t0:.1f}s")               # 1.0s  (3 burst + 10 more at 10/s)

import threading
import time


class Limits:
    """Acquire one request token and `tokens` TPM tokens atomically enough for a client-side limiter."""

    def __init__(self, rpm: float, tpm: float) -> None:
        self.rpm_rate, self.tpm_rate = rpm / 60, tpm / 60
        self.req, self.tok = rpm / 60, tpm / 60          # start with one second of budget
        self.cap_req, self.cap_tok = rpm / 60, tpm / 60
        self.t = time.monotonic()
        self.lock = threading.Lock()

    def acquire(self, tokens: int) -> float:
        if tokens > self.cap_tok:
            raise ValueError("single request exceeds one second of TPM budget; raise the cap")
        waited = 0.0
        while True:
            with self.lock:
                now = time.monotonic()
                dt, self.t = now - self.t, now
                self.req = min(self.cap_req, self.req + dt * self.rpm_rate)
                self.tok = min(self.cap_tok, self.tok + dt * self.tpm_rate)
                if self.req >= 1 and self.tok >= tokens:
                    self.req -= 1
                    self.tok -= tokens
                    return waited
                need = max((1 - self.req) / self.rpm_rate, (tokens - self.tok) / self.tpm_rate)
            time.sleep(need)
            waited += need


limits = Limits(rpm=6000, tpm=60_000)                    # 100 req/s, 1000 tok/s
waits = [limits.acquire(tokens=400) for _ in range(5)]
print(round(sum(waits), 1))                              # 1.0  (TPM is the binding limit here)

import threading
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from urllib.parse import urlsplit

_host_slots = defaultdict(lambda: threading.BoundedSemaphore(2))
_host_lock = threading.Lock()
_active = defaultdict(int)
_peak = defaultdict(int)


def slots_for(url: str) -> threading.BoundedSemaphore:
    with _host_lock:                                    # defaultdict insert is check-then-set
        return _host_slots[urlsplit(url).netloc]


def fetch(url: str) -> str:
    host = urlsplit(url).netloc
    with slots_for(url):                                # at most 2 concurrent requests per host
        with _host_lock:
            _active[host] += 1
            _peak[host] = max(_peak[host], _active[host])
        try:
            return url                                  # requests.get(url, timeout=...) here
        finally:
            with _host_lock:
                _active[host] -= 1


urls = [f"https://site{i % 3}.example/p{i}" for i in range(60)]
with ThreadPoolExecutor(max_workers=12) as ex:
    list(ex.map(fetch, urls))
print(max(_peak.values()) <= 2)                         # True