init

Author: lesnik512

.github/workflows/ci.yml

@@ -0,0 +1,52 @@
+name: main
+
+on:
+  push:
+    branches:
+      - main
+  pull_request: {}
+
+concurrency:
+  group: ${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: extractions/setup-just@v2
+      - uses: astral-sh/setup-uv@v3
+        with:
+          enable-cache: true
+          cache-dependency-glob: "**/pyproject.toml"
+      - run: uv python install 3.11
+      - run: just install lint-ci
+
+  pytest:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version:
+          - "3.11"
+          - "3.12"
+          - "3.13"
+          - "3.14"
+    steps:
+      - uses: actions/checkout@v4
+      - uses: extractions/setup-just@v2
+      - uses: astral-sh/setup-uv@v3
+        with:
+          enable-cache: true
+          cache-dependency-glob: "**/pyproject.toml"
+      - run: uv python install ${{ matrix.python-version }}
+      - run: just install
+      - run: just test . --cov=. --cov-report xml
+      - uses: codecov/codecov-action@v4.0.1
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+        with:
+          files: ./coverage.xml
+          flags: unittests
+          name: codecov-${{ matrix.python-version }}

.github/workflows/publish.yml

@@ -0,0 +1,17 @@
+name: Publish Package
+
+on:
+  release:
+    types:
+      - published
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: extractions/setup-just@v2
+      - uses: astral-sh/setup-uv@v3
+      - run: just publish
+        env:
+          PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}

.gitignore

@@ -0,0 +1,20 @@
+# Generic things
+*.pyc
+*~
+__pycache__/*
+*.swp
+*.sqlite3
+*.map
+.vscode
+.idea
+.DS_Store
+.env
+.pytest_cache
+.ruff_cache
+.coverage
+htmlcov/
+coverage.xml
+pytest.xml
+.python-version
+.venv
+dist

CLAUDE.md

@@ -0,0 +1,48 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Commands
+
+```bash
+just install       # lock + sync all deps (run after pulling or changing pyproject.toml)
+just lint          # eof-fixer, ruff format, ruff check --fix, ty check
+just lint-ci       # same but no auto-fix (used in CI)
+just test          # pytest with coverage
+just test-branch   # pytest with branch coverage
+just publish       # bump version to $GITHUB_REF_NAME, build, publish to PyPI
+```
+
+Run a single test file or test by name:
+```bash
+uv run --no-sync pytest tests/committer/test_kafka_committer.py
+uv run --no-sync pytest -k test_committer_logs_task_exceptions
+```
+
+## Architecture
+
+The library provides concurrent Kafka message processing for FastStream. Three modules are exposed:
+
+**`processing.py` — `KafkaConcurrentHandler` (singleton)**
+The core engine. Implements the singleton pattern (one instance per process) using `__new__` + `threading.Lock`. Manages:
+- An `asyncio.Semaphore` for concurrency limiting (`concurrency_limit=0` disables it)
+- A set of in-flight `asyncio.Task`s
+- A background observer task that periodically calls `_check_tasks_health()` to discard stale completed tasks
+- Signal handlers (SIGTERM/SIGINT/SIGQUIT) that trigger graceful shutdown
+- Optional integration with `KafkaBatchCommitter` when `enable_batch_commit=True`
+
+Key design: `handle_task()` fires-and-forgets coroutines as asyncio tasks. Message filtering by consumer group is done via the `topic_group` header — messages whose header doesn't match the consumer's `_group_id` are skipped.
+
+**`middleware.py` — FastStream middleware + lifecycle functions**
+- `KafkaConcurrentProcessingMiddleware`: FastStream `BaseMiddleware` subclass. Its `consume_scope` wraps each incoming message in a task submitted to `KafkaConcurrentHandler` (retrieved from FastStream's context).
+- `initialize_concurrent_processing(context, ...)`: call on app startup to create and start the handler, storing it in FastStream's global context.
+- `stop_concurrent_processing(context)`: call on app shutdown; resets the singleton so it can be re-initialized (important for tests).
+
+**`batch_committer.py` — `KafkaBatchCommitter`**
+Runs as a background asyncio task (spawned via `spawn()`). Collects `KafkaCommitTask` objects from a queue, batches them by topic-partition, waits for each task's asyncio future to complete, then commits the max offset per partition to Kafka. Batching is triggered by timeout or batch size, whichever comes first. `CommitterIsDeadError` is raised to callers if the committer's main task has died.
+
+## Key patterns
+
+- **Singleton reset in tests**: `KafkaConcurrentHandler._initialized = False` and `._instance = None` must be reset between tests (done in conftest fixtures via `stop_concurrent_processing`).
+- **Type suppression**: use `# ty: ignore[rule-name]` (not `# type: ignore`) for ty type checker suppressions.
+- **No `from __future__ import annotations`**: annotations are evaluated eagerly; `typing.Self`/`typing.Never` are used directly (requires Python ≥ 3.11).

Justfile

@@ -0,0 +1,29 @@
+default: install lint test
+
+install:
+    uv lock --upgrade
+    uv sync --all-extras --frozen --group lint
+
+lint:
+    uv run eof-fixer .
+    uv run ruff format
+    uv run ruff check --fix
+    uv run ty check
+
+lint-ci:
+    uv run eof-fixer . --check
+    uv run ruff format --check
+    uv run ruff check --no-fix
+    uv run ty check
+
+test *args:
+    uv run --no-sync pytest {{ args }}
+
+test-branch:
+    @just test --cov-branch
+
+publish:
+    rm -rf dist
+    uv version $GITHUB_REF_NAME
+    uv build
+    uv publish --token $PYPI_TOKEN

README.md

@@ -0,0 +1 @@
+# kafka-utils

faststream_concurrent_aiokafka/__init__.py

@@ -0,0 +1,12 @@
+from faststream_concurrent_aiokafka.middleware import (
+    KafkaConcurrentProcessingMiddleware,
+    initialize_concurrent_processing,
+    stop_concurrent_processing,
+)
+
+
+__all__ = [
+    "KafkaConcurrentProcessingMiddleware",
+    "initialize_concurrent_processing",
+    "stop_concurrent_processing",
+]

faststream_concurrent_aiokafka/batch_committer.py

@@ -0,0 +1,197 @@
+import asyncio
+import contextlib
+import dataclasses
+import itertools
+import logging
+import threading
+import typing
+
+from faststream.kafka import TopicPartition
+
+
+if typing.TYPE_CHECKING:
+    from aiokafka import AIOKafkaConsumer
+
+
+logger = logging.getLogger(__name__)
+
+
+SHUTDOWN_TIMEOUT_SEC: typing.Final = 20
+
+
+class CommitterIsDeadError(Exception): ...
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True, slots=True)
+class KafkaCommitTask:
+    asyncio_task: asyncio.Task[typing.Any]
+    topic_partition: TopicPartition
+    offset: int
+    consumer: typing.Any
+
+
+class KafkaBatchCommitter:
+    def __init__(
+        self,
+        commit_batch_timeout_sec: float = 10.0,
+        commit_batch_size: int = 10,
+    ) -> None:
+        self._messages_queue: asyncio.Queue[KafkaCommitTask] = asyncio.Queue()
+        self._asyncio_commit_process_task: asyncio.Task[typing.Any] | None = None
+        self._flush_batch_event = asyncio.Event()
+
+        self._commit_batch_timeout_sec = commit_batch_timeout_sec
+        self._commit_batch_size = commit_batch_size
+        self._shutdown_timeout = SHUTDOWN_TIMEOUT_SEC
+
+        self._spawn_lock = threading.Lock()
+
+    def _check_is_commit_task_running(self) -> None:
+        is_commit_task_running: typing.Final[bool] = bool(
+            self._asyncio_commit_process_task
+            and not self._asyncio_commit_process_task.cancelled()
+            and not self._asyncio_commit_process_task.done(),
+        )
+        if not is_commit_task_running:
+            msg: typing.Final = "Committer main task is not running"
+            raise CommitterIsDeadError(msg)
+
+    def _flush_tasks_queue(self) -> list[KafkaCommitTask]:
+        tasks_to_return: typing.Final[list[KafkaCommitTask]] = []
+        while not self._messages_queue.empty():
+            tasks_to_return.append(self._messages_queue.get_nowait())
+        return tasks_to_return
+
+    async def _populate_commit_batch(self) -> tuple[list[KafkaCommitTask], bool]:
+        uncommited_tasks: typing.Final[list[KafkaCommitTask]] = []
+        should_shutdown = False
+        queue_get_task: asyncio.Task[typing.Any] | None = None
+        flush_wait_task: asyncio.Task[typing.Any] | None = None
+        timeout_task: asyncio.Task[typing.Any] | None = None
+        try:
+            timeout_task = asyncio.create_task(asyncio.sleep(self._commit_batch_timeout_sec))
+            while len(uncommited_tasks) < self._commit_batch_size:
+                queue_get_task = asyncio.create_task(self._messages_queue.get())
+                flush_wait_task = asyncio.create_task(self._flush_batch_event.wait())
+                await asyncio.wait([queue_get_task, flush_wait_task, timeout_task], return_when=asyncio.FIRST_COMPLETED)
+
+                if queue_get_task.done():
+                    uncommited_tasks.append(queue_get_task.result())
+                else:
+                    queue_get_task.cancel()
+
+                # commit_all is called
+                if flush_wait_task.done():
+                    queue_get_task.cancel()
+                    uncommited_tasks.extend(self._flush_tasks_queue())
+                    self._flush_batch_event.clear()
+                    timeout_task.cancel()
+                    should_shutdown = True
+                    break
+                flush_wait_task.cancel()
+
+                if timeout_task.done():
+                    logger.debug("Timeout exceeded, batch contains %s elements", len(uncommited_tasks))
+                    break
+
+            logger.debug("Batch condition reached with %s elements", len(uncommited_tasks))
+        except asyncio.CancelledError:
+            should_shutdown = True
+            uncommited_tasks.extend(self._flush_tasks_queue())
+
+        for task in (queue_get_task, flush_wait_task, timeout_task):
+            task and task.cancel()
+
+        return uncommited_tasks, should_shutdown
+
+    async def _call_committer(
+        self, tasks_batch: list[KafkaCommitTask], partitions_to_offsets: dict[TopicPartition, int]
+    ) -> bool:
+        if not partitions_to_offsets:
+            return True
+        commit_succeeded = True
+        consumer: typing.Final[AIOKafkaConsumer] = tasks_batch[0].consumer
+        try:
+            await consumer.commit(partitions_to_offsets)
+        except Exception as exc:
+            commit_succeeded = False
+            logger.exception("Error during commit to kafka", exc_info=exc)
+            for task in tasks_batch:
+                await self._messages_queue.put(task)
+        return commit_succeeded
+
+    async def _commit_tasks_batch(self, tasks_batch: list[KafkaCommitTask]) -> bool:
+        partitions_to_tasks: typing.Final = itertools.groupby(
+            sorted(tasks_batch, key=lambda x: x.topic_partition), lambda x: x.topic_partition
+        )
+
+        results: typing.Final = await asyncio.gather(
+            *[task.asyncio_task for task in tasks_batch], return_exceptions=True
+        )
+        for result in results:
+            if isinstance(result, BaseException):
+                logger.error("Task has finished with an exception", exc_info=result)
+
+        partitions_to_offsets: typing.Final[dict[TopicPartition, int]] = {}
+        partition: TopicPartition
+        tasks: typing.Iterator[KafkaCommitTask]
+        for partition, tasks in partitions_to_tasks:
+            max_message_offset: int | None = None
+            for task in tasks:
+                if max_message_offset is None or task.offset > max_message_offset:
+                    max_message_offset = task.offset
+
+            if max_message_offset is not None:
+                partitions_to_offsets[partition] = max_message_offset + 1
+
+        commit_succeeded: typing.Final = await self._call_committer(tasks_batch, partitions_to_offsets)
+        for _ in tasks_batch:
+            self._messages_queue.task_done()
+        return commit_succeeded
+
+    async def _run_commit_process(self) -> None:
+        should_shutdown = False
+        while not should_shutdown:
+            commit_batch, should_shutdown = await self._populate_commit_batch()
+            if commit_batch:
+                await self._commit_tasks_batch(commit_batch)
+
+    async def commit_all(self) -> None:
+        """Commit all without shutting down the main process."""
+        self._flush_batch_event.set()
+        await self._messages_queue.join()
+
+    async def send_task(self, new_task: KafkaCommitTask) -> None:
+        self._check_is_commit_task_running()
+        await self._messages_queue.put(
+            new_task,
+        )
+
+    def spawn(self) -> None:
+        with self._spawn_lock:
+            if not self._asyncio_commit_process_task:
+                self._asyncio_commit_process_task = asyncio.create_task(self._run_commit_process())
+            else:
+                logger.error("Committer main task already running")
+
+    async def close(self) -> None:
+        """Close committer."""
+        if not self._asyncio_commit_process_task:
+            logger.error("Committer main task is not running, cannot close committer properly")
+            return
+
+        self._flush_batch_event.set()
+        try:
+            await asyncio.wait_for(self._asyncio_commit_process_task, timeout=self._shutdown_timeout)
+        except TimeoutError:
+            logger.exception("Committer main task shutdown timed out, forcing cancellation")
+            self._asyncio_commit_process_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._asyncio_commit_process_task
+        except Exception as exc:
+            logger.exception("Committer task failed during shutdown", exc_info=exc)
+            raise
+
+    @property
+    def is_healthy(self) -> bool:
+        return self._asyncio_commit_process_task is not None and not self._asyncio_commit_process_task.done()