| 303 | |
| 304 | |
| 305 | class ConsumptionScheduler: |
| 306 | def __init__(self): |
| 307 | """Schedules when to consume a desired amount""" |
| 308 | self._tokens_to_scheduled_consumption = {} |
| 309 | self._total_wait = 0 |
| 310 | |
| 311 | def is_scheduled(self, token): |
| 312 | """Indicates if a consumption request has been scheduled |
| 313 | |
| 314 | :type token: RequestToken |
| 315 | :param token: The token associated to the consumption |
| 316 | request that is used to identify the request. |
| 317 | """ |
| 318 | return token in self._tokens_to_scheduled_consumption |
| 319 | |
| 320 | def schedule_consumption(self, amt, token, time_to_consume): |
| 321 | """Schedules a wait time to be able to consume an amount |
| 322 | |
| 323 | :type amt: int |
| 324 | :param amt: The amount of bytes scheduled to be consumed |
| 325 | |
| 326 | :type token: RequestToken |
| 327 | :param token: The token associated to the consumption |
| 328 | request that is used to identify the request. |
| 329 | |
| 330 | :type time_to_consume: float |
| 331 | :param time_to_consume: The desired time it should take for that |
| 332 | specific request amount to be consumed in regardless of previously |
| 333 | scheduled consumption requests |
| 334 | |
| 335 | :rtype: float |
| 336 | :returns: The amount of time to wait for the specific request before |
| 337 | actually consuming the specified amount. |
| 338 | """ |
| 339 | self._total_wait += time_to_consume |
| 340 | self._tokens_to_scheduled_consumption[token] = { |
| 341 | 'wait_duration': self._total_wait, |
| 342 | 'time_to_consume': time_to_consume, |
| 343 | } |
| 344 | return self._total_wait |
| 345 | |
| 346 | def process_scheduled_consumption(self, token): |
| 347 | """Processes a scheduled consumption request that has completed |
| 348 | |
| 349 | :type token: RequestToken |
| 350 | :param token: The token associated to the consumption |
| 351 | request that is used to identify the request. |
| 352 | """ |
| 353 | scheduled_retry = self._tokens_to_scheduled_consumption.pop(token) |
| 354 | self._total_wait = max( |
| 355 | self._total_wait - scheduled_retry['time_to_consume'], 0 |
| 356 | ) |
| 357 | |
| 358 | |
| 359 | class BandwidthRateTracker: |