Python stuffs

tokenbucket.py

class TokenBucket:
    def __init__(self, bucketname, hostname, portnumber):
        self.bucketname = bucketname
        self.capacity = 1000 # max tokens in bucket
        self.windowTimeMin = 10 # how many mins to populate bucket with max tokens
        self.refillTokenPerMin = self.capacity / self.windowTimeMin # e.g. cap = 1000 tokens, window = 10 mins, then add 100 tokens/min
        self.memcachedClient = Client((hostname, portnumber))
        self.tokenBalance = 0 # initial value

    def getLastRefill(self):
        result = int(float(self.memcachedClient.get('last_refill_' + self.bucketname)))
        return result

    def setLastRefill(self):
        result = self.memcachedClient.set('last_refill_' + self.bucketname, int(datetime.now().timestamp()))
        return result

    def getTokens(self):
        result = int(float(self.memcachedClient.get('tokens_' + self.bucketname)))
        return result

    def setTokens(self, tokens):
        result = self.memcachedClient.set('tokens_' + self.bucketname, tokens)
        return result
    
    def refill(self):
        result = True
        lastRefillTimestamp = self.getLastRefill()
        if not lastRefillTimestamp:
            setTokens = self.capacity # bucket not exists, provide max tokens instantly
            self.memcachedClient.set('tokens_' + self.bucketname, setTokens)
        else:
            refillTokens = ( int(datetime.now().timestamp()) - lastRefillTimestamp ) / 60 * self.refillTokenPerMin
            if refillTokens > self.capacity:
                refillTokens = self.capacity
            setTokens = self.getTokens() + refillTokens
        self.setLastRefill()
        self.setTokens(setTokens)
        return result

    def consumeTokens(self, tokenCount):
        self.refill()
        balance = self.getTokens()
        if balance >= tokenCount: # Still enough balance
            balance -= tokenCount
            self.setTokens(balance)
        else: # not enough balance, empty the bucket
            result = False
            balance = 0
            self.setTokens(balance)
        return balance