diff --git a/common.py b/common.py index 39a1a85..df6c78f 100644 --- a/common.py +++ b/common.py @@ -5,17 +5,20 @@ import os import pickle import re import sys +import time import urllib.request - +from collections import namedtuple +from threading import Lock from local_config import conf +RATE_NO_LIMIT = 0x00 RATE_GLOBAL = 0x01 RATE_NO_SILENCE = 0x02 RATE_INTERACTIVE = 0x04 RATE_CHAT = 0x08 RATE_URL = 0x10 - -rate_limit_classes = (RATE_CHAT, RATE_GLOBAL, RATE_NO_SILENCE, RATE_INTERACTIVE, RATE_URL) +RATE_EVENT = 0x20 +RATE_FUN = 0x40 BUFSIZ = 8192 EVENTLOOP_DELAY = 0.100 # seconds @@ -41,22 +44,98 @@ def conf_load(): return {} +Bucket = namedtuple("BucketConfig", ["history", "period", "max_hist_len"]) + +buckets = { + # everything else + RATE_GLOBAL: Bucket(history=[], period=60, max_hist_len=10), + # bot writes with no visible stimuli + RATE_NO_SILENCE: Bucket(history=[], period=10, max_hist_len=5), + # interactive stuff like ping + RATE_INTERACTIVE: Bucket(history=[], period=30, max_hist_len=5), + # chitty-chat, master volume control + RATE_CHAT: Bucket(history=[], period=10, max_hist_len=5), + # reacting on URLs + RATE_URL: Bucket(history=[], period=10, max_hist_len=5), + # triggering events + RATE_EVENT: Bucket(history=[], period=60, max_hist_len=10), + # bot blames people, produces cake and entertains + RATE_FUN: Bucket(history=[], period=180, max_hist_len=5), +} + +rate_limit_classes = buckets.keys() + + +def rate_limit(rate_class=RATE_GLOBAL): + """ + Remember N timestamps, + if N[0] newer than now()-T then do not output, do not append. + else pop(0); append() + + :param rate_class: the type of message to verify + :return: False if blocked, True if allowed + """ + if rate_class not in rate_limit_classes: + return all(rate_limit(c) for c in rate_limit_classes if c & rate_class) + + now = time.time() + bucket = buckets[rate_class] + logging.getLogger(__name__).debug("[ratelimit][bucket=%x][time=%s]%s" % (rate_class, now, bucket.history)) + + if len(bucket.history) >= bucket.max_hist_len and bucket.history[0] > (now - bucket.period): + # print("blocked") + return False + else: + if bucket.history and len(bucket.history) > bucket.max_hist_len: + bucket.history.pop(0) + bucket.history.append(now) + return True + + +def rate_limited(max_per_second): + """ + very simple flow control context manager + :param max_per_second: how many events per second may be executed - more are delayed + :return: + """ + min_interval = 1.0 / float(max_per_second) + + def decorate(func): + lasttimecalled = [0.0] + + def ratelimitedfunction(*args, **kargs): + elapsed = time.clock() - lasttimecalled[0] + lefttowait = min_interval - elapsed + if lefttowait > 0: + time.sleep(lefttowait) + ret = func(*args, **kargs) + lasttimecalled[0] = time.clock() + return ret + + return ratelimitedfunction + + return decorate + + def get_version_git(): import subprocess cmd = ['git', 'log', '--oneline', '--abbrev-commit'] - p = subprocess.Popen(cmd, bufsize=1, stdout=subprocess.PIPE) - first_line = p.stdout.readline() - line_count = len(p.stdout.readlines()) + 1 + try: + p = subprocess.Popen(cmd, bufsize=1, stdout=subprocess.PIPE) + first_line = p.stdout.readline() + line_count = len(p.stdout.readlines()) + 1 - if 0 == p.wait(): - # skip this 1st, 2nd, 3rd stuff and use always [0-9]th - return "version (Git, %dth rev) '%s'" % ( - line_count, str(first_line.strip(), encoding='utf8') - ) - else: - return "(unknown version)" + if 0 == p.wait(): + # skip this 1st, 2nd, 3rd stuff and use always [0-9]th + return "version (Git, %dth rev) '%s'" % ( + line_count, str(first_line.strip(), encoding='utf8') + ) + else: + return "(unknown version)" + except: + return "cannot determine version" VERSION = get_version_git() @@ -120,8 +199,7 @@ def extract_title(url): if result: match = result.groups()[0] - if not parser: - parser = html.parser.HTMLParser() + parser = html.parser.HTMLParser() try: expanded_html = parser.unescape(match) except UnicodeDecodeError as e: # idk why this can happen, but it does