SmartRetryMiddleware: retry delay compounds without schedule_source #716

Taskiq version

0.13.0 (same on main at f08a54a)

Python version

Python 3.12

OS

macOS

What happened?

SmartRetryMiddleware sends a different delay on each retry depending on whether schedule_source is set.

Without schedule_source, on_send re-kicks the task with with_labels(delay=delay). That sets the message's delay label to the delay it just computed. On the next failure make_delay reads that label back as the base and works from it again. With schedule_source the label is left alone, so every retry starts from the original value.

Six failed attempts with default_delay=5 and the default max_delay_exponent=60. The broker rows are the delay label sent with each retry, and the schedule_source rows are the scheduled delay:

broker, exponent:        [5.0, 10.0, 30.0, 60.0, 60.0, 60.0]
schedule_source, exp.:   [5.0, 10.0, 15.0, 20.0, 25.0, 30.0]
broker, jitter:          [5.8, 6.6, 7.0, 7.3, 7.8, 8.2]
schedule_source, jitter: [5.8, 5.8, 5.4, 5.3, 5.5, 5.4]
broker, defaults:        [5.0, 5.0, 5.0, 5.0, 5.0, 5.0]

The label only changes the real wait on brokers that follow the delay convention, like taskiq-sqs. On those, use_delay_exponent hits the cap on the fourth retry. With use_jitter, each retry adds its random part on top of the last one, so the delay keeps creeping up. I get the same labels from a real InMemoryBroker worker loop, not just the harness below.

Neither path is exponential in the usual sense. With schedule_source each retry is base * n (5, 10, 15). Without it, each retry multiplies the previous result by the retry count (5, 10, 30, then 120 capped to 60), so it grows as base * n!. The problem is that the result depends on whether schedule_source is set. The docs describe the task's delay argument as "Initial delay before retrying the task, in seconds."

I'm happy to send a PR, but I wasn't sure which fix you'd want. The broker docs (docs/extending-taskiq/broker.md, Conventions) say a broker should delay execution by the delay label, so the retry still has to set it. One way is to keep the first delay in its own private label, the way _retries is kept, and compute from that. If the growth on the broker path is what you meant, the schedule_source path is the one to change instead.

Broker initialization code

Full repro, no services needed:

import asyncio
import datetime
import random

from taskiq import InMemoryBroker, ScheduleSource, SmartRetryMiddleware
from taskiq.message import TaskiqMessage
from taskiq.result import TaskiqResult


class Broker(InMemoryBroker):
    async def kick(self, message):
        self.last = message


class Source(ScheduleSource):
    async def get_schedules(self):
        return []

    async def add_schedule(self, schedule):
        self.last = schedule


async def delays(source=None, **kwargs):
    mw = SmartRetryMiddleware(
        default_retry_label=True, default_retry_count=8, default_delay=5,
        schedule_source=source, **kwargs,
    )
    broker = Broker().with_middlewares(mw)
    labels, out = {}, []
    for _ in range(6):
        now = datetime.datetime.now(datetime.timezone.utc)
        msg = TaskiqMessage(task_id="t", task_name="t", labels=labels, args=[], kwargs={})
        result = TaskiqResult(is_err=True, return_value=None, execution_time=0)
        await mw.on_error(msg, result, ValueError())
        if source is None:  # the next worker gets the re-kicked message
            labels = dict(broker.last.labels)
            out.append(round(float(labels["delay"]), 1))
        else:
            labels = dict(source.last.labels)
            out.append(round((source.last.time - now).total_seconds(), 1))
    return out


async def main():
    print("broker, exponent:      ", await delays(use_delay_exponent=True))
    print("schedule_source, exp.: ", await delays(Source(), use_delay_exponent=True))
    random.seed(0)
    print("broker, jitter:        ", await delays(use_jitter=True))
    random.seed(0)
    print("schedule_source, jitter:", await delays(Source(), use_jitter=True))
    print("broker, defaults:      ", await delays())


asyncio.run(main())