ratelimiter example

main.py

import boto3
import time

from botocore.session import Session as CoreSession
from botocore.hooks import HierarchicalEmitter

import traceback

import token_bucket


class ObservedEmitter(HierarchicalEmitter):

    StackTrack = False

    def emit(self, event_name, **kwargs):
        print "event_name", event_name
        if self.StackTrack:
            traceback.print_stack()
        return self._emit(event_name, kwargs)


def s3_usage(s):
    client = s.client('s3')
    client.list_buckets()


def ec2_usage(s, loop=100):
    client = s.client('ec2')

    for i in range(loop):
        client.describe_vpcs()


def iam_usage(s):
    client = s.client('iam')
    client.list_users()


class ServiceLimiter(object):

    def __init__(self):
        self.service_limits = {}

    def get_limit(self, service):
        if service not in self.service_limits:
            storage = token_bucket.MemoryStorage()
            self.service_limits[service] = s = token_bucket.Limiter(
                1, 1, storage)
            return s
        return self.service_limits[service]

    def enforce_rate_limit(self, **kw):
        _, svc, op = kw.get('event_name', '').split('.')
        limiter = self.get_limit(svc)
        while not limiter.consume(svc):
            print "limiting", svc
            time.sleep(1)
        return


def main():
    events = ObservedEmitter()
    limiter = ServiceLimiter()
    events.register_first('request-created.ec2', limiter.enforce_rate_limit)
    session = CoreSession(event_hooks=events)
    s = boto3.Session(botocore_session=session)

    #s3_usage(s)
    ec2_usage(s)
    #iam_usage(s)


if __name__ == '__main__':
    main()