From b3c5e7085010c63bb721fe1afc2adbeb552919b7 Mon Sep 17 00:00:00 2001
From: Kumar Aditya <59607654+kumaraditya303@users.noreply.github.com>
Date: Fri, 23 Dec 2022 17:20:16 +0530
Subject: [PATCH] Add `asyncio_tcp` benchmark (#254)

---
 pyperformance/data-files/benchmarks/MANIFEST  |  1 +
 .../benchmarks/bm_asyncio_tcp/pyproject.toml  |  9 ++++
 .../bm_asyncio_tcp/run_benchmark.py           | 44 +++++++++++++++++++
 3 files changed, 54 insertions(+)
 create mode 100644 pyperformance/data-files/benchmarks/bm_asyncio_tcp/pyproject.toml
 create mode 100644 pyperformance/data-files/benchmarks/bm_asyncio_tcp/run_benchmark.py

diff --git a/pyperformance/data-files/benchmarks/MANIFEST b/pyperformance/data-files/benchmarks/MANIFEST
index 6c5d6e07..671efd15 100644
--- a/pyperformance/data-files/benchmarks/MANIFEST
+++ b/pyperformance/data-files/benchmarks/MANIFEST
@@ -7,6 +7,7 @@ async_tree	<local>
 async_tree_cpu_io_mixed	<local:async_tree>
 async_tree_io	<local:async_tree>
 async_tree_memoization	<local:async_tree>
+asyncio_tcp	<local>
 concurrent_imap	<local>
 coroutines	<local>
 coverage	<local>
diff --git a/pyperformance/data-files/benchmarks/bm_asyncio_tcp/pyproject.toml b/pyperformance/data-files/benchmarks/bm_asyncio_tcp/pyproject.toml
new file mode 100644
index 00000000..2e5f1e47
--- /dev/null
+++ b/pyperformance/data-files/benchmarks/bm_asyncio_tcp/pyproject.toml
@@ -0,0 +1,9 @@
+[project]
+name = "pyperformance_bm_asyncio_tcp"
+requires-python = ">=3.8"
+dependencies = ["pyperf"]
+urls = {repository = "https://github.com/python/pyperformance"}
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "asyncio_tcp"
diff --git a/pyperformance/data-files/benchmarks/bm_asyncio_tcp/run_benchmark.py b/pyperformance/data-files/benchmarks/bm_asyncio_tcp/run_benchmark.py
new file mode 100644
index 00000000..e400fcd1
--- /dev/null
+++ b/pyperformance/data-files/benchmarks/bm_asyncio_tcp/run_benchmark.py
@@ -0,0 +1,44 @@
+"""
+Benchmark for asyncio TCP server and client performance
+transferring 10MB of data.
+
+Author: Kumar Aditya
+"""
+
+
+import asyncio
+from pyperf import Runner
+
+
+CHUNK_SIZE = 1024 ** 2 * 10
+
+
+async def handle_echo(reader: asyncio.StreamReader,
+                      writer: asyncio.StreamWriter) -> None:
+    data = b'x' * CHUNK_SIZE
+    for _ in range(100):
+        writer.write(data)
+        await writer.drain()
+    writer.close()
+    await writer.wait_closed()
+
+
+async def main() -> None:
+    server = await asyncio.start_server(handle_echo, '127.0.0.1', 8882)
+
+    async with server:
+        asyncio.create_task(server.start_serving())
+        reader, writer = await asyncio.open_connection('127.0.0.1', 8882)
+        data_len = 0
+        while True:
+            data = await reader.read(CHUNK_SIZE)
+            if not data:
+                break
+            data_len += len(data)
+        assert data_len == CHUNK_SIZE * 100
+        writer.close()
+        await writer.wait_closed()
+
+if __name__ == '__main__':
+    runner = Runner()
+    runner.bench_async_func('asyncio_tcp', main)