Queued rate limiter attempt

queued_rateLimiter.py

import time
from functools import wraps
import threading
import queue
from typing import List, Callable
import string
import random
import asyncio


class QueuedRateLimiter:
    """ Rate limit execution of tasks in a queue

    Example:
        max_exec_per_limit = 10 and window_size_in_seconds = 60 sec
        Only 10 executions is allowed during a 60 second window.
        When x execs is reached during n second window, exec is paused and resumed when allowed again.

    """

    def __init__(self, max_exec_per_limit: int = 10, window_size_in_seconds: int = 60):
        self.max_exec_per_limit = max_exec_per_limit
        self.window_size_in_seconds = window_size_in_seconds  # Seconds
        self.time_until_next = 0
        self.is_rate_limited = False
        self.time_bucket: List[float] = []
        self.func_queue = queue.Queue()  # Is thread safe! No need for locks and such

    def rate_limiter(self):
        """ Rate limiting
            Decides when to rate limit.
        """

        now = time.perf_counter()

        if self.is_rate_limited and now < self.time_until_next:
            # Ignore
            return
        elif self.is_rate_limited and now > self.time_until_next:
            self.is_rate_limited = False
            self.time_bucket.clear()

        if len(self.time_bucket) >= self.max_exec_per_limit:

            diff_first = now - self.time_bucket[0]
            print(f"{diff_first=}")

            if diff_first <= self.window_size_in_seconds:
                # Rate limit
                self.time_until_next = self.window_size_in_seconds
                self.is_rate_limited = True
                return
            else:
                # Clear time_bucket
                self.time_bucket.clear()

        self.time_bucket.append(now)

        if self.func_queue.empty():
            return

    def run(self) -> None:
        while True:
            while not self.func_queue.empty():
                self.rate_limiter()
                print("Looking for a job")
                
                if self.is_rate_limited:
                    print(f"RATE LIMITED {self.time_until_next=} sec")
                    time.sleep(self.time_until_next)
                else:
                    # Its ok to execute a task now..
                    task = self.func_queue.get(timeout=1)
                    func: Callable = task["func"]
                    func_args = task["func_args"]
                    func(func_args)

            print("---- Queue is empty ----")
            time.sleep(1)
            print("Waiting for jobs..")

    async def run_async(self) -> None:
        while True:
            while not self.func_queue.empty():
                self.rate_limiter()
                print("Looking for a job")
                
                if self.is_rate_limited:
                    print(f"RATE LIMITED {self.time_until_next=} sec")
                    await asyncio.sleep(self.time_until_next)
                else:
                    # Its ok to execute a task now..
                    task = self.func_queue.get(timeout=1)
                    func: Callable = task["func"]
                    func_args = task["func_args"]
                    await func(func_args)

            print("---- Queue is empty ----")
            await asyncio.sleep(1)
            print("Waiting for jobs..")


def get_random_string(length=5):
    letters = string.ascii_lowercase
    result_str = ''.join(random.choice(letters) for i in range(length))
    return f"someMsg-{result_str}"


def send_msg_test(msg):
    print(f"Sent {msg=}")


async def send_msg_test_async(msg):
    print(f"Sent_async {msg=}")


async def async_main():
    rt = QueuedRateLimiter(window_size_in_seconds=4)

    for i in range(12):
        rt.func_queue.put({"func": send_msg_test_async, "func_args": get_random_string()})

    loop.create_task(rt.run_async())  # Essential to get async version running
    
    print("SLEEPING for 5..")
    await asyncio.sleep(2)

    for i in range(12):
        rt.func_queue.put({"func": send_msg_test_async, "func_args": get_random_string()})


def sync_main():
    rt = QueuedRateLimiter(window_size_in_seconds=4)

    for i in range(12):
        rt.func_queue.put({"func": send_msg_test, "func_args": get_random_string()})

    th = threading.Thread(target=rt.run)
    th.start()
    
    print("Sleeping for 5..")
    time.sleep(2)

    for i in range(12):
        rt.func_queue.put({"func": send_msg_test, "func_args": get_random_string()})


if __name__ == "__main__":
    # Run async version
    loop = asyncio.get_event_loop()
    loop.run_until_complete(async_main())
    loop.run_forever()  # Essential to get async version running

    # Run sync version
    #sync_main()