84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
# (c) Copyright 2014,2015 Hewlett-Packard Development Company, L.P.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
|
|
import socket
|
|
import time
|
|
|
|
|
|
class ThrottledSocket(object):
|
|
def __init__(self, download_limit, upload_limit, wrappedsock):
|
|
self.__dict__['_wrappedsock'] = wrappedsock
|
|
self.__dict__['download_limit'] = download_limit
|
|
self.__dict__['upload_limit'] = upload_limit
|
|
|
|
def __getattr__(self, attr):
|
|
return getattr(self._wrappedsock, attr)
|
|
|
|
def __setattr__(self, attr, value):
|
|
return setattr(self._wrappedsock, attr, value)
|
|
|
|
def recv(self, *args):
|
|
start = time.time()
|
|
buf = self._wrappedsock.recv(*args)
|
|
self._sleep(len(buf), self.download_limit, start, time.time())
|
|
return buf
|
|
|
|
def sendall(self, *args):
|
|
start = time.time()
|
|
res = self._wrappedsock.send(*args)
|
|
self._sleep(res, self.upload_limit, start, time.time())
|
|
|
|
@staticmethod
|
|
def _sleep_duration(transmitted, limit, start_time, end_time):
|
|
exp_duration = float(transmitted) / limit
|
|
expected_end = start_time + exp_duration
|
|
return expected_end - end_time
|
|
|
|
@staticmethod
|
|
def _sleep(transmitted, limit, start_time, end_time):
|
|
if limit > 1:
|
|
sleep_duration = ThrottledSocket._sleep_duration(
|
|
transmitted,
|
|
limit,
|
|
start_time,
|
|
end_time)
|
|
if sleep_duration > 0:
|
|
time.sleep(sleep_duration)
|
|
|
|
def makefile(self, mode='r', bufsize=-1):
|
|
return socket._fileobject(self, mode, bufsize)
|
|
|
|
|
|
def monkeypatch_bandwidth(download_bytes_per_sec, upload_bytes_per_sec):
|
|
"""
|
|
Monkey patch socket to ensure that all
|
|
new sockets created are throttled.
|
|
"""
|
|
if upload_bytes_per_sec > -1 or download_bytes_per_sec > - 1:
|
|
def make_throttled_socket(*args, **kwargs):
|
|
return ThrottledSocket(
|
|
download_bytes_per_sec,
|
|
upload_bytes_per_sec,
|
|
socket._realsocket(*args, **kwargs))
|
|
|
|
socket.socket = make_throttled_socket
|
|
socket.SocketType = ThrottledSocket
|
|
|
|
|
|
def monkeypatch_socket_bandwidth(backup_opt_dict):
|
|
download_limit = backup_opt_dict.download_limit
|
|
upload_limit = backup_opt_dict.upload_limit
|
|
monkeypatch_bandwidth(download_limit, upload_limit)
|