"""After a ProcessPoolExecutor worker dies, shutdown(wait=False) and submit()
block for as long as another worker outlives SIGTERM.

    python repro.py ignore    # workers' SIGTERM handler does nothing
    python repro.py exit      # workers' SIGTERM handler calls sys.exit(0)
    python repro.py default   # control: default SIGTERM action

Bounded: after BOUND seconds the script SIGKILLs the workers itself.
"""
import concurrent.futures
import multiprocessing
import os
import signal
import sys
import threading
import time

BOUND = 10.0


def install_handler(mode):
    if mode == "ignore":
        signal.signal(signal.SIGTERM, lambda signum, frame: None)
    elif mode == "exit":
        signal.signal(signal.SIGTERM, lambda signum, frame: sys.exit(0))


def work(seconds):
    try:
        time.sleep(seconds)
    finally:
        time.sleep(0.05)  # a task's cleanup, e.g. closing a file
    return os.getpid()


def die(delay):
    time.sleep(delay)
    os._exit(1)


def in_thread(fn):
    outcome = []

    def run():
        try:
            fn()
            outcome.append("returned")
        except Exception as exc:
            outcome.append(f"raised {type(exc).__name__}")

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    return thread, outcome


def main():
    mode = sys.argv[1]
    ex = concurrent.futures.ProcessPoolExecutor(
        max_workers=2, mp_context=multiprocessing.get_context("spawn"),
        initializer=install_handler, initargs=(mode,))
    assert len({f.result() for f in [ex.submit(work, 0.5) for _ in range(2)]}) == 2
    workers = list(ex._processes.values())
    manager = ex._executor_manager_thread

    # One worker is running a task when the other one dies.
    busy = ex.submit(work, 1.0)
    ex.submit(die, 0.2)
    try:
        busy.result()
    except concurrent.futures.process.BrokenProcessPool:
        print("pool is broken")

    calls = {"submit()": in_thread(lambda: ex.submit(work, 0)),
             "shutdown(wait=False)": in_thread(lambda: ex.shutdown(wait=False))}
    if hasattr(ex, "kill_workers"):  # 3.14+
        calls["kill_workers()"] = in_thread(ex.kill_workers)
    deadline = time.monotonic() + BOUND
    for thread, _ in calls.values():
        thread.join(max(0.0, deadline - time.monotonic()))

    blocked = [name for name, (thread, _) in calls.items() if thread.is_alive()]
    for name, (_, outcome) in calls.items():
        print(f"{name}: {outcome[0] if outcome else f'still blocked after {BOUND:.0f} s'}")
    if blocked:
        frame = sys._current_frames().get(manager.ident)
        stack = []
        while frame is not None:
            stack.append(frame.f_code.co_name)
            frame = frame.f_back
        print("executor manager thread:", " <- ".join(stack))
        alive = [p for p in workers if p.is_alive()]
        print("workers still alive:", [p.pid for p in alive])
        for p in alive:
            p.kill()
        for thread, _ in calls.values():
            thread.join(BOUND)
        for name in blocked:
            outcome = calls[name][1]
            print(f"after SIGKILL: {name} {outcome[0] if outcome else 'STILL BLOCKED'}")
    for child in multiprocessing.active_children():
        child.kill()
    if any(thread.is_alive() for thread, _ in calls.values()):
        os._exit(2)  # the atexit hook would join the manager thread
    sys.exit(1 if blocked else 0)


if __name__ == "__main__":
    main()