punkshell/scriptlib/stdin_race.py

import sys
import time
import asyncio
from subprocess import Popen, PIPE

RESET = "\033[0m"
C = "\033[32m"  # Child color green
P = "\033[33m"  # Parent color yellow


def usage():
    print("Usage:")
    print("       stdin_race.py pump <persecond> <maxcount>")
    print("       stdin_race.py parent")
    print("       stdin_race.py child <delay_ms>")
    print("\nExample:")
    print("       python stdin_race.py pump 35 50 | python stdin_race.py parent")
    sys.exit(1)


async def pump(persecond, maxcount):
    if persecond > 1000:
        print("WARNING: (>1000) sub-millisecond scheduling not available - will go full speed", file=sys.stderr)
        await asyncio.sleep(0.5)

    counter = -1
    ms = 1000 / persecond

    async def pump_emit():
        nonlocal counter
        try:
            counter += 1
            print(f".{counter}", end="", flush=True)
        except BrokenPipeError:
            return False
        return True

    async def pump_schedule():
        nonlocal counter
        while maxcount <= 0 or counter < maxcount - 1:
            if not await pump_emit():
                break
            await asyncio.sleep(ms / 1000)
        print("pump-done", file=sys.stderr, flush=True)

    await pump_schedule()


async def parent():
    print(f"{P}parent{RESET}", file=sys.stderr, flush=True)
    await asyncio.sleep(0.25)

    # Read the first chunk from stdin
    parent_chunk1 = sys.stdin.read(8)
    print(f"{P}{parent_chunk1}{RESET}", file=sys.stderr, flush=True)

    # Launch the child process
    child_proc = Popen(
        [sys.executable, __file__, "child", "150"],
        stdin=sys.stdin, stdout=PIPE, stderr=sys.stderr, text=True
    )

    # Forward stdin to the child process
    #while True:
    #    chunk = sys.stdin.read(1)
    #    if not chunk:
    #        break
    #    try:
    #        child_proc.stdin.write(chunk)
    #        #child_proc.stdin.flush()
    #    except BrokenPipeError:
    #        print("Broken pipe error, child process may have exited.", file=sys.stderr)
    #        break
    child_proc.wait()

    #child_proc.wait()

    print("parent-tail-read", file=sys.stderr, flush=True)
    while True:
        chunk = sys.stdin.read(1)
        if not chunk:
            break
        print(chunk, end="", file=sys.stderr, flush=True)

    print(f"\n{P}parent-done{RESET}", flush=True)


async def child(delay_ms):
    print(f"\n{C}child{RESET}", file=sys.stderr, flush=True)
    await asyncio.sleep(delay_ms / 1000)

    # Read exactly 16 characters from stdin
    chunk = sys.stdin.read(16)
    print(f"{C}{chunk}{RESET}", file=sys.stderr, flush=True)

    print("child-done", file=sys.stderr, flush=True)
    sys.exit(0)


def main():
    if len(sys.argv) < 2:
        usage()

    role = sys.argv[1]
    if role == "pump":
        if len(sys.argv) != 4:
            usage()
        persecond = int(sys.argv[2])
        maxcount = int(sys.argv[3])
        asyncio.run(pump(persecond, maxcount))
    elif role == "parent":
        if len(sys.argv) != 2:
            usage()
        asyncio.run(parent())
    elif role == "child":
        if len(sys.argv) != 3:
            usage()
        delay_ms = int(sys.argv[2])
        asyncio.run(child(delay_ms))
    else:
        usage()


if __name__ == "__main__":
    main()