Add asyncio_tcp benchmark (#254)

kumaraditya303 · web-flow · commit b3c5e7085010 · 2022-12-23T17:20:16.000+05:30
diff --git 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
@@ -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
@@ -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)
