1
0
mirror of https://github.com/bspeice/metrik synced 2025-07-03 06:45:07 -04:00

Add initial rate-limit functionality

Likely needs more tests, but that's all I'm getting done tonight.
This commit is contained in:
Bradlee Speice
2016-08-23 22:23:35 -04:00
parent d1d58a1bd7
commit 4d36403c59
3 changed files with 152 additions and 6 deletions

View File

@ -1,12 +1,16 @@
from __future__ import print_function
import logging
import datetime
from time import sleep
from luigi import Task
from luigi.parameter import DateMinuteParameter, BoolParameter
from pymongo import MongoClient
from metrik.targets.mongo import MongoTarget
from metrik.targets.noop import NoOpTarget
from metrik.conf import MONGO_HOST, MONGO_PORT, MONGO_DATABASE
class MongoCreateTask(Task):
@ -61,3 +65,59 @@ class MongoNoBackCreateTask(MongoCreateTask):
# wish to persist for the future.
if self.live:
return super(MongoNoBackCreateTask, self).run()
class MongoRateLimit(object):
rate_limit_collection = 'rate_limit'
def __init__(self, service, limit, interval, max_tries=5, backoff=.5):
"""
:param present:
:type present: datetime.datetime
:param service:
:param limit:
:param interval:
:type interval: datetime.timedelta
:param max_tries:
:param backoff:
"""
self.service = service
self.limit = limit
self.interval = interval
self.max_tries = max_tries
self.backoff = backoff
self.db = MongoClient(host=MONGO_HOST, port=MONGO_PORT)[MONGO_DATABASE]
def get_present(self):
return datetime.datetime.now()
def query_locks(self, present):
return self.db[self.rate_limit_collection].find(
{'_created_at': {'$gt': present - self.interval},
'service': self.service}).count()
def save_lock(self, present):
self.db[self.rate_limit_collection].save({
'_created_at': present, 'service': self.service
})
def sleep_until(self, present):
future_time = present + self.interval * self.backoff
return (future_time - present).total_seconds()
def acquire_lock(self):
num_tries = 0
while num_tries < self.max_tries:
num_tries += 1
num_locks = self.query_locks(self.get_present())
if num_locks < self.limit:
self.save_lock(self.get_present())
return True
elif num_tries < self.max_tries:
sleep_amount = self.sleep_until(self.get_present())
sleep(sleep_amount)
return False