"""
Demonstrate hang in `proc.stdin.wait_closed()`. The hang only occurs on
Windows.

This code works fine on MacOS, Linux.

Correct output:

```
BrokenPipe 1 - correct
BrokenPipe 2 - correct
```
"""

import asyncio
import sys

PIPE_MAX_SIZE = 4 * 1024 * 1024 + 1


async def _start_proc():
    "Start process that sleeps for 2 seconds then exits."
    return await asyncio.create_subprocess_exec(
        sys.executable,
        "-c",
        "import time; time.sleep(2)",
        stdin=asyncio.subprocess.PIPE,
    )


async def race_condition1():
    "It's random whether this code will trigger the hang in `wait_closed`"
    proc = await _start_proc()
    try:
        data = b"a" * PIPE_MAX_SIZE
        proc.stdin.write(data)
        await proc.stdin.drain()
    except BrokenPipeError:
        pass
    proc.stdin.close()

    try:
        # `wait_closed` hangs if we are unlucky.
        await proc.stdin.wait_closed()
    except BrokenPipeError:
        print("BrokenPipe 1 - correct")


async def race_condition2():
    "This code forces the transport to be aborted _after_ the proc.stdin.close()."
    proc = await _start_proc()
    try:
        data = b"a" * PIPE_MAX_SIZE
        proc.stdin.write(data)
        await asyncio.wait_for(proc.stdin.drain(), 1)
    except asyncio.TimeoutError:
        pass
    proc.stdin.close()

    # Here, the process exits and proc.stdin's transport is aborted.
    await asyncio.sleep(2)

    try:
        # `wait_closed` will hang on Windows every time.
        await proc.stdin.wait_closed()
    except BrokenPipeError:
        print("BrokenPipe 2 - correct")


asyncio.run(race_condition1())
asyncio.run(race_condition2())