Merge pull request #38 from taskiq-python/chore-add-benchmarks

chore: add benchmarks wth codspeed

Author: danfimov

.github/workflows/codspeed.yml

@@ -0,0 +1,51 @@
+name: CodSpeed
+
+on:
+  push:
+    branches: [ main ]
+    paths-ignore:
+      - 'docs/**'
+      - '**/*.md'
+  pull_request:
+    branches: [ main ]
+    paths-ignore:
+      - 'docs/**'
+      - '**/*.md'
+
+permissions:
+  contents: read
+  id-token: write # required for OIDC authentication with CodSpeed
+
+jobs:
+  benchmarks:
+    name: Benchmarks
+    runs-on: ubuntu-latest
+    services:
+      ministack:
+        image: ministackorg/ministack:1.3.53
+        ports:
+          - 4566:4566
+        env:
+          AWS_DEFAULT_REGION: us-east-1
+          GATEWAY_PORT: 4566
+          MINISTACK_ACCOUNT_ID: "000000000000"
+          MINISTACK_REGION: us-east-1
+          LOG_LEVEL: INFO
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          persist-credentials: false
+      - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
+        with:
+          python-version: "3.13"
+      - uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2
+        with:
+          enable-cache: true
+          version: "latest"
+      - name: Install dependencies
+        run: uv sync --all-extras
+      - name: Run the benchmarks
+        uses: CodSpeedHQ/action@4296e51e7041e24dadb86d1d6e8b9320d223dbe8 # v5.0.3
+        with:
+          mode: simulation
+          run: uv run pytest tests/ --codspeed

docker-compose.yml

@@ -12,7 +12,7 @@ services:
       LOG_LEVEL: INFO
       PERSIST_STATE: "1"
     healthcheck:
-      test: ["CMD", "curl", "-f", "http://localhost:4566/_ministack/health"]
+      test: ["CMD", "busybox", "wget", "-q", "-O", "/dev/null", "http://127.0.0.1:4566/_ministack/health"]
       interval: 2s
       timeout: 5s
       retries: 30

pyproject.toml

@@ -44,19 +44,20 @@ dev = [
     {include-group = "lint"},
     {include-group = "test"},
     {include-group = "types"},
-    "prek>=0.4.3",
+    "prek>=0.5.0",
 ]
 test = [
     "pytest>=9.0.3",
     "pytest-asyncio>=0.23.8",
+    "pytest-codspeed>=5.0.3",
 ]
 lint = [
     "bandit>=1.9.4",
-    "ruff>=0.15.15",
-    "zizmor>=1.25.2",
+    "ruff>=0.16.5",
+    "zizmor>=1.29.0",
 ]
 types = [
-    "mypy>=2.1.0",
+    "mypy>=2.3.1",
     "types-aiobotocore[essential]>=3.7.0",
 ]
 examples = [
@@ -82,6 +83,7 @@ asyncio_mode = "auto"
 markers = [
     "unit: marks unit tests",
     "integration: marks tests with real infrastructure env",
+    "benchmark: marks performance benchmarks (run with --codspeed)",
 ]
 
 [tool.coverage.report]
@@ -129,6 +131,7 @@ ignore = [
 
     "EM101",
     "TRY003",
+    "CPY001",
 ]
 
 [tool.ruff.lint.per-file-ignores]

tests/benchmarks/__init__.py

@@ -0,0 +1,62 @@
+from collections.abc import AsyncGenerator
+
+import pytest
+from taskiq.acks import AckableMessage
+from taskiq.result import TaskiqResult
+
+from tests.conftest import AWSCredentials, _queue_name_from_url
+
+from taskiq_sqs import S3ResultBackend, SQSBroker
+
+
+BATCH_SIZE = 10
+RESULT_TASK_ID = "benchmark-task"
+
+
+async def ack(message: AckableMessage) -> None:
+    result = message.ack()
+    if result is not None:
+        await result
+
+
+@pytest.fixture
+async def bench_broker(
+    aws_credentials: AWSCredentials,
+    sqs_queue: str,
+) -> AsyncGenerator[SQSBroker, None]:
+    """The shared `sqs_broker`, but tuned for measuring a round-trip.
+
+    Two settings differ from the default broker: long polling, because with `wait_time_seconds=0` a receive returns
+    immediately and usually empty, so a round-trip would measure idle polls; and a batch size that lets a single
+    `ReceiveMessage` call serve the whole batch.
+    """
+    broker = SQSBroker(
+        queue_name=_queue_name_from_url(sqs_queue),
+        wait_time_seconds=1,
+        max_number_of_messages=BATCH_SIZE,
+        **aws_credentials,
+    )
+    await broker.startup()
+    yield broker
+    await broker.shutdown()
+
+
+@pytest.fixture
+async def listener(bench_broker: SQSBroker) -> AsyncGenerator[AsyncGenerator[AckableMessage, None], None]:
+    """A single `listen()` generator, created outside the measured test body."""
+    messages = bench_broker.listen()
+    yield messages
+    await messages.aclose()
+
+
+@pytest.fixture
+async def stored_result(s3_backend: S3ResultBackend[str]) -> TaskiqResult[str]:
+    """A result already written to the bucket, so read benchmarks only measure reads."""
+    result: TaskiqResult[str] = TaskiqResult(
+        is_err=False,
+        return_value="benchmark",
+        execution_time=0.1,
+        log=None,
+    )
+    await s3_backend.set_result(RESULT_TASK_ID, result)
+    return result

tests/benchmarks/conftest.py

@@ -0,0 +1,43 @@
+from collections.abc import AsyncGenerator
+
+import pytest
+from taskiq import BrokerMessage
+from taskiq.acks import AckableMessage
+
+from .conftest import BATCH_SIZE, ack
+from taskiq_sqs import SQSBroker
+
+
+@pytest.mark.benchmark
+async def test_build_kick_kwargs(bench_broker: SQSBroker, broker_message: BrokerMessage) -> None:
+    await bench_broker._build_kick_kwargs(broker_message)
+
+
+@pytest.mark.benchmark
+async def test_kick(bench_broker: SQSBroker, broker_message: BrokerMessage) -> None:
+    await bench_broker.kick(broker_message)
+
+
+@pytest.mark.benchmark
+async def test_kick_and_listen_roundtrip(
+    bench_broker: SQSBroker,
+    broker_message: BrokerMessage,
+    listener: AsyncGenerator[AckableMessage, None],
+) -> None:
+    await bench_broker.kick(broker_message)
+    message = await anext(listener)
+    await ack(message)
+
+
+@pytest.mark.benchmark
+async def test_kick_and_listen_roundtrip_batch(
+    bench_broker: SQSBroker,
+    broker_message: BrokerMessage,
+    listener: AsyncGenerator[AckableMessage, None],
+) -> None:
+    for _ in range(BATCH_SIZE):
+        await bench_broker.kick(broker_message)
+     # SQS may return fewer messages than requested per call, so keep pulling until the whole batch is back.
+    for _ in range(BATCH_SIZE):
+        message = await anext(listener)
+        await ack(message)

tests/benchmarks/test_broker.py

@@ -0,0 +1,27 @@
+import pytest
+from taskiq.result import TaskiqResult
+
+from tests.benchmarks.conftest import RESULT_TASK_ID
+
+from taskiq_sqs import S3ResultBackend
+
+
+@pytest.mark.benchmark
+async def test_set_result(s3_backend: S3ResultBackend[str], stored_result: TaskiqResult[str]) -> None:
+    await s3_backend.set_result(RESULT_TASK_ID, stored_result)
+
+
+@pytest.mark.benchmark
+async def test_get_result(
+    s3_backend: S3ResultBackend[str],
+    stored_result: TaskiqResult[str],  # noqa: ARG001
+) -> None:
+    await s3_backend.get_result(RESULT_TASK_ID)
+
+
+@pytest.mark.benchmark
+async def test_is_result_ready(
+    s3_backend: S3ResultBackend[str],
+    stored_result: TaskiqResult[str],  # noqa: ARG001
+) -> None:
+    await s3_backend.is_result_ready(RESULT_TASK_ID)