lbry-sdk/lbrynet/core/RateLimiter.py

155 lines
4.3 KiB
Python
Raw Normal View History

2016-11-03 20:42:45 +01:00
import logging
2015-08-20 17:27:15 +02:00
from zope.interface import implements
from lbrynet.interfaces import IRateLimiter
from twisted.internet import task
2015-08-20 17:27:15 +02:00
2016-11-03 20:42:45 +01:00
log = logging.getLogger(__name__)
2015-08-20 17:27:15 +02:00
class DummyRateLimiter(object):
def __init__(self):
self.dl_bytes_this_second = 0
self.ul_bytes_this_second = 0
self.total_dl_bytes = 0
self.total_ul_bytes = 0
self.target_dl = 0
self.target_ul = 0
self.tick_call = None
2015-08-20 17:27:15 +02:00
def start(self):
self.tick_call = task.LoopingCall(self.tick)
self.tick_call.start(1)
2015-08-20 17:27:15 +02:00
def tick(self):
2015-08-20 17:27:15 +02:00
self.dl_bytes_this_second = 0
self.ul_bytes_this_second = 0
def stop(self):
if self.tick_call is not None:
self.tick_call.stop()
self.tick_call = None
2015-08-20 17:27:15 +02:00
def set_dl_limit(self, limit):
pass
def set_ul_limit(self, limit):
pass
def report_dl_bytes(self, num_bytes):
self.dl_bytes_this_second += num_bytes
self.total_dl_bytes += num_bytes
def report_ul_bytes(self, num_bytes):
self.ul_bytes_this_second += num_bytes
self.total_ul_bytes += num_bytes
class RateLimiter(object):
"""This class ensures that upload and download rates don't exceed specified maximums"""
implements(IRateLimiter)
#called by main application
def __init__(self, max_dl_bytes=None, max_ul_bytes=None):
self.max_dl_bytes = max_dl_bytes
self.max_ul_bytes = max_ul_bytes
self.dl_bytes_this_interval = 0
self.ul_bytes_this_interval = 0
2015-08-20 17:27:15 +02:00
self.total_dl_bytes = 0
self.total_ul_bytes = 0
self.tick_call = None
self.tick_interval = 0.1
2015-08-20 17:27:15 +02:00
self.dl_throttled = False
self.ul_throttled = False
self.protocols = []
def start(self):
2016-11-03 20:42:45 +01:00
log.info("Starting %s", self)
self.tick_call = task.LoopingCall(self.tick)
self.tick_call.start(self.tick_interval)
2015-08-20 17:27:15 +02:00
def tick(self):
self.dl_bytes_this_interval = 0
self.ul_bytes_this_interval = 0
2015-08-20 17:27:15 +02:00
self.unthrottle_dl()
self.unthrottle_ul()
def stop(self):
2016-11-03 20:42:45 +01:00
log.info("Stopping %s", self)
if self.tick_call is not None:
self.tick_call.stop()
self.tick_call = None
2015-08-20 17:27:15 +02:00
def set_dl_limit(self, limit):
self.max_dl_bytes = limit
def set_ul_limit(self, limit):
self.max_ul_bytes = limit
#throttling
def check_dl(self):
if self.max_dl_bytes is not None and self.dl_bytes_this_interval > self.max_dl_bytes * self.tick_interval:
from twisted.internet import reactor
reactor.callLater(0, self.throttle_dl)
2015-08-20 17:27:15 +02:00
def check_ul(self):
if self.max_ul_bytes is not None and self.ul_bytes_this_interval > self.max_ul_bytes * self.tick_interval:
from twisted.internet import reactor
reactor.callLater(0, self.throttle_ul)
2015-08-20 17:27:15 +02:00
def throttle_dl(self):
if self.dl_throttled is False:
for protocol in self.protocols:
protocol.throttle_download()
self.dl_throttled = True
def throttle_ul(self):
if self.ul_throttled is False:
for protocol in self.protocols:
protocol.throttle_upload()
self.ul_throttled = True
def unthrottle_dl(self):
if self.dl_throttled is True:
for protocol in self.protocols:
protocol.unthrottle_download()
self.dl_throttled = False
def unthrottle_ul(self):
if self.ul_throttled is True:
for protocol in self.protocols:
protocol.unthrottle_upload()
self.ul_throttled = False
#called by protocols
def report_dl_bytes(self, num_bytes):
self.dl_bytes_this_interval += num_bytes
2015-08-20 17:27:15 +02:00
self.total_dl_bytes += num_bytes
self.check_dl()
2015-08-20 17:27:15 +02:00
def report_ul_bytes(self, num_bytes):
self.ul_bytes_this_interval += num_bytes
2015-08-20 17:27:15 +02:00
self.total_ul_bytes += num_bytes
self.check_ul()
2015-08-20 17:27:15 +02:00
def register_protocol(self, protocol):
if protocol not in self.protocols:
self.protocols.append(protocol)
if self.dl_throttled is True:
protocol.throttle_download()
if self.ul_throttled is True:
protocol.throttle_upload()
def unregister_protocol(self, protocol):
if protocol in self.protocols:
2016-11-03 20:42:57 +01:00
self.protocols.remove(protocol)