| 59 | |
| 60 | |
| 61 | class BandwidthLimiter: |
| 62 | def __init__(self, leaky_bucket, time_utils=None): |
| 63 | """Limits bandwidth for shared S3 transfers |
| 64 | |
| 65 | :type leaky_bucket: LeakyBucket |
| 66 | :param leaky_bucket: The leaky bucket to use limit bandwidth |
| 67 | |
| 68 | :type time_utils: TimeUtils |
| 69 | :param time_utils: Time utility to use for interacting with time. |
| 70 | """ |
| 71 | self._leaky_bucket = leaky_bucket |
| 72 | self._time_utils = time_utils |
| 73 | if time_utils is None: |
| 74 | self._time_utils = TimeUtils() |
| 75 | |
| 76 | def get_bandwith_limited_stream( |
| 77 | self, fileobj, transfer_coordinator, enabled=True |
| 78 | ): |
| 79 | """Wraps a fileobj in a bandwidth limited stream wrapper |
| 80 | |
| 81 | :type fileobj: file-like obj |
| 82 | :param fileobj: The file-like obj to wrap |
| 83 | |
| 84 | :type transfer_coordinator: s3transfer.futures.TransferCoordinator |
| 85 | param transfer_coordinator: The coordinator for the general transfer |
| 86 | that the wrapped stream is a part of |
| 87 | |
| 88 | :type enabled: boolean |
| 89 | :param enabled: Whether bandwidth limiting should be enabled to start |
| 90 | """ |
| 91 | stream = BandwidthLimitedStream( |
| 92 | fileobj, self._leaky_bucket, transfer_coordinator, self._time_utils |
| 93 | ) |
| 94 | if not enabled: |
| 95 | stream.disable_bandwidth_limiting() |
| 96 | return stream |
| 97 | |
| 98 | |
| 99 | class BandwidthLimitedStream: |