Add output rate limiting

This commit is contained in:
dave 2017-11-24 16:06:55 -08:00
parent 2f88dc28c6
commit 3d047bb0cc
5 changed files with 97 additions and 10 deletions

View File

@ -40,7 +40,11 @@ Instance Configuration
],
"force_ipv6": false,
"force_ipv4": false,
"bind": ["1.2.3.4", 5678]
"bind": ["1.2.3.4", 5678],
"rate_limit": {
"rate_max": 5.0,
"rate_int":1.1
}
},
"modules":[
"PingResponder",
@ -98,6 +102,11 @@ options:
To bind to an address but no specific port, set the second tuple entry to `null`.
.. cmdoption:: connection.rate_limit
Set to false to disable rate limiting. Otherwise, a dict containing two floats keyed: `rate_max`: how many messages
may be bursted at once, and `rate_int`: after bursting, how many seconds between messages.
.. cmdoption:: modules
A list of modules to load. Modules are loaded in the order they are listed

View File

@ -13,7 +13,11 @@
["dickson.freenode.net", 6667],
["morgan.freenode.net", 6667]
],
"force_ipv6": false
"force_ipv6": false,
"rate_limit": {
"rate_max": 5.0,
"rate_int":1.1
}
},
"modules":[
"PingResponder",

41
pyircbot/common.py Normal file
View File

@ -0,0 +1,41 @@
from time import time
from math import floor
class burstbucket(object):
def __init__(self, maximum, interval):
"""
Burst bucket class for rate limiting
:param maximum: maximum value in the bucket
:param interval: how often a whole item is added to the bucket
"""
# How many messages can be bursted
self.bucket_max = maximum
# how often the bucket has 1 item added
self.bucket_period = interval
# last time the burst bucket was filled
self.bucket_lastfill = time()
self.bucket = self.bucket_max
def get(self):
"""
Return 0 if no sleeping is necessary to rate limit. Otherwise, return the number of seconds to sleep. This
method should be called again by the user after sleeping
"""
# First, update the bucket
# Check if $period time has passed since the bucket was filled
since_fill = time() - self.bucket_lastfill
if since_fill > self.bucket_period:
# How many complete points are credited
fills = floor(since_fill / self.bucket_period)
self.bucket += fills
if self.bucket > self.bucket_max:
self.bucket = self.bucket_max
# Advance the lastfill time appropriately
self.bucket_lastfill += self.bucket_period * fills
if self.bucket >= 1:
self.bucket -= 1
return 0
return self.bucket_period - since_fill

View File

@ -12,7 +12,7 @@ import logging
import traceback
import sys
from inspect import getargspec
from time import sleep
from pyircbot.common import burstbucket
from collections import namedtuple
from io import StringIO
@ -23,7 +23,13 @@ ServerPrefix = namedtuple("ServerPrefix", "hostname")
class IRCCore(object):
def __init__(self, servers):
def __init__(self, servers, loop, rate_limit=True, rate_max=5.0, rate_int=1.1):
self._loop = loop
# rate limiting options
self.rate_limit = rate_limit
self.rate_max = float(rate_max)
self.rate_int = float(rate_int)
self.connected = False
"""If we're connected or not"""
@ -53,6 +59,9 @@ class IRCCore(object):
# Set up hooks for modules
self.initHooks()
self.outputq = asyncio.Queue()
self._loop.call_soon(asyncio.ensure_future, self.outputqueue())
async def loop(self, loop):
while self.alive:
try:
@ -88,6 +97,26 @@ class IRCCore(object):
logging.info("Reconnecting in 3s...")
await asyncio.sleep(3)
async def outputqueue(self):
bucket = burstbucket(self.rate_max, self.rate_int)
while True:
prio, line = await self.outputq.get()
# sleep until the bucket allows us to send
if self.rate_limit:
while True:
s = bucket.get()
if s == 0:
break
else:
await asyncio.sleep(s, loop=self._loop)
self.fire_hook('_SEND', args=None, prefix=None, trailing=None)
self.log.debug(">>> {}".format(repr(line)))
try:
self.writer.write((line + "\r\n").encode("UTF-8"))
except Exception as e: # Probably fine if we drop messages while offline
print(e)
print(self.trace())
async def kill(self, message="Help! Another thread is killing me :(", forever=True):
"""Send quit message, flush queue, and close the socket
@ -137,9 +166,7 @@ class IRCCore(object):
self.fire_hook(command, args=args, prefix=prefix, trailing=trailing)
def sendRaw(self, data):
self.log.debug(">>> {}".format(repr(data)))
self.fire_hook('_SEND', args=None, prefix=None, trailing=None)
self.writer.write((data + "\r\n").encode("UTF-8"))
asyncio.run_coroutine_threadsafe(self.outputq.put((5, data, )), self._loop)
" Module related code "
def initHooks(self):

View File

@ -31,6 +31,8 @@ class PyIRCBot(object):
self.log = logging.getLogger('PyIRCBot')
"""Reference to logger object"""
self.loop = asyncio.get_event_loop()
"""saved copy of the instance config"""
self.botconfig = botconfig
@ -43,8 +45,14 @@ class PyIRCBot(object):
"""Reference to BotRPC thread"""
self.rpc = BotRPC(self)
ratelimit = self.botconfig["connection"].get("rate_limit", None) or dict(rate_max=5.0, rate_int=1.1)
"""IRC protocol handler"""
self.irc = IRCCore(servers=self.botconfig["connection"]["servers"])
self.irc = IRCCore(servers=self.botconfig["connection"]["servers"],
loop=self.loop,
rate_limit=True if ratelimit else False,
rate_max=ratelimit["rate_max"],
rate_int=ratelimit["rate_int"])
if self.botconfig.get("connection").get("force_ipv6", False):
self.irc.connection_family = AF_INET6
elif self.botconfig.get("connection").get("force_ipv4", False):
@ -72,8 +80,6 @@ class PyIRCBot(object):
self.irc.addHook("PRIVMSG", self._irchook_internal)
def run(self):
self.loop = asyncio.get_event_loop()
self.client = asyncio.ensure_future(self.irc.loop(self.loop), loop=self.loop)
try:
self.loop.set_debug(True)