token-bucket-as-a-service.py

import time, contextlib, struct, json
import uwsgi

def token_bucket(interval, burst=1):
	token_fmt = '>dL'
	token_len = struct.calcsize(token_fmt)
	token_get = lambda: struct.unpack(token_fmt, uwsgi.sharedarea_read(0, token_len))
	token_set = lambda *v: uwsgi.sharedarea_write(0, struct.pack(token_fmt, *v))

	token_set(time.monotonic(), burst)
	assert uwsgi.sharedarea_read(0, token_len), 'Enable uwsgi --sharedarea plz'

	@contextlib.contextmanager
	def token_check():
		uwsgi.sharedarea_wlock()
		try: yield (token_get, token_set)
		finally: uwsgi.sharedarea_unlock()

	rate = interval**-1

	def get_delay(grab=1):
		ts = time.monotonic()
		with token_check() as (tget, tset):
			ts_sync, tokens = tget()
			ts_sync, tokens = ts, min(burst, tokens + (ts - ts_sync) * rate)
			if tokens >= grab: delay, tokens = 0, tokens - grab
			else: delay = (grab - tokens) / rate
			tset(ts_sync, tokens)
		return delay

	return get_delay


# interval, burst = uwsgi.opt[...]
interval, burst = 3600, 10 # hardcoded 10 req/h
get_delay = token_bucket(interval, burst)

def application(env, start_response):
	start_response('200 OK', [('Content-Type','application/json')])
	return [json.dumps(dict(delay=get_delay()))]