Task concurrency / limiter question

[lembregtse]

What would be the best approach to do the following. We have a task in taskiq that is used to fetch data from a data source. Now per data source we want to limit the simultaneous tasks to be for example "2". No more than 2 tasks should be fetching from the data source at the same time. Workers are spread across different servers.

In our previous setup we had a decorator function that would randomly delay the task if a Redis "increment" "decrement" value was at the max concurrent amount.

I suppose the best way would be to implement middleware with the "pre_execute" function and the "post_execute" / "on_error" functions?

I was wondering if there is a better approach that might be more integrated with the broker? Waiting should not be blocking since we don't want to hold-up the different worker processes. The unique key to know the "per data-source" value could be a label that we set when kiqing the task.

We are using Redis and RabbitMQ as result and broker.

[s3rius]

Hello and thanks for asking this question.

You can actually create a dependency, that waits for redis-based semaphore for release.

As an example:

from taskiq import Context, TaskiqDepends

class ConcurencyLimiter:
	def __init__(self, counter_name: Optional[str] = None, limit: int = 2):
		self.limit = limit
		self.counter_name = counter_name

	async def __call__(self, redis_pool: ConnectionPool = Depends(get_redis_pool), context: Context = TaskiqDepends()) -> AsyncGenerator[None, None]:
		counter_name = self.counter_name or f"{context.message.task_name}.RATELIMIT"
		cur_value = self.limit + 1
		while cur_val >= self.limit:
			# here you just get current value from redis
			# and check that it's available.
			...
		# You increase your counter before task execution.
		redis.incr(counter_name)
		
		yield

		# You decrease it after the function is complete.
		redis.decr(counter_name)

In your tasks, you can easily depend on this limiter.

async def my_limited_task(limiter: None = TaskiqDepends(ConcurencyLimiter(4))):
	# do things.

I can make a full working example this evening, if you want.
Also, in future updates we're going to introduce error propagation to generator dependencies, so it's better to wrap yield with try.

try:
	yield
finally:
	redis.decr(counter_name)

[s3rius]

Currently, I don't see a cool solution that might cancel task execution completely from within middleware or dependency.

I've created an issue for that. #133.

[lembregtse]

Yes the proposed solution is similarly as to what we currently did with Celery. So would be great if we can achieve that with #133. I'll subscribe to the issue, thanks.

[s3rius]

Hi! Sory, for such a late reply. We are proud to announce that we have implemented all primitives to solve your specific case.

Here is an example project with implemented limiter. You can adopt it and start using.

mweq.zip

It uses new requeue method of a context. Also, I disabled logger of taskiq-dependencies, because throwing an exception in dependencies triggers resolver to catch them and print. logging.getLogger("taskiq.dependencies.ctx").setLevel(logging.ERROR)

[s3rius]

The only problem I found here is that it cannot be delayed in some cases. Of course, you can change delay labels of a message before calling requeue.

    async def __call__(
        self,
        redis: Redis = TaskiqDepends(get_redis_connection),
        context: Context = TaskiqDepends(),
    ) -> AsyncGenerator[None, None]:
        counter_name = self.counter_name or f"{context.message.task_name}.RATELIMIT"
        # You increase your counter before task execution.
        cur_val = int(await redis.incr(counter_name) or 0)

        if cur_val >= self.limit:
            await redis.decr(counter_name)
            context.message.labels["delay"] = "3"
            await context.requeue()

        yield

        # You decrease it after the function is complete.
        await redis.decr(counter_name)

Here's fixed version of limiter. Now it's a bit more precise and has delays for messages.

I have to admit, that delay label is not supported by some brokers. But aio-pika can delay messages based on this label.

[s3rius]

@lembregtse, if you think that it works for you, please choose the reply that helped you the most and mark it as answer.

[lembregtse]

Thanks a lot!