Skip to content

Middleware

Middleware wraps every task a worker runs. Lifecycle hooks are events: they run once when a worker or scheduler starts or stops. A middleware layer is a wrapper: it runs around each task, sees the message before the task runs and the result after it, and can change or veto the outcome. Use it for timing, metrics, tracing, vetoes and per-task cleanup.

Middleware runs in workers only; a scheduler runs no tasks.

A minimal layer

A layer is an async def that takes the call and next, and returns a TaskResult:

async def layer(call: TaskCall, next: Next) -> TaskResult: ...

call.envelope is the message being handled and call.delivery the delivery it arrived on. await next(call) runs the rest of the chain and the task, and returns its result. Register a layer with @app.middleware:

"""A middleware layer that logs each task before and after it runs."""

import asyncio

from asynk import App, MemoryStore, Next, TaskCall, TaskResult, Worker
from asynk.transport.memory import memory_transport

producer, consumer = memory_transport()
app = App("minimal", producer=producer, consumer=consumer, store=MemoryStore())


@app.middleware
async def log_tasks(call: TaskCall, next: Next) -> TaskResult:
    print(f"starting {call.envelope.task} (attempt {call.envelope.attempt})")
    result = await next(call)
    print(f"{call.envelope.task} ended with {result.status.value}")
    return result


@app.task()
async def greet(name: str) -> str:
    return f"hello, {name}"


async def main() -> None:
    worker = Worker(app)
    run = asyncio.create_task(worker.run())
    handle = await greet.delay("asynk")
    print("result:", await handle.result(timeout=5))
    worker.request_shutdown()
    await run


if __name__ == "__main__":
    asyncio.run(main())

The layer usually returns the result it got from next. It may also return another one, for example a TaskResult with status SUCCESS and a fallback value, which the worker then records as the task's result.

Ordering

The last registered layer is the outermost: it runs first on the way in and last on the way out. The first registered layer sits closest to the task.

"""The last registered middleware layer is the outermost."""

import asyncio

from asynk import App, MemoryStore, Next, TaskCall, TaskResult, Worker
from asynk.transport.memory import memory_transport

producer, consumer = memory_transport()
app = App("order", producer=producer, consumer=consumer, store=MemoryStore())

order: list[str] = []


@app.middleware
async def first(call: TaskCall, next: Next) -> TaskResult:
    order.append("first: in")
    result = await next(call)
    order.append("first: out")
    return result


@app.middleware
async def second(call: TaskCall, next: Next) -> TaskResult:
    order.append("second: in")
    result = await next(call)
    order.append("second: out")
    return result


@app.task()
async def work() -> None:
    order.append("task")


async def main() -> None:
    worker = Worker(app)
    run = asyncio.create_task(worker.run())
    handle = await work.delay()
    await handle.result(timeout=5)
    worker.request_shutdown()
    await run
    print("\n".join(order))


if __name__ == "__main__":
    asyncio.run(main())

This prints:

second: in
first: in
task
first: out
second: out

The chain is built for each delivery, so a layer registered after the worker was created still applies. Registering the same layer twice has no effect.

Vetoing a task

To keep a task from running, return call.skip() without calling next. The task is recorded as SKIPPED, result() returns None, and the delivery is settled like a success, so the message is not redelivered.

"""A middleware layer that vetoes paused tasks with `call.skip()`."""

import asyncio

from asynk import App, MemoryStore, Next, TaskCall, TaskResult, Worker
from asynk.transport.memory import memory_transport

producer, consumer = memory_transport()
app = App("veto", producer=producer, consumer=consumer, store=MemoryStore())

paused = {"send_newsletter"}


@app.middleware
async def skip_paused(call: TaskCall, next: Next) -> TaskResult:
    if call.envelope.task in paused:
        return call.skip()  # the task does not run
    return await next(call)


@app.task(name="send_newsletter")
async def send_newsletter() -> str:
    return "sent"


async def main() -> None:
    worker = Worker(app)
    run = asyncio.create_task(worker.run())
    handle = await send_newsletter.delay()
    print("result:", await handle.result(timeout=5))
    print("state:", (await handle.state()).value)
    worker.request_shutdown()
    await run


if __name__ == "__main__":
    asyncio.run(main())

Cleaning up

Put cleanup in a try/finally around next, so it runs whatever happens inside:

@app.middleware
async def with_connection(call: TaskCall, next: Next) -> TaskResult:
    connection = await pool.acquire()
    try:
        return await next(call)
    finally:
        await pool.release(connection)

This includes shutdown. When a task is still running at the end of shutdown_timeout, the worker cancels it; the cancellation unwinds every layer, running their finally blocks, and the delivery is rejected with requeue so the broker redelivers it. Let asyncio.CancelledError propagate: do not swallow it in a layer.

When a layer raises

What happens depends on whether the task already ran.

  • Before the task ran (the layer raised before next, or a layer further in did): the error is logged and goes through the task's retry policy, like a failing task (see Retries). With retries left the result is RETRYING: the worker publishes envelope.retry() and acks the original delivery. With retries used up it is RETRIES_EXCEEDED, and with no policy, or an error not in retry_on, it is FAILED.
  • After the task ran (the layer raised after next returned): the error is logged and the task's own outcome stands. It is recorded and settled as if the layer had returned it, and the layer's error causes no redelivery.

With max_retries=None, a layer that always raises before the task retries forever, waiting the policy's backoff between attempts.

Settling the delivery yourself

A layer may settle the delivery itself with call.delivery.ack() or call.delivery.reject(). The worker sees that the delivery is settled and does not settle it again. For example, a layer can ack deliveries whose retries are used up, instead of letting the consumer dead-letter them:

@app.middleware
async def ack_exhausted(call: TaskCall, next: Next) -> TaskResult:
    result = await next(call)
    if result.status is TaskStatus.RETRIES_EXCEEDED:
        await call.delivery.ack()
    return result

The worker's own retry happens after the chain returns. A layer that settles the delivery of a RETRYING result takes over the retry: the worker publishes nothing, and the state stays RETRYING until another attempt records a new one.

A layer that settles the delivery and then raises is only logged: the message is not redelivered.

Pitfalls

  • Forgetting next skips the task. A layer that returns without awaiting next(call) vetoes the task, whether it meant to or not.
  • Calling next twice runs the task twice. Each call runs the rest of the chain and the task again.
  • Sync layers fail type-checking only. asynk does not check a layer when it is registered; mypy reports a plain def. At runtime, a plain def that returns a TaskResult raises TypeError when the worker awaits it, which is handled like a layer that raised before the task ran.
  • Layers registered later wrap earlier ones. A layer meant to see everything, such as metrics, must be registered last, after the layers that veto or replace results.
  • Unknown tasks reach the layers. A message naming a task the app has not registered still runs through the chain, and next returns a result with status UNKNOWN_TASK. call.envelope.task is then whatever name the producer sent.

A message that cannot be decoded into an envelope runs no layer: it is logged and rejected without requeue.

Timing and counting tasks (Prometheus)

A metrics layer times each task with time.perf_counter() and counts runs per task and status. Here the numbers go into a small class, and a veto layer shows that the metrics layer, registered last, also sees skipped tasks:

"""Timing and counting tasks with a middleware layer, using only the standard library."""

import asyncio
import contextlib
import time
from collections import defaultdict
from typing import Any

from asynk import App, MemoryStore, Next, TaskCall, TaskError, TaskHandle, TaskResult, Worker
from asynk.transport.memory import memory_transport

producer, consumer = memory_transport()
app = App("metrics", producer=producer, consumer=consumer, store=MemoryStore())


class TaskMetrics:
    """Counts runs per task and status, and adds up the seconds each task took."""

    def __init__(self) -> None:
        self.counts: defaultdict[tuple[str, str], int] = defaultdict(int)
        self.seconds: defaultdict[str, float] = defaultdict(float)

    def observe(self, task: str, status: str, seconds: float) -> None:
        self.counts[task, status] += 1
        self.seconds[task] += seconds


metrics = TaskMetrics()
paused = {"cleanup"}


@app.middleware
async def skip_paused(call: TaskCall, next: Next) -> TaskResult:
    if call.envelope.task in paused:
        return call.skip()
    return await next(call)


# Registered last, so it is the outermost layer and also sees skipped tasks.
@app.middleware
async def measure(call: TaskCall, next: Next) -> TaskResult:
    started = time.perf_counter()
    result = await next(call)
    metrics.observe(call.envelope.task, result.status.value, time.perf_counter() - started)
    return result


@app.task(name="add")
async def add(a: int, b: int) -> int:
    return a + b


@app.task(name="divide")
async def divide(a: int, b: int) -> float:
    return a / b


@app.task(name="cleanup")
async def cleanup() -> None: ...


async def main() -> TaskMetrics:
    worker = Worker(app)
    run = asyncio.create_task(worker.run())
    handles: list[TaskHandle[Any]] = [
        await add.delay(1, 2),
        await add.delay(3, 4),
        await cleanup.delay(),
    ]
    for handle in handles:
        await handle.result(timeout=5)
    failing = await divide.delay(1, 0)
    with contextlib.suppress(TaskError):
        await failing.result(timeout=5)
    worker.request_shutdown()
    await run
    for (task, status), count in sorted(metrics.counts.items()):
        print(f"{task} {status}: {count}")
    for task, seconds in sorted(metrics.seconds.items()):
        print(f"{task}: {seconds:.6f}s in total")
    return metrics


if __name__ == "__main__":
    asyncio.run(main())

It prints:

add success: 2
cleanup skipped: 1
divide failed: 1

followed by the total seconds per task.

The same layer with prometheus_client exposes a counter labelled by task and status and a histogram of durations labelled by task, served on http://localhost:8000/metrics. It needs pip install prometheus-client:

import time

from prometheus_client import Counter, Histogram, start_http_server

from asynk import Next, Runner, TaskCall, TaskResult, TaskStatus

TASKS = Counter("asynk_tasks_total", "Tasks handled, by task and status.", ["task", "status"])
DURATION = Histogram("asynk_task_duration_seconds", "Task duration in seconds.", ["task"])


@app.before_startup
def serve_metrics(runner: Runner) -> None:
    if runner.role == "worker":
        start_http_server(8000)


@app.middleware
async def measure(call: TaskCall, next: Next) -> TaskResult:
    started = time.perf_counter()
    result = await next(call)
    task = "unknown" if result.status is TaskStatus.UNKNOWN_TASK else call.envelope.task
    TASKS.labels(task, result.status.value).inc()
    DURATION.labels(task).observe(time.perf_counter() - started)
    return result

Keep labels low-cardinality: every distinct label value is a new time series. Task name and status are bounded; task arguments or ids are not. The name of an unknown task comes from the message, so the layer above reports all unknown tasks under one unknown label.