负载削减保护系统

from concurrent.futures import ThreadPoolExecutor
from threading import Barrier, BoundedSemaphore, Event, Lock


def exercise(gated):
    slots = BoundedSemaphore(3)  # Bounds queued plus running work, not just threads.
    started, release, lock = Barrier(3), Event(), Lock()
    running, began = set(), []

    def held(job):
        with lock:
            running.add(job)
            began.append(job)
        try:
            if job < 2:
                started.wait(timeout=5)
            assert release.wait(timeout=5), "test coordinator did not release workers"
            return job
        finally:
            with lock:
                running.remove(job)

    with ThreadPoolExecutor(max_workers=2) as pool:
        def submit(function, *args):
            if not gated:
                return pool.submit(function, *args)
            if not slots.acquire(blocking=False):
                return None  # Explicit local admission refusal; no future was enqueued.
            try:
                future = pool.submit(function, *args)
            except RuntimeError:
                slots.release()
                raise
            # Completion, failure, and cancellation before execution all return the slot.
            future.add_done_callback(lambda _: slots.release())
            return future

        futures = [submit(held, job) for job in range(12)]
        accepted = [future for future in futures if future is not None]
        replacement = None
        try:
            started.wait(timeout=5)
            with lock:
                assert running == {0, 1} and set(began) == {0, 1}
            assert all(not future.done() for future in accepted)
            if gated:
                assert len(accepted) == 3 and futures[3:] == [None] * 9
                assert futures[2].cancel()  # Both workers are held, so job 2 has not started.
                replacement = submit(held, 12)
                assert replacement is not None
                assert submit(held, 13) is None
                print("gated burst: 2 running + 1 queued, 9 initial refusals; cancellation returns a slot")
            else:
                assert len(accepted) == 12
                print("ungated burst: 2 workers accept 12 unfinished tasks; the other 10 still wait")
        finally:
            release.set()
        completed = [future.result(timeout=5) for future in accepted if not future.cancelled()]
        if replacement is not None:
            completed.append(replacement.result(timeout=5))
        assert sorted(completed) == ([0, 1, 12] if gated else list(range(12)))

        if gated:
            def fail():
                raise ValueError("injected task failure")

            failed = submit(fail)
            assert failed is not None
            try:
                failed.result(timeout=5)
            except ValueError as error:
                assert str(error) == "injected task failure"
            else:
                raise AssertionError("the failed task must preserve its failure")

    if gated:
        # The executor is now closed; a failed submission must not consume a slot.
        try:
            submit(lambda: None)
        except RuntimeError:
            pass
        else:
            raise AssertionError("submission to a closed executor must fail")
        acquired = [slots.acquire(blocking=False) for _ in range(3)]
        try:
            assert acquired == [True, True, True]
            assert not slots.acquire(blocking=False)
        finally:
            for success in acquired:
                if success:
                    slots.release()
        print("completion, task failure, cancellation, and closed-executor submission leave all 3 slots reusable")


exercise(False)
exercise(True)