| 210 | |
| 211 | |
| 212 | class LeakyBucket: |
| 213 | def __init__( |
| 214 | self, |
| 215 | max_rate, |
| 216 | time_utils=None, |
| 217 | rate_tracker=None, |
| 218 | consumption_scheduler=None, |
| 219 | ): |
| 220 | """A leaky bucket abstraction to limit bandwidth consumption |
| 221 | |
| 222 | :type rate: int |
| 223 | :type rate: The maximum rate to allow. This rate is in terms of |
| 224 | bytes per second. |
| 225 | |
| 226 | :type time_utils: TimeUtils |
| 227 | :param time_utils: The time utility to use for interacting with time |
| 228 | |
| 229 | :type rate_tracker: BandwidthRateTracker |
| 230 | :param rate_tracker: Tracks bandwidth consumption |
| 231 | |
| 232 | :type consumption_scheduler: ConsumptionScheduler |
| 233 | :param consumption_scheduler: Schedules consumption retries when |
| 234 | necessary |
| 235 | """ |
| 236 | self._max_rate = float(max_rate) |
| 237 | self._time_utils = time_utils |
| 238 | if time_utils is None: |
| 239 | self._time_utils = TimeUtils() |
| 240 | self._lock = threading.Lock() |
| 241 | self._rate_tracker = rate_tracker |
| 242 | if rate_tracker is None: |
| 243 | self._rate_tracker = BandwidthRateTracker() |
| 244 | self._consumption_scheduler = consumption_scheduler |
| 245 | if consumption_scheduler is None: |
| 246 | self._consumption_scheduler = ConsumptionScheduler() |
| 247 | |
| 248 | def consume(self, amt, request_token): |
| 249 | """Consume an a requested amount |
| 250 | |
| 251 | :type amt: int |
| 252 | :param amt: The amount of bytes to request to consume |
| 253 | |
| 254 | :type request_token: RequestToken |
| 255 | :param request_token: The token associated to the consumption |
| 256 | request that is used to identify the request. So if a |
| 257 | RequestExceededException is raised the token should be used |
| 258 | in subsequent retry consume() request. |
| 259 | |
| 260 | :raises RequestExceededException: If the consumption amount would |
| 261 | exceed the maximum allocated bandwidth |
| 262 | |
| 263 | :rtype: int |
| 264 | :returns: The amount consumed |
| 265 | """ |
| 266 | with self._lock: |
| 267 | time_now = self._time_utils.time() |
| 268 | if self._consumption_scheduler.is_scheduled(request_token): |
| 269 | return self._release_requested_amt_for_scheduled_request( |