mirror of https://github.com/zulip/zulip.git
67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
|
from __future__ import absolute_import
|
||
|
|
|
||
|
|
from zephyr.models import UserProfile, get_user_profile_by_id
|
||
|
|
from zephyr.lib.rate_limiter import redis_key, client, max_api_calls, max_api_window
|
||
|
|
|
||
|
|
from django.core.management.base import BaseCommand
|
||
|
|
from django.conf import settings
|
||
|
|
from optparse import make_option
|
||
|
|
|
||
|
|
import os, time, logging
|
||
|
|
|
||
|
|
class Command(BaseCommand):
|
||
|
|
help = """Checks redis to make sure our rate limiting system hasn't grown a bug and left redis with a bunch of data
|
||
|
|
|
||
|
|
Usage: ./manage.py [--trim] check_redis"""
|
||
|
|
|
||
|
|
option_list = BaseCommand.option_list + (
|
||
|
|
make_option('-t', '--trim',
|
||
|
|
dest='trim',
|
||
|
|
default=False,
|
||
|
|
action='store_true',
|
||
|
|
help="Actually trim excess"),
|
||
|
|
)
|
||
|
|
|
||
|
|
def _check_within_range(self, key, count_func, trim_func):
|
||
|
|
user_id = int(key.split(':')[1])
|
||
|
|
try:
|
||
|
|
user = get_user_profile_by_id(user_id)
|
||
|
|
except:
|
||
|
|
user = None
|
||
|
|
max_calls = max_api_calls(user=user)
|
||
|
|
|
||
|
|
count = count_func()
|
||
|
|
if count > max_calls:
|
||
|
|
logging.error("Redis health check found key with more elements \
|
||
|
|
than max_api_calls! (trying to trim) %s %s" % (key, count))
|
||
|
|
client.expire(key, max_api_window(user=user))
|
||
|
|
if self.trim:
|
||
|
|
trim_func(key, max_calls)
|
||
|
|
|
||
|
|
def handle(self, *args, **options):
|
||
|
|
if not settings.RATE_LIMITING:
|
||
|
|
print "This machine is not using redis or rate limiting, aborting"
|
||
|
|
exit(1)
|
||
|
|
|
||
|
|
# Find all keys, and make sure they're all within size constraints
|
||
|
|
wildcard_list = "ratelimit:*:*:list"
|
||
|
|
wildcard_zset = "ratelimit:*:*:zset"
|
||
|
|
|
||
|
|
self.trim = options['trim']
|
||
|
|
|
||
|
|
lists = client.keys(wildcard_list)
|
||
|
|
for list_name in lists:
|
||
|
|
self._check_within_range(list_name,
|
||
|
|
lambda: client.llen(list_name),
|
||
|
|
lambda key, max_calls: client.ltrim(key, 0, max_calls - 1))
|
||
|
|
|
||
|
|
zsets = client.keys(wildcard_zset)
|
||
|
|
for zset in zsets:
|
||
|
|
now = time.time()
|
||
|
|
# We can warn on our zset being too large, but we don't know what
|
||
|
|
# elements to trim. We'd have to go through every list item and take
|
||
|
|
# the intersection. The best we can do is expire it
|
||
|
|
self._check_within_range(zset,
|
||
|
|
lambda: client.zcount(zset, 0, now),
|
||
|
|
lambda key, max_calls: None)
|