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
Complete runnable implementation
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))
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
-
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/exceptaround the refund is the whole safety mechanism, and it must refund before re-raising. -
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.
-
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.
-
Vertex count is a proxy, not a cost. It correlates well with
shapelyruntime 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. -
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_pausesreadsassignment()every cycle rather than caching it. -
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.
Verification
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.
Related
- Backpressure & Flow Control for Spatial Consumers — the topic this guide belongs to
- Sizing a Prefetch Window for Mixed Geometry Payloads — choosing the budget this mechanism enforces
- Shedding Spatial Load by Geographic Priority — what to do when pausing is no longer enough
- Consumer Lag & Partition Skew Monitoring — the signal that says a consumer is behind rather than merely busy