限流是共享预算

redis_lab=$(mktemp -d)
curl --fail --location --output "$redis_lab/redis.tar.gz" https://download.redis.io/releases/redis-8.10.1.tar.gz
tar -xzf "$redis_lab/redis.tar.gz" -C "$redis_lab"
make -C "$redis_lab/redis-8.10.1" build redis -j2 MALLOC=libc
export REDIS_SERVER="$redis_lab/redis-8.10.1/src/redis-server"
python3 -m venv "$redis_lab/venv"
"$redis_lab/venv/bin/python" -m pip install redis==8.1.0

from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
import os
from pathlib import Path
import subprocess
from tempfile import TemporaryDirectory
from threading import Barrier
import time

import redis
from redis.backoff import NoBackoff
from redis.retry import Retry


def connection(endpoint):
    return redis.Redis(unix_socket_path=endpoint, decode_responses=False,
                       socket_timeout=2, retry=Retry(NoBackoff(), 0))


@contextmanager
def instance():
    with TemporaryDirectory(prefix="redis-budget-", dir="/tmp") as directory:
        endpoint = str(Path(directory) / "redis.sock")
        process = subprocess.Popen([
            os.environ["REDIS_SERVER"], "--port", "0", "--unixsocket", endpoint,
            "--unixsocketperm", "700", "--dir", directory,
            "--save", "", "--appendonly", "no",
        ], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
        client = connection(endpoint)
        try:
            deadline = time.monotonic() + 5
            while True:
                assert process.poll() is None, "Redis exited during startup"
                try:
                    if client.ping():
                        break
                except redis.ConnectionError:
                    assert time.monotonic() < deadline, "Redis did not start"
                    time.sleep(0.02)
            assert client.info("server")["process_id"] == process.pid
            yield client, endpoint
        finally:
            client.close()
            if process.poll() is None:
                process.terminate()
            try:
                _, errors = process.communicate(timeout=5)
            except subprocess.TimeoutExpired:
                process.kill()
                process.communicate()
                raise
            assert process.returncode == 0, errors.decode()


acquire = """
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
if current >= tonumber(ARGV[1]) then
    return {0, current, redis.call('PTTL', KEYS[1])}
end
if current == 0 then
    redis.call('SET', KEYS[1], '1', 'PX', ARGV[2])
else
    redis.call('INCR', KEYS[1])
end
return {1, current + 1, redis.call('PTTL', KEYS[1])}
"""

LIMIT, WINDOW_MS, WORKERS = 5, 60000, 16

with instance() as (control, endpoint):
    def concurrent(operation):
        barrier = Barrier(WORKERS)

        def run(_):
            with connection(endpoint) as client:
                return operation(client, barrier)

        with ThreadPoolExecutor(max_workers=WORKERS) as pool:
            return list(pool.map(run, range(WORKERS)))

    control.set("budget:unsafe", 0)

    def split_check(client, barrier):
        observed = int(client.get("budget:unsafe"))
        barrier.wait(timeout=5)  # All callers read before any caller increments.
        if observed >= LIMIT:
            return False
        client.incr("budget:unsafe")
        return True

    unsafe = concurrent(split_check)
    assert sum(unsafe) == WORKERS
    assert int(control.get("budget:unsafe")) == WORKERS
    print("split GET/INCR admits 16 against a budget of 5")

    def atomic_check(client, barrier):
        barrier.wait(timeout=5)
        return client.eval(acquire, 1, "budget:shared", LIMIT, WINDOW_MS)

    results = concurrent(atomic_check)
    assert sum(reply[0] for reply in results) == LIMIT
    assert sorted(reply[1] for reply in results if reply[0]) == list(range(1, LIMIT + 1))
    assert all(reply[1] == LIMIT for reply in results if not reply[0])
    assert all(0 < reply[2] <= WINDOW_MS for reply in results)
    assert control.get("budget:shared") == b"5"
    print("one atomic script admits 5 and rejects 11 across independent client connections")

    ttl_before = control.pttl("budget:shared")
    assert control.eval(acquire, 1, "budget:shared", LIMIT, WINDOW_MS)[:2] == [0, LIMIT]
    assert 0 < control.pttl("budget:shared") <= ttl_before
    assert control.get("budget:shared") == b"5"
    print("a rejection neither spends more budget nor refreshes the window")

    assert control.pexpire("budget:shared", 50)
    deadline = time.monotonic() + 3
    while control.exists("budget:shared"):
        assert time.monotonic() < deadline, "test deadline did not expire"
        time.sleep(0.01)
    first = control.eval(acquire, 1, "budget:shared", LIMIT, WINDOW_MS)
    assert first[:2] == [1, 1] and 0 < first[2] <= WINDOW_MS
    print("after deliberately shortening and expiring the test window, the next admission starts at 1")