Repository: celery/celery
File: celery/worker/strategy.py (Python)

Code before:
    # Track the number of ETA tasks the worker is holding
    eta_task_limit = app.conf.worker_eta_task_limit
    # task event related
    # (optimized to avoid calling request.send_event)
    eventer = consumer.event_dispatcher
    events = eventer and eventer.enabled
    send_event = eventer and eventer.send
    task_sends_events = events and task.send_events

    call_at = consumer.timer.call_at
    apply_eta_task = consumer.apply_eta_task
    rate_limits_enabled = not consumer.disable_rate_limits
    get_bucket = consumer.task_buckets.__getitem__
    handle = consumer.on_task_request
    limit_task = consumer._limit_task
    limit_post_eta = consumer._limit_post_eta
    Request = symbol_by_name(task.Request)
    Req = create_request_cls(Request, task, consumer.pool, hostname, eventer, app=app)

    revoked_tasks = consumer.controller.state.revoked
    eta_task_limit = getattr(app.conf, 'worker_eta_task_limit', None)

    class ETATaskTracker:
        """Track ETA tasks to enforce limits."""
        def __init__(self):
            self.count = 0

        def check_limit(self, req):
            if eta_task_limit is not None and self.count >= eta_task_limit:
                info("ETA task limit reached (%s/%s), rejecting task %s",
                     self.count, eta_task_limit, req.id)
                req.reject(requeue=True)
                return True

            self.count += 1
            return False

        def create_callback(self, req, original_callback):
            def eta_callback(*args):
                self.count = max(0, self.count - 1)
                return original_callback(*args)
            eta_callback.__name__ = 'eta_callback'
            return eta_callback

    eta_tracker = ETATaskTracker()

    def task_message_handler(message, body, ack, reject, callbacks,
                             to_timestamp=to_timestamp):
        if body is None and 'args' not in message.payload:
            body, headers, decoded, utc = (
                message.body, message.headers, False, app.uses_utc_timezone(),

Review comment (suggestion):
Couldn't we use kombu's native prefetch count tracking for this , and cap the count there , instead of maintaining another count in celery?  I am not very familiar with celery's code base but I wonder if that would be a cleaner approach ?

[here](https://github.com/celery/kombu/blob/v5.5.4/kombu/common.py#L402) is where it is maintained in kombu.

Code after:
    # task event related
    # (optimized to avoid calling request.send_event)
    eventer = consumer.event_dispatcher
    events = eventer and eventer.enabled
    send_event = eventer and eventer.send
    task_sends_events = events and task.send_events

    call_at = consumer.timer.call_at
    apply_eta_task = consumer.apply_eta_task
    rate_limits_enabled = not consumer.disable_rate_limits
    get_bucket = consumer.task_buckets.__getitem__
    handle = consumer.on_task_request
    limit_task = consumer._limit_task
    limit_post_eta = consumer._limit_post_eta
    Request = symbol_by_name(task.Request)
    Req = create_request_cls(Request, task, consumer.pool, hostname, eventer,
                             app=app)

    revoked_tasks = consumer.controller.state.revoked

    def task_message_handler(message, body, ack, reject, callbacks,
                             to_timestamp=to_timestamp):
        if body is None and 'args' not in message.payload:
            body, headers, decoded, utc = (
                message.body, message.headers, False, app.uses_utc_timezone(),
            )
        else:
            if 'args' in message.payload:
                body, headers, decoded, utc = hybrid_to_proto2(message,
                                                               message.payload)
            else:
                body, headers, decoded, utc = proto1_to_proto2(message, body)

        req = Req(
            message,
            on_ack=ack, on_reject=reject, app=app, hostname=hostname,
            eventer=eventer, task=task, connection_errors=connection_errors,
            body=body, headers=headers, decoded=decoded, utc=utc,
        )
        if _does_info:
            # Similar to `app.trace.info()`, we pass the formatting args as the
            # `extra` kwarg for custom log handlers
            context = {
                'id': req.id,
                'name': req.name,
                'args': req.argsrepr,
                'kwargs': req.kwargsrepr,
                'eta': req.eta,
            }
            info(_app_trace.LOG_RECEIVED, context, extra={'data': context})
        if (req.expires or req.id in revoked_tasks) and req.revoked():