Add benches to nats-client

Signed-off-by: Casper Beyer <[email protected]>

Author: caspervonb

.github/workflows/bench.yml

@@ -0,0 +1,118 @@
+name: bench
+
+on:
+    pull_request:
+        paths:
+            - "nats-client/**"
+    push:
+        branches:
+            - main
+        paths:
+            - "nats-client/**"
+    workflow_dispatch:
+
+jobs:
+    head:
+        name: Head
+        runs-on: ubuntu-latest
+        strategy:
+            matrix:
+                python-version: ["3.11", "3.12", "3.13", "pypy-3.10"]
+        steps:
+            - name: Checkout head branch
+              uses: actions/checkout@v5
+
+            - name: Set up Python
+              uses: actions/setup-python@v6
+              with:
+                  python-version: ${{ matrix.python-version }}
+
+            - name: Install uv
+              uses: astral-sh/setup-uv@v6
+
+            - name: Install dependencies
+              run: uv sync --dev
+              working-directory: nats-client
+
+            - name: Run benchmarks
+              run: |
+                  uv run pytest benches/ -k bench_ --benchmark-json=../head-bench-${{ matrix.python-version }}.json -q
+              working-directory: nats-client
+
+            - name: Upload head benchmark results
+              uses: actions/upload-artifact@v4
+              with:
+                  name: head-benchmark-${{ matrix.python-version }}
+                  path: head-bench-${{ matrix.python-version }}.json
+
+    base:
+        name: Base
+        runs-on: ubuntu-latest
+        strategy:
+            matrix:
+                python-version: ["3.11", "3.12", "3.13", "pypy-3.10"]
+        steps:
+            - name: Checkout base branch
+              uses: actions/checkout@v5
+              with:
+                  ref: ${{ github.base_ref || 'main' }}
+
+            - name: Set up Python
+              uses: actions/setup-python@v6
+              with:
+                  python-version: ${{ matrix.python-version }}
+
+            - name: Install uv
+              uses: astral-sh/setup-uv@v6
+
+            - name: Install dependencies
+              run: uv sync --dev
+              working-directory: nats-client
+
+            - name: Run benchmarks
+              run: |
+                  uv run pytest benches/ -k bench_ --benchmark-json=../base-bench-${{ matrix.python-version }}.json -q
+              working-directory: nats-client
+
+            - name: Upload base benchmark results
+              uses: actions/upload-artifact@v4
+              with:
+                  name: base-benchmark-${{ matrix.python-version }}
+                  path: base-bench-${{ matrix.python-version }}.json
+
+    compare:
+        name: Compare
+        runs-on: ubuntu-latest
+        needs: [head, base]
+        strategy:
+            matrix:
+                python-version: ["3.11", "3.12", "3.13", "pypy-3.10"]
+        steps:
+            - name: Checkout code
+              uses: actions/checkout@v5
+
+            - name: Set up Python
+              uses: actions/setup-python@v6
+              with:
+                  python-version: ${{ matrix.python-version }}
+
+            - name: Install uv
+              uses: astral-sh/setup-uv@v6
+
+            - name: Install dependencies
+              run: uv sync --dev
+              working-directory: nats-client
+
+            - name: Download head benchmark results
+              uses: actions/download-artifact@v4
+              with:
+                  name: head-benchmark-${{ matrix.python-version }}
+
+            - name: Download base benchmark results
+              uses: actions/download-artifact@v4
+              with:
+                  name: base-benchmark-${{ matrix.python-version }}
+
+            - name: Compare benchmarks
+              run: |
+                  uv run pytest-benchmark compare base-bench-${{ matrix.python-version }}.json head-bench-${{ matrix.python-version }}.json --group-by=name --sort=mean --columns=mean,median

.github/workflows/test.yml

@@ -42,13 +42,8 @@ jobs:
 
             - name: Run tests
               run: |
-<<<<<<< HEAD
                   uv run flake8 --ignore="W391, W503, W504, E501" ./nats/src/nats/js/
                   uv run pytest -x -vv -s --continue-on-collection-errors ./nats/tests
-=======
-                  pipenv run flake8 --ignore="W391, W503, W504" ./nats/js/
-                  pipenv run pytest -x -vv -s --continue-on-collection-errors ./tests
->>>>>>> 44fb982 (Add nats-server package)
               env:
                   PATH: $HOME/nats-server:$PATH
 

nats-client/benches/__init__.py

@@ -0,0 +1 @@
+"""Benchmarks for nats-client package."""

nats-client/benches/bench_protocol.py

@@ -0,0 +1,103 @@
+"""Benchmarks for NATS protocol parsing operations."""
+
+import asyncio
+
+import pytest
+from nats.client.protocol.message import parse, parse_headers
+
+
+class MockReader:
+    """Mock reader for testing protocol parsing."""
+
+    def __init__(self, data: bytes):
+        self.data = data
+        self.pos = 0
+
+    async def readline(self) -> bytes:
+        """Read a line from the data."""
+        if self.pos >= len(self.data):
+            return b""
+
+        start = self.pos
+        while self.pos < len(self.data) and self.data[self.pos:self.pos + 1] != b"\n":
+            self.pos += 1
+
+        if self.pos < len(self.data):
+            self.pos += 1
+
+        return self.data[start:self.pos]
+
+    async def readexactly(self, n: int) -> bytes:
+        """Read exactly n bytes from the data."""
+        if self.pos + n > len(self.data):
+            raise asyncio.IncompleteReadError(
+                self.data[self.pos:], n - (len(self.data) - self.pos)
+            )
+
+        result = self.data[self.pos:self.pos + n]
+        self.pos += n
+        return result
+
+
+def bench_parse_msg(benchmark):
+    """Benchmark parsing MSG protocol message."""
+    msg_data = b"MSG test.subject 1 5\r\nhello\r\n"
+
+    def run_parse():
+        async def _parse():
+            reader = MockReader(msg_data)
+            return await parse(reader)
+        return asyncio.run(_parse())
+
+    benchmark(run_parse)
+
+
+def bench_parse_hmsg(benchmark):
+    """Benchmark parsing HMSG protocol message."""
+    msg_data = b"HMSG test.subject 1 reply 44 49\r\nNATS/1.0\r\nContent-Type: application/json\r\n\r\nhello\r\n"
+
+    def run_parse():
+        async def _parse():
+            reader = MockReader(msg_data)
+            return await parse(reader)
+        return asyncio.run(_parse())
+
+    benchmark(run_parse)
+
+
+def bench_parse_headers_simple(benchmark):
+    """Benchmark parsing simple headers."""
+    header_data = b"NATS/1.0\r\nContent-Type: application/json\r\n\r\n"
+
+    def run_parse():
+        return parse_headers(header_data)
+
+    benchmark(run_parse)
+
+
+def bench_parse_headers_with_status(benchmark):
+    """Benchmark parsing headers with status code."""
+    header_data = b"NATS/1.0 503 No Responders\r\n\r\n"
+
+    def run_parse():
+        return parse_headers(header_data)
+
+    benchmark(run_parse)
+
+
+def bench_parse_headers_multiple(benchmark):
+    """Benchmark parsing multiple headers."""
+    header_data = (
+        b"NATS/1.0\r\n"
+        b"Content-Type: application/json\r\n"
+        b"X-Custom-Header: value1\r\n"
+        b"X-Another-Header: value2\r\n"
+        b"X-Multi-Value: a\r\n"
+        b"X-Multi-Value: b\r\n"
+        b"\r\n"
+    )
+
+    def run_parse():
+        return parse_headers(header_data)
+
+    benchmark(run_parse)