Unpickler BUILD can traverse a concurrently mutated shared state dictionary #157124

#!/usr/bin/env python3
"""Race Unpickler BUILD iteration with mutation of its shared state dict."""

import pickle
import sys
import threading
import time


MODE = sys.argv[2] if len(sys.argv) > 2 else "state"
if MODE not in {"state", "slot"}:
    raise ValueError("mode must be 'state' or 'slot'")

SHARED = {f"field_{i}": object() for i in range(4096)}
KEYS = tuple(SHARED)
STOP = threading.Event()
FAILURES = []


def get_shared():
    return SHARED


class SharedState(dict):
    def __reduce__(self):
        return get_shared, ()


class Victim:
    def __reduce__(self):
        state = SharedState()
        if MODE == "slot":
            state = (None, state)
        return Victim, (), state


PAYLOAD = pickle.dumps(Victim(), protocol=2)


def mutate():
    while not STOP.is_set():
        SHARED.clear()
        SHARED.update({key: object() for key in KEYS})


def load_worker(deadline):
    count = 0
    try:
        while time.monotonic() < deadline and not STOP.is_set():
            obj = pickle.loads(PAYLOAD)
            count += 1
    except BaseException as exc:
        FAILURES.append((type(exc).__name__, repr(exc), count))
        STOP.set()


def main():
    seconds = float(sys.argv[1]) if len(sys.argv) > 1 else 10.0
    deadline = time.monotonic() + seconds
    threads = [threading.Thread(target=mutate) for _ in range(2)]
    threads += [threading.Thread(target=load_worker, args=(deadline,)) for _ in range(8)]
    for thread in threads:
        thread.start()
    for thread in threads[2:]:
        thread.join()
    STOP.set()
    for thread in threads[:2]:
        thread.join()
    print("mode=", MODE, "gil=", sys._is_gil_enabled(), "failures=", FAILURES)
    return bool(FAILURES)


raise SystemExit(main())