Example #1
0
def rate_limit(state, task_name, rate_limit, **kwargs):
    """Tell worker(s) to modify the rate limit for a task by type.

    See Also:
        :attr:`celery.task.base.Task.rate_limit`.

    Arguments:
        task_name (str): Type of task to set rate limit for.
        rate_limit (int, str): New rate limit.
    """
    try:
        rate(rate_limit)
    except ValueError as exc:
        return nok('Invalid rate limit string: {0!r}'.format(exc))

    try:
        state.app.tasks[task_name].rate_limit = rate_limit
    except KeyError:
        logger.error('Rate limit attempt for unknown task %s',
                     task_name,
                     exc_info=True)
        return nok('unknown task')

    state.consumer.reset_rate_limits()

    if not rate_limit:
        logger.info('Rate limits disabled for tasks of type %s', task_name)
        return ok('rate limit disabled successfully')

    logger.info('New rate limit for tasks of type %s: %s.', task_name,
                rate_limit)
    return ok('new rate limit set successfully')
Example #2
0
def rate_limit(state, task_name, rate_limit, **kwargs):
    """Tell worker(s) to modify the rate limit for a task by type.

    See Also:
        :attr:`celery.task.base.Task.rate_limit`.

    Arguments:
        task_name (str): Type of task to set rate limit for.
        rate_limit (int, str): New rate limit.
    """
    # pylint: disable=redefined-outer-name
    # XXX Note that this redefines `terminate`:
    #     Outside of this scope that is a function.
    try:
        rate(rate_limit)
    except ValueError as exc:
        return nok('Invalid rate limit string: {0!r}'.format(exc))

    try:
        state.app.tasks[task_name].rate_limit = rate_limit
    except KeyError:
        logger.error('Rate limit attempt for unknown task %s',
                     task_name, exc_info=True)
        return nok('unknown task')

    state.consumer.reset_rate_limits()

    if not rate_limit:
        logger.info('Rate limits disabled for tasks of type %s', task_name)
        return ok('rate limit disabled successfully')

    logger.info('New rate limit for tasks of type %s: %s.',
                task_name, rate_limit)
    return ok('new rate limit set successfully')
def rate_limit(state, task_name, rate_limit, **kwargs):
    """Tell worker(s) to modify the rate limit for a task by type.

    See Also:
        :attr:`celery.task.base.Task.rate_limit`.

    Arguments:
        task_name (str): Type of task to set rate limit for.
        rate_limit (int, str): New rate limit.
    """
    # pylint: disable=redefined-outer-name
    # XXX Note that this redefines `terminate`:
    #     Outside of this scope that is a function.
    try:
        rate(rate_limit)
    except ValueError as exc:
        return nok(f'Invalid rate limit string: {exc!r}')

    try:
        state.app.tasks[task_name].rate_limit = rate_limit
    except KeyError:
        logger.error('Rate limit attempt for unknown task %s',
                     task_name,
                     exc_info=True)
        return nok('unknown task')

    state.consumer.reset_rate_limits()

    if not rate_limit:
        logger.info('Rate limits disabled for tasks of type %s', task_name)
        return ok('rate limit disabled successfully')

    logger.info('New rate limit for tasks of type %s: %s.', task_name,
                rate_limit)
    return ok('new rate limit set successfully')
Example #4
0
    def _context(self,
                 sig,
                 rate_limits=True,
                 events=True,
                 utc=True,
                 limit=None):
        assert sig.type.Strategy
        assert sig.type.Request

        reserved = Mock()
        consumer = Mock()
        consumer.task_buckets = defaultdict(lambda: None)
        if limit:
            bucket = TokenBucket(rate(limit), capacity=1)
            consumer.task_buckets[sig.task] = bucket
        consumer.controller.state.revoked = set()
        consumer.disable_rate_limits = not rate_limits
        consumer.event_dispatcher.enabled = events
        s = sig.type.start_strategy(self.app, consumer, task_reserved=reserved)
        assert s

        message = self.task_message_from_sig(
            self.app,
            sig,
            utc=utc,
            TaskMessage=self.get_message_class(),
        )
        message = self.prepare_message(message)
        yield self.Context(sig, s, reserved, consumer, message)
 def __init__(self, state, freq=1.0, maxrate=None,
              cleanup_freq=3600.0, timer=None, app=None):
     self.app = app_or_default(app)
     self.state = state
     self.freq = freq
     self.cleanup_freq = cleanup_freq
     self.timer = timer or self.timer or Timer()
     self.logger = logger
     self.maxrate = maxrate and TokenBucket(rate(maxrate))
Example #6
0
    def test_conversion(self):
        self.assertEqual(rate(999), 999)
        self.assertEqual(rate(7.5), 7.5)
        self.assertEqual(rate('2.5/s'), 2.5)
        self.assertEqual(rate('1456/s'), 1456)
        self.assertEqual(rate('100/m'),
                         100 / 60.0)
        self.assertEqual(rate('10/h'),
                         10 / 60.0 / 60.0)

        for zero in (0, None, '0', '0/m', '0/h', '0/s', '0.0/s'):
            self.assertEqual(rate(zero), 0)
Example #7
0
 def __init__(self,
              state,
              freq=1.0,
              maxrate=None,
              cleanup_freq=3600.0,
              timer=None,
              app=None):
     self.app = app_or_default(app)
     self.state = state
     self.freq = freq
     self.cleanup_freq = cleanup_freq
     self.timer = timer or self.timer or Timer()
     self.logger = logger
     self.maxrate = maxrate and TokenBucket(rate(maxrate))
Example #8
0
    def _context(self, sig,
                 rate_limits=True, events=True, utc=True, limit=None):
        assert sig.type.Strategy

        reserved = Mock()
        consumer = Mock()
        consumer.task_buckets = defaultdict(lambda: None)
        if limit:
            bucket = TokenBucket(rate(limit), capacity=1)
            consumer.task_buckets[sig.task] = bucket
        consumer.controller.state.revoked = set()
        consumer.disable_rate_limits = not rate_limits
        consumer.event_dispatcher.enabled = events
        s = sig.type.start_strategy(self.app, consumer, task_reserved=reserved)
        assert s

        message = self.task_message_from_sig(
            self.app, sig, utc=utc, TaskMessage=self.get_message_class(),
        )
        message = self.prepare_message(message)
        yield self.Context(sig, s, reserved, consumer, message)
Example #9
0
 def bucket_for_task(self, type):
     limit = rate(getattr(type, 'rate_limit', None))
     return TokenBucket(limit, capacity=1) if limit else None
Example #10
0
 def bucket_for_task(self, type):
     limit = rate(getattr(type, 'rate_limit', None))
     return TokenBucket(limit, capacity=1) if limit else None
Example #11
0
def test_rate_limit_string(s, expected):
    assert rate(s) == expected
Example #12
0
def test_rate_limit_string(s, expected):
    assert rate(s) == expected