Mercurial
diff async_multi_threads/bucket_try_2.py @ 69:551d9fc0a2ba
Updated wrong names.
| author | June Park <parkjune1995@gmail.com> |
|---|---|
| date | Thu, 25 Dec 2025 20:07:46 -0800 |
| parents | |
| children |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/async_multi_threads/bucket_try_2.py Thu Dec 25 20:07:46 2025 -0800 @@ -0,0 +1,65 @@ +# Context +# +# Please implement two functions: refillTokenBucket and useTokens. Two classes are already defined: DistributedCache and TokenBucket. +# +# refillTokenBucket(user_id): Refill tokens for the specified user's bucket. +# +# useTokens(user_id, tokens): Check if there are enough tokens in the specified user's token bucket. If so, update the remaining token count and return true; otherwise, return false. You are required to use instances of the existing classes to implement these functions. + +import asyncio +from typing import Dict + + +class DistributedCache: + + def __init__(self): + self._hashmap: Dict[str, TokenBucket] = {} + + def add_user(self, user_id: str): + self._hashmap[user_id] = TokenBucket() + + def get_token_bucket(self, user_id: str): + return self._hashmap[user_id] + +class TokenBucket: + + def __init__(self, initial_token_value: int = 100): + self._tokens = initial_token_value + self.initial_token_value = initial_token_value + self._lock = asyncio.Lock() + + async def consume(self, token: int): + async with self._lock: + tokens = self.get_token() + await asyncio.sleep(0.1) + if tokens < token: + return False + tokens -= token + self.set_token(tokens) + return True + + async def refill(self): + async with self._lock: + self._tokens = self.initial_token_value + + def get_token(self): + return self._tokens + + def set_token(self, token: int): + self._tokens = token + +distCache = DistributedCache() +distCache.add_user("june") + +async def refillTokenBucket(user_id: str): + await distCache.get_token_bucket(user_id).refill() + +async def useTokens(user_id: str, token: int): + return await distCache.get_token_bucket(user_id).consume(token) + +async def main(): + await asyncio.gather(useTokens("june", 10), useTokens("june", 10), useTokens("june", 10)) + print(distCache.get_token_bucket("june").get_token()) + +asyncio.run(main()) +