I'd like to implement a multi-threaded token bucket. The producer will produce tokens for the container whereas the consumer will try to take the tokens from the container. If the number of tokens in the container is equal or greater than the tokens that the consumer asks for, then it will return True and take the tokens; otherwise, the request will be rejected and return False.
I'm trying to use the Semaphore to implement it but found it is a bit weird to access the number of tokens in the container using self.container._Semaphore__value. The reason I want to know the value of the semaphore is I want to make sure the consumer can get enough tokens from the bucket; otherwise, I'll reject it.
Maybe I have some misunderstanding of the whole thing. Can anyone point out a better way to implement a multi-threaded token bucket?
from threading import BoundedSemaphore, Thread, Lock
class TokenBucket(object):
def __init__(self, max_capacity, rate):
self.container = BoundedSemaphore(max_capacity)
self.rate = rate
self.lock = Lock()
def consume(self, n):
self.lock.acquire()
if n <= self.container._Semaphore__value:
for _ in range(n):
self.container.acquire()
return True
else:
return False
self.lock.release()
def produce(self):
self.lock.acquire()
for i in range(self.fill_rate):
try:
self.container.release()
except ValueError:
print("Full, skipping.")
self.lock.release()