26 lines
1.0 KiB
Python
26 lines
1.0 KiB
Python
 
|
import asyncio,time
|
||
|
|
BASE,MAX,STEP,CALM,CAP=0.5,2.0,0.25,5.0,20
|
||
|
|
class State:
|
||
|
|
__slots__=("slot","interval","last","tasks")
|
||
|
|
def __init__(s):s.slot=0.0;s.interval=BASE;s.last=0.0;s.tasks=set()
|
||
|
|
class Limiter:
|
||
|
|
def __init__(s):s.users={}
|
||
|
|
def admit(s,uid):
|
||
|
|
now=time.monotonic();st=s.users.get(uid)
|
||
|
|
if st is None:st=s.users[uid]=State()
|
||
|
|
if len(st.tasks)>=CAP:return st,None,True
|
||
|
|
if now-st.last>CALM:st.interval=BASE
|
||
|
|
wait=max(0.0,st.slot-now)
|
||
|
|
if wait>0:st.interval=min(MAX,st.interval+STEP)
|
||
|
|
st.slot=max(now,st.slot)+st.interval;st.last=now
|
||
|
|
return st,wait,False
|
||
|
|
def flush(s,st):
|
||
|
|
n=len(st.tasks)
|
||
|
|
for t in list(st.tasks):t.cancel()
|
||
|
|
st.tasks.clear();st.slot=0.0;st.interval=BASE;return n
|
||
|
|
def track(s,uid,st,task):
|
||
|
|
st.tasks.add(task)
|
||
|
|
def done(t):
|
||
|
|
st.tasks.discard(t)
|
||
|
|
if not st.tasks and time.monotonic()-st.last>CALM*2 and s.users.get(uid) is st:s.users.pop(uid,None)
|
||
|
|
task.add_done_callback(done)
|