add rate limiting per engine

searx/shared/__init__.py

@@ -11,7 +11,7 @@ try:
     uwsgi = importlib.import_module('uwsgi')
 except:
     # no uwsgi
-    from .shared_simple import SimpleSharedDict as SharedDict, schedule
+    from .shared_simple import SimpleSharedDict as SharedDict, schedule, run_locked
 
     logger.info('Use shared_simple implementation')
 else:
@@ -32,7 +32,7 @@ else:
 
     else:
         # uwsgi
-        from .shared_uwsgi import UwsgiCacheSharedDict as SharedDict, schedule
+        from .shared_uwsgi import UwsgiCacheSharedDict as SharedDict, schedule, run_locked
 
         logger.info('Use shared_uwsgi implementation')
 

searx/shared/shared_abstract.py

@@ -10,7 +10,7 @@ class SharedDict(ABC):
         pass
 
     @abstractmethod
-    def set_int(self, key: str, value: int):
+    def set_int(self, key: str, value: int, expire: Optional[int] = None):
         pass
 
     @abstractmethod
@@ -18,5 +18,5 @@ class SharedDict(ABC):
         pass
 
     @abstractmethod
-    def set_str(self, key: str, value: str):
+    def set_str(self, key: str, value: str, expire: Optional[int] = None):
         pass

searx/shared/shared_simple.py

@@ -16,14 +16,28 @@ class SimpleSharedDict(shared_abstract.SharedDict):
     def get_int(self, key: str) -> Optional[int]:
         return self.d.get(key, None)
 
-    def set_int(self, key: str, value: int):
+    def set_int(self, key: str, value: int, expire: Optional[int] = None):
         self.d[key] = value
+        if expire:
+            self._expire(key, expire)
 
     def get_str(self, key: str) -> Optional[str]:
         return self.d.get(key, None)
 
-    def set_str(self, key: str, value: str):
+    def set_str(self, key: str, value: str, expire: Optional[int] = None):
         self.d[key] = value
+        if expire:
+            self._expire(key, expire)
+
+    def _expire(self, key: str, expire: int):
+        t = threading.Timer(expire, lambda k, d: d.pop(k), args=[key, self.d])
+        t.daemon = True
+        t.start()
+
+
+def run_locked(func, *args):
+    # SimpleSharedDict is not actually shared, so no locking needed
+    return func(*args)
 
 
 def schedule(delay, func, *args):

searx/shared/shared_uwsgi.py

@@ -2,6 +2,7 @@
 
 import time
 from typing import Optional
+import threading
 import uwsgi  # pyright: ignore # pylint: disable=E0401
 from . import shared_abstract
 
@@ -17,9 +18,11 @@ class UwsgiCacheSharedDict(shared_abstract.SharedDict):
         else:
             return int.from_bytes(value, 'big')
 
-    def set_int(self, key: str, value: int):
+    def set_int(self, key: str, value: int, expire: Optional[int] = None):
         b = value.to_bytes(4, 'big')
         uwsgi.cache_update(key, b)
+        if expire:
+            self._expire(key, expire)
 
     def get_str(self, key: str) -> Optional[str]:
         value = uwsgi.cache_get(key)
@@ -28,9 +31,26 @@ class UwsgiCacheSharedDict(shared_abstract.SharedDict):
         else:
             return value.decode('utf-8')
 
-    def set_str(self, key: str, value: str):
+    def set_str(self, key: str, value: str, expire: Optional[int] = None):
         b = value.encode('utf-8')
         uwsgi.cache_update(key, b)
+        if expire:
+            self._expire(key, expire)
+
+    def _expire(self, key: str, expire: int):
+        t = threading.Timer(expire, uwsgi.cache_del, args=[key])
+        t.daemon = True
+        t.start()
+
+
+def run_locked(func, *args):
+    result = None
+    uwsgi.lock()
+    try:
+        result = func(*args)
+    finally:
+        uwsgi.unlock()
+    return result
 
 
 def schedule(delay, func, *args):