Testing, fixes

libretranslate/app.py

@@ -21,7 +21,7 @@ from werkzeug.exceptions import HTTPException
 from werkzeug.http import http_date
 from flask_babel import Babel
 
-from libretranslate import flood, secret, remove_translated_files, security, storage
+from libretranslate import scheduler, flood, secret, remove_translated_files, security, storage
 from libretranslate.language import detect_languages, improve_translation_formatting
 from libretranslate.locales import (_, _lazy, get_available_locales, get_available_locale_codes, gettext_escaped, 
         gettext_html, lazy_swag, get_alternate_locale_links)
@@ -204,11 +204,12 @@ def create_app(args):
 
         limiter = Limiter()
 
-    if args.req_flood_threshold > 0:
-        flood.setup(args.req_flood_threshold)
-    if args.api_keys and args.require_api_key_secret:
-        secret.setup()
+    if not "gunicorn" in os.environ.get("SERVER_SOFTWARE", ""):
+      # Gunicorn starts the scheduler in the master process
+      scheduler.setup(args)
 
+    flood.setup(args)
+    secret.setup(args)
 
     measure_request = None
     gauge_request = None

libretranslate/flood.py

@@ -1,11 +1,5 @@
-import atexit
-from multiprocessing import Value
-
 from libretranslate.storage import get_storage
-from apscheduler.schedulers.background import BackgroundScheduler
 
-
-setup_scheduler = Value('b', False)
 active = False
 threshold = -1
 
@@ -20,30 +14,18 @@ def forgive_banned():
         if banned[ip] <= 0:
             clear_list.append(ip)
         else:
-            banned[ip] = min(threshold, banned[ip]) - 1
+            s.set_hash_int("banned", ip, min(threshold, banned[ip]) - 1)
 
     for ip in clear_list:
         s.del_hash("banned", ip)
 
-def setup(violations_threshold=100):
+def setup(args):
     global active
     global threshold
 
+    if args.req_flood_threshold > 0:
         active = True
-    threshold = violations_threshold
-
-    # Only setup the scheduler and secrets on one process
-    if not setup_scheduler.value:
-        setup_scheduler.value = True
-
-        scheduler = BackgroundScheduler()
-        scheduler.add_job(func=forgive_banned, trigger="interval", minutes=30)
-
-        scheduler.start()
-
-        # Shut down the scheduler when exiting the app
-        atexit.register(lambda: scheduler.shutdown())
-
+        threshold = args.req_flood_threshold
 
 def report(request_ip):
     if active:

libretranslate/scheduler.py

@@ -0,0 +1,23 @@
+import atexit
+from apscheduler.schedulers.background import BackgroundScheduler
+scheduler = None
+
+def setup(args):
+    from libretranslate.flood import forgive_banned
+    from libretranslate.secret import rotate_secrets
+
+    global scheduler
+
+    if scheduler is None:
+        scheduler = BackgroundScheduler()
+
+        if args.req_flood_threshold > 0:
+            scheduler.add_job(func=forgive_banned, trigger="interval", minutes=10)
+
+        if args.api_keys and args.require_api_key_secret:
+            scheduler.add_job(func=rotate_secrets, trigger="interval", minutes=30)
+        
+        scheduler.start()
+
+        # Shut down the scheduler when exiting the app
+        atexit.register(lambda: scheduler.shutdown())

libretranslate/secret.py

@@ -1,7 +1,7 @@
 import atexit
 import random
 import string
-from multiprocessing import Value
+from multiprocessing.dummy import Value
 
 from libretranslate.storage import get_storage
 from apscheduler.schedulers.background import BackgroundScheduler
@@ -16,6 +16,9 @@ def rotate_secrets():
     secret_1 = s.get_str("secret_1")
     s.set_str("secret_0", secret_1)
     s.set_str("secret_1", generate_secret())
+    print(s.get_str("secret_0"))
+    print(s.get_str("secret_1"))
+    
 
 def secret_match(secret):
     s = get_storage()
@@ -24,19 +27,8 @@ def secret_match(secret):
 def get_current_secret():
     return get_storage().get_str("secret_1")
 
-def setup():
-    # Only setup the scheduler and secrets on one process
-    if not setup_secrets.value:
-        setup_secrets.value = True
-        
+def setup(args):
+    if args.api_keys and args.require_api_key_secret:
         s = get_storage()
         s.set_str("secret_0", generate_secret())
         s.set_str("secret_1", generate_secret())
-        
-        scheduler = BackgroundScheduler()
-        scheduler.add_job(func=rotate_secrets, trigger="interval", minutes=30)
-        
-        scheduler.start()
-
-        # Shut down the scheduler when exiting the app
-        atexit.register(lambda: scheduler.shutdown())

libretranslate/storage.py

@@ -20,7 +20,7 @@ class Storage:
     def get_str(self, key):
         raise Exception("not implemented")
 
-    def set_hash_value(self, ns, key, value):
+    def set_hash_int(self, ns, key, value):
         raise Exception("not implemented")
     def get_hash_int(self, ns, key):
         raise Exception("not implemented")
@@ -56,6 +56,11 @@ class MemoryStorage(Storage):
     def get_str(self, key):
         return str(self.store.get(key, ""))
 
+    def set_hash_int(self, ns, key, value):
+        if ns not in self.store:
+            self.store[ns] = {}
+        self.store[ns][key] = int(value)
+
     def get_hash_int(self, ns, key):
         d = self.store.get(ns, {})
         return int(d.get(key, 0))
@@ -79,7 +84,10 @@ class MemoryStorage(Storage):
             self.store[ns][key] -= 1
 
     def get_all_hash_int(self, ns):
+        if ns in self.store:
             return [{str(k): int(v)} for k,v in self.store[ns].items()]
+        else:
+            return []
 
     def del_hash(self, ns, key):
         del self.store[ns][key]
@@ -123,6 +131,9 @@ class RedisStorage(Storage):
         else:
             return int(v)
     
+    def set_hash_int(self, ns, key, value):
+        self.conn.hset(ns, key, value)
+
     def inc_hash_int(self, ns, key):
         return int(self.conn.hincrby(ns, key))
 
@@ -130,10 +141,10 @@ class RedisStorage(Storage):
         return int(self.conn.hincrby(ns, key, -1))
 
     def get_all_hash_int(self, ns):
-        return [{k.decode("utf-8"): int(v)} for k,v in self.conn.hgetall(ns).items()]
+        return {k.decode("utf-8"): int(v) for k,v in self.conn.hgetall(ns).items()}
 
     def del_hash(self, ns, key):
-        conn.hdel(ns, key)
+        self.conn.hdel(ns, key)
 
 def setup(storage_uri):
     global storage

scripts/gunicorn_conf.py

@@ -1,4 +1,42 @@
 from prometheus_client import multiprocess
+import re
+import sys
 
 def child_exit(server, worker):
     multiprocess.mark_process_dead(worker.pid)
+
+def on_starting(server):
+    # Parse command line arguments
+    proc_name = server.cfg.default_proc_name
+    kwargs = {}
+    if proc_name.startswith("wsgi:app"):
+        str_args = re.sub('wsgi:app\s*\(\s*(.*)\s*\)', '\\1', proc_name).strip().split(",")
+        for a in str_args:
+            if "=" in a:
+                k,v = a.split("=")
+                k = k.strip()
+                v = v.strip()
+
+                if v.lower() in ["true", "false"]:
+                    v = v.lower() == "true"
+                elif v[0] == '"':
+                    v = v[1:-1]
+                kwargs[k] = v
+
+    from libretranslate.main import get_args
+    sys.argv = ['--wsgi']
+    for k in kwargs:
+        ck = k.replace("_", "-")
+        if isinstance(kwargs[k], bool) and kwargs[k]:
+            sys.argv.append("--" + ck)
+        else:
+            sys.argv.append("--" + ck)
+            sys.argv.append(kwargs[k])
+
+    args = get_args()
+    
+    from libretranslate import storage, scheduler, flood, secret
+    storage.setup(args.shared_storage)
+    scheduler.setup(args)
+    flood.setup(args)
+    secret.setup(args)