| 43 | } |
| 44 | |
| 45 | pub fn is_allowed(&self, client_ip: IpAddr) -> bool { |
| 46 | let now = Instant::now(); |
| 47 | let mut clients = self.clients.lock(); |
| 48 | |
| 49 | if let Some(state) = clients.get_mut(&client_ip) { |
| 50 | let elapsed = now.duration_since(state.last_update); |
| 51 | let elapsed_ms = elapsed.as_millis(); |
| 52 | let refill = elapsed_ms.saturating_mul(self.refill_rate); |
| 53 | state.millitokens = state |
| 54 | .millitokens |
| 55 | .saturating_add(refill) |
| 56 | .min(self.max_millitokens); |
| 57 | state.last_update = now; |
| 58 | |
| 59 | if state.millitokens >= MILLITOKENS_PER_TOKEN { |
| 60 | state.millitokens -= MILLITOKENS_PER_TOKEN; |
| 61 | true |
| 62 | } else { |
| 63 | false |
| 64 | } |
| 65 | } else { |
| 66 | let state = ClientState { |
| 67 | millitokens: self.max_millitokens.saturating_sub(MILLITOKENS_PER_TOKEN), |
| 68 | last_update: now, |
| 69 | }; |
| 70 | clients.insert(client_ip, state); |
| 71 | true |
| 72 | } |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | pub type SharedRateLimiter = Option<Arc<RateLimiter>>; |