|
51
|
1 # Context
|
|
|
2 #
|
|
|
3 # Please implement two functions: refillTokenBucket and useTokens. Two classes are already defined: DistributedCache and TokenBucket.
|
|
|
4 #
|
|
|
5 # refillTokenBucket(user_id): Refill tokens for the specified user's bucket.
|
|
|
6 #
|
|
|
7 # 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.
|
|
|
8
|
|
|
9 import asyncio
|
|
|
10 from typing import Dict
|
|
|
11
|
|
|
12
|
|
|
13 class DistributedCache:
|
|
|
14
|
|
|
15 def __init__(self):
|
|
|
16 self._hashmap: Dict[str, TokenBucket] = {}
|
|
|
17
|
|
|
18 def add_user(self, user_id: str):
|
|
|
19 self._hashmap[user_id] = TokenBucket()
|
|
|
20
|
|
|
21 def get_token_bucket(self, user_id: str):
|
|
|
22 return self._hashmap[user_id]
|
|
|
23
|
|
|
24 class TokenBucket:
|
|
|
25
|
|
|
26 def __init__(self, initial_token_value: int = 100):
|
|
|
27 self._tokens = initial_token_value
|
|
|
28 self.initial_token_value = initial_token_value
|
|
|
29 self._lock = asyncio.Lock()
|
|
|
30
|
|
|
31 async def consume(self, token: int):
|
|
|
32 async with self._lock:
|
|
|
33 tokens = self.get_token()
|
|
|
34 await asyncio.sleep(0.1)
|
|
|
35 if tokens < token:
|
|
|
36 return False
|
|
|
37 tokens -= token
|
|
|
38 self.set_token(tokens)
|
|
|
39 return True
|
|
|
40
|
|
|
41 async def refill(self):
|
|
|
42 async with self._lock:
|
|
|
43 self._tokens = self.initial_token_value
|
|
|
44
|
|
|
45 def get_token(self):
|
|
|
46 return self._tokens
|
|
|
47
|
|
|
48 def set_token(self, token: int):
|
|
|
49 self._tokens = token
|
|
|
50
|
|
|
51 distCache = DistributedCache()
|
|
|
52 distCache.add_user("june")
|
|
|
53
|
|
|
54 async def refillTokenBucket(user_id: str):
|
|
|
55 await distCache.get_token_bucket(user_id).refill()
|
|
|
56
|
|
|
57 async def useTokens(user_id: str, token: int):
|
|
|
58 return await distCache.get_token_bucket(user_id).consume(token)
|
|
|
59
|
|
|
60 async def main():
|
|
|
61 await asyncio.gather(useTokens("june", 10), useTokens("june", 10), useTokens("june", 10))
|
|
|
62 print(distCache.get_token_bucket("june").get_token())
|
|
|
63
|
|
|
64 asyncio.run(main())
|
|
|
65
|
