# Time one large bytes write that has to queue behind data already buffered
# in an asyncio socket transport.
# Usage: python bench_write_copy.py [size in MiB, default 48]
import asyncio
import socket
import sys
import time

SIZE = int(sys.argv[1] if len(sys.argv) > 1 else 48) * 1024 * 1024
PREFIX = 2 * 1024 * 1024


class Reader(asyncio.Protocol):
    """Counts received bytes and resolves `done` once all have arrived."""

    def __init__(self, expected, done):
        self.received = 0
        self.expected = expected
        self.done = done

    def data_received(self, data):
        self.received += len(data)
        if self.received == self.expected:
            self.done.set_result(None)


async def main():
    loop = asyncio.get_running_loop()
    sock_a, sock_b = socket.socketpair()
    done = loop.create_future()
    reader, _ = await loop.connect_accepted_socket(
        lambda: Reader(PREFIX + SIZE, done), sock_b)
    writer, _ = await loop.connect_accepted_socket(asyncio.Protocol, sock_a)

    payload = b"x" * SIZE
    start = time.perf_counter()
    writer.write(b"p" * PREFIX)  # fills the socket, so the next write is buffered
    writer.write(payload)        # sent in many partial chunks
    await done
    elapsed = time.perf_counter() - start

    writer.close()
    reader.close()
    await asyncio.sleep(0)
    print(f"{SIZE >> 20} MiB write: {elapsed * 1000:.1f} ms")


asyncio.run(main())