Applying Backpressure When a Spatial Consumer Falls Behind

Pause the partition while continuing to poll, rather than slowing the handlers — a consumer that is busy instead of polling is presumed dead, so the naive fix trades a backlog for a rebalance that redelivers the batch to a member with identical settings.

This guide sits under Backpressure & Flow Control for Spatial Consumers, within Queue Management, Retries & Delivery Guarantees. That topic explains why work rather than message count is the right unit; this guide is the pause-and-resume mechanism built on it.

When to use this pattern

  • Consumer lag on one partition grows while others are idle, which for a geographically partitioned stream is the normal shape of trouble.
  • Rebalances happen under load, which is the signature of a consumer that stopped polling.
  • Payload cost varies by orders of magnitude, so no static prefetch is right.

Slowing down is not backpressure

The rebalance storm, and how it sustains itself A consumer meets a batch of heavy multipolygons and responds by processing them slowly, blocking inside its handler. Its next poll call therefore does not happen for ninety seconds, which exceeds the broker's maximum poll interval of thirty seconds. The broker presumes the member dead and triggers a rebalance, reassigning its partitions. The batch it was part-way through is uncommitted, so it is redelivered — to another member of the same group, running the same code with the same settings, which meets the same batch and blocks in the same way. That member is then evicted too, and the cycle continues: each rebalance costs the whole group a pause while assignments are recomputed, so total throughput falls while the backlog that caused the problem keeps growing. Nothing in this loop is a bug in the broker or in the batch; it is entirely a consequence of treating slow processing as a form of flow control. Pausing the partition instead keeps the poll loop running, so the member stays alive, keeps its assignment, and simply stops accepting new work until it has capacity. blocking in the handler · max.poll.interval.ms = 30 s handler busy 90 s poll deadline missed member evicted, rebalance batch redelivered …to a member with the same code and the same settings Each cycle pauses the whole group while assignments are recomputed Throughput falls while the backlog keeps growing. This is not a broker bug or a bad batch — it is treating slow processing as flow control. pausing the partition poll() still called, returns nothing session stays alive, assignment kept no new work accepted until there is capacity
Figure 1. The loop is self-sustaining: every eviction hands the same batch to a member that will be evicted by it. Nothing stops until the batch is smaller than the poll interval, which is what backpressure arranges.

Complete runnable implementation

python
import asyncio
from dataclasses import dataclass, field

from aiokafka import AIOKafkaConsumer, TopicPartition
from prometheus_client import Counter, Gauge

PAUSED = Gauge("consumer_partitions_paused", "Partitions currently paused")
BUDGET_USED = Gauge("consumer_vertex_budget_used", "Vertices in flight")
REFUNDS = Counter("consumer_budget_refunds_total", "Refunds", ("outcome",))

# Vertices, not messages. A partition can carry 40-vertex pings and
# 41 000-vertex multipolygons, and no count is right for both.
VERTEX_BUDGET = 250_000
# Resume only once the budget has genuinely recovered, not at the same line
# that triggered the pause — otherwise the partition oscillates.
RESUME_AT = 0.6


@dataclass(slots=True)
class WorkBudget:
    total: int
    in_flight: int = 0
    _waiters: list = field(default_factory=list)

    @property
    def exhausted(self) -> bool:
        return self.in_flight >= self.total

    @property
    def recovered(self) -> bool:
        return self.in_flight <= self.total * RESUME_AT

    def charge(self, vertices: int) -> None:
        self.in_flight += vertices
        BUDGET_USED.set(self.in_flight)

    def refund(self, vertices: int, outcome: str) -> None:
        # Refund on failure too. A handler that raises without refunding
        # leaks budget, and the partition stays paused forever — a stall
        # that looks exactly like a slow consumer.
        self.in_flight = max(0, self.in_flight - vertices)
        REFUNDS.labels(outcome=outcome).inc()
        BUDGET_USED.set(self.in_flight)


class BackpressuredConsumer:
    def __init__(self, consumer: AIOKafkaConsumer, process,
                 budget: int = VERTEX_BUDGET) -> None:
        self._consumer = consumer
        self._process = process
        self._budget = WorkBudget(total=budget)
        self._paused: set[TopicPartition] = set()

    async def run(self) -> None:
        while True:
            # Poll ALWAYS, even while paused. This is the call that keeps the
            # member alive; a paused partition simply returns no records.
            batch = await self._consumer.getmany(timeout_ms=500, max_records=50)

            for tp, records in batch.items():
                for record in records:
                    cost = _vertex_count(record)
                    self._budget.charge(cost)
                    asyncio.create_task(self._handle(record, cost))

            self._reconcile_pauses()

    async def _handle(self, record, cost: int) -> None:
        try:
            await self._process(record)
            self._budget.refund(cost, "ok")
        except Exception:
            self._budget.refund(cost, "failed")
            raise

    def _reconcile_pauses(self) -> None:
        assigned = self._consumer.assignment()

        if self._budget.exhausted:
            to_pause = assigned - self._paused
            if to_pause:
                self._consumer.pause(*to_pause)
                self._paused |= to_pause
        elif self._budget.recovered and self._paused:
            self._consumer.resume(*self._paused)
            self._paused.clear()

        PAUSED.set(len(self._paused))


def _vertex_count(record) -> int:
    """Read the precomputed count from the envelope.

    Deserialising the geometry to measure it defeats the purpose: the whole
    point is to decide before paying the parsing cost.
    """
    return int(record.headers_dict.get("vertex-count", 1000))
One threshold oscillates; two settle Budget occupancy is traced over time for a consumer under sustained heavy load. With a single threshold, the partition pauses the moment the budget is exhausted, one task completes, occupancy drops just below the line, the partition resumes, the next fetched multipolygon consumes the freed budget immediately, and the partition pauses again — several cycles a second, each one a broker round trip, and the pause metric becomes unreadable noise. With a resume band set at sixty per cent, the partition stays paused until enough work has genuinely completed, then resumes and takes a substantial batch before pausing again. The cycle period goes from milliseconds to seconds, the broker round trips fall by orders of magnitude, and the paused metric becomes something an operator can interpret. The cost is slightly lower average utilisation, because the consumer sometimes has spare budget it is not using — which is a good trade for a signal that means something and a broker that is not being asked to pause and resume a partition forty times a second. budget occupancy under sustained heavy load exhausted resume at 60% single threshold — pause/resume several times a second each cycle is a broker round trip; the paused metric is unreadable noise resume band — the cycle period goes from milliseconds to seconds The cost is slightly lower average utilisation — the consumer sometimes holds spare budget it is not using. In exchange the broker is not asked to pause and resume forty times a second, and the paused gauge becomes something an operator can read.
Figure 2. The oscillation is not merely inefficient — it destroys the one metric that tells an operator the consumer is under pressure.

Parameter reference

Name Type Spatial constraint Default
VERTEX_BUDGET int Vertices in flight, measured against resident memory per vertex 250000
RESUME_AT float Fraction of budget; equal to 1.0 means oscillation 0.6
vertex-count header int Precomputed at ingest — measuring it here defeats the purpose 1000
max_records int Small; the budget is the real bound, this just caps one fetch 50
Refund on failure Mandatory, or a raising handler leaks budget and stalls forever
poll while paused Required; it is what keeps the member in the group

Gotchas and spatial edge cases

  1. A handler that raises without refunding stalls the consumer permanently. The budget never recovers, the partition never resumes, and the symptom is a consumer that stopped consuming with no error after the first one. The try/except around the refund is the whole safety mechanism, and it must refund before re-raising.

  2. Pausing every assigned partition is coarse but usually right. The budget is a property of the process, so a heavy geometry on one partition legitimately blocks the others. Per-partition budgets are possible and rarely worth the complexity, because the memory that runs out is shared.

  3. A single message larger than the budget deadlocks. If one multipolygon has more vertices than the whole budget, it can never be charged, the partition pauses, and nothing ever completes to refund it. Either allow one over-budget message through when nothing is in flight, or route oversized geometry through the claim check in Streaming & Chunking Large Geometry Payloads.

  4. Vertex count is a proxy, not a cost. It correlates well with shapely runtime and poorly with anything involving a spatial join against a large index. If the handler’s cost is dominated by something else, charge that instead — the mechanism does not care what the unit is, only that it is knowable before dispatch.

  5. Pause state is lost across a rebalance. New assignments arrive unpaused, so the reconcile step must run against the current assignment rather than a remembered set — which is why _reconcile_pauses reads assignment() every cycle rather than caching it.

  6. Backpressure and retries push in opposite directions. A saturated consumer that starts timing out generates retries that arrive as additional load on the consumer that was already saturated. Gate retries on the same signal, as Tuning Retry Budgets for Webhook Provider SLAs describes.

A missing refund is a stall with no error after the first one Four heavy messages are dispatched and charged against the budget. Three complete and refund. The fourth raises — a corrupt geometry, a timeout, anything — and if the refund happens only on the success path, its charge stays in flight forever. The budget is now permanently short by that amount. A few more failures over the following hours and the in-flight total never falls below the pause threshold again, so the partition is paused indefinitely with no work actually running. What an operator sees is a consumer that stopped consuming, with one exception in the log from hours earlier and nothing since — because the consumer is not doing anything to fail at. The lag graph rises steadily and the error rate is zero. Refunding in a finally-style path before re-raising costs one line and turns the whole class of failure into an ordinary error the retry machinery already handles. budget in flight, with a refund only on the success path pause threshold a handler raises and its charge is never refunded in-flight never falls again · paused indefinitely, with nothing running What an operator sees: lag rising steadily, error rate zero, one exception hours earlier — because the consumer has nothing left to fail at.
Figure 3. The consumer is not broken in any way it can report. Refunding before re-raising costs one line and turns this into an ordinary retryable error.

Verification

python
import asyncio
import pytest


def test_budget_pauses_and_recovers_with_hysteresis():
    budget = WorkBudget(total=1000)
    budget.charge(1000)
    assert budget.exhausted and not budget.recovered

    budget.refund(100, "ok")            # 900 in flight — still above 60%
    assert not budget.recovered

    budget.refund(500, "ok")            # 400 in flight — below 60%
    assert budget.recovered


def test_failed_work_is_refunded():
    """The leak that stalls a consumer with no error after the first."""
    budget = WorkBudget(total=1000)

    async def failing(_record):
        raise ValueError("bad geometry")

    consumer = BackpressuredConsumer(_FakeKafka(), failing)
    consumer._budget = budget
    budget.charge(400)

    with pytest.raises(ValueError):
        asyncio.run(consumer._handle(_record(vertices=400), 400))
    assert budget.in_flight == 0


@pytest.mark.asyncio
async def test_poll_continues_while_paused():
    """The property that keeps the member in the group."""
    kafka = _FakeKafka(records=_heavy_batch(vertices=300_000))
    consumer = BackpressuredConsumer(kafka, _slow_process)

    task = asyncio.create_task(consumer.run())
    await asyncio.sleep(0.3)
    task.cancel()

    assert kafka.paused, "the partition should be paused"
    assert kafka.polls > 3, "poll must keep being called while paused"

The last test is the one that distinguishes this design from the naive one. A consumer that blocks in its handler also stops accepting work — the difference is entirely in whether poll keeps being called, and nothing else in the code makes that visible.