Add rate limiter decorator.
authorScott <[email protected]>
Sat, 8 Jan 2022 23:21:25 +0000 (15:21 -0800)
committerScott <[email protected]>
Sat, 8 Jan 2022 23:21:25 +0000 (15:21 -0800)
decorator_utils.py

index 4f98a6d4062c4fc5bfcbcac0a2fddfc00c607cf5..480543ae97e26498b9d4314c927c5ff85b076757 100644 (file)
@@ -2,7 +2,6 @@
 
 """Decorators."""
 
-import datetime
 import enum
 import functools
 import inspect
@@ -15,7 +14,7 @@ import sys
 import threading
 import time
 import traceback
-from typing import Callable, Optional
+from typing import Any, Callable, Optional
 import warnings
 
 # This module is commonly used by others in here and should avoid
@@ -58,6 +57,31 @@ def invocation_logged(func: Callable) -> Callable:
     return wrapper_invocation_logged
 
 
+def rate_limited(n_per_second: int) -> Callable:
+    """Limit invocation of a wrapped function to n calls per second.
+    Thread safe.
+
+    """
+    min_interval = 1.0 / float(n_per_second)
+
+    def wrapper_rate_limited(func: Callable) -> Callable:
+        last_invocation_time = [0.0]
+
+        def wrapper_wrapper_rate_limited(*args, **kargs) -> Any:
+            while True:
+                elapsed = time.clock_gettime(0) - last_invocation_time[0]
+                wait_time = min_interval - elapsed
+                if wait_time > 0.0:
+                    time.sleep(wait_time)
+                else:
+                    break
+            ret = func(*args, **kargs)
+            last_invocation_time[0] = time.clock_gettime(0)
+            return ret
+        return wrapper_wrapper_rate_limited
+    return wrapper_rate_limited
+
+
 def debug_args(func: Callable) -> Callable:
     """Print the function signature and return value at each call."""