Mercurial

comparison grok_interview/bucket_try_2.py @ 51:68fa88ac73fe

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Interview prep for xAI
author June Park <parkjune1995@gmail.com>
date Mon, 15 Dec 2025 19:55:17 -0800
parents
children
comparison
equal deleted inserted replaced
46:b9a40c633c93 51:68fa88ac73fe
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