From b522a0f76386490f980cbae0079d62a03be9b0ba Mon Sep 17 00:00:00 2001
From: Piero Toffanin <pt@masseranolabs.com>
Date: Mon, 5 Feb 2024 15:03:27 -0500
Subject: [PATCH] Add --req-time-cost

---
 libretranslate/app.py            | 23 +++++++++++++++++++----
 libretranslate/default_values.py |  5 +++++
 libretranslate/main.py           |  7 +++++++
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/libretranslate/app.py b/libretranslate/app.py
index 17f15c5..4ea7be0 100644
--- a/libretranslate/app.py
+++ b/libretranslate/app.py
@@ -1,4 +1,5 @@
 import io
+import math
 import os
 import re
 import tempfile
@@ -219,6 +220,13 @@ def create_app(args):
 
         from flask_limiter import Limiter
 
+        def limits_cost():
+          req_cost = getattr(request, 'req_cost', 1)
+          if args.req_time_cost > 0:
+            return max(req_cost, int(math.ceil(getattr(request, 'duration', 0) / args.req_time_cost)))
+          else:
+            return req_cost
+
         limiter = Limiter(
             key_func=get_remote_address,
             default_limits=get_routes_limits(
@@ -226,7 +234,7 @@ def create_app(args):
             ),
             storage_uri=args.req_limit_storage,
             default_limits_deduct_when=lambda req: True, # Force cost to be called after the request
-            default_limits_cost=lambda: getattr(request, 'req_cost', 1)
+            default_limits_cost=limits_cost
         )
     else:
         from .no_limiter import Limiter
@@ -325,12 +333,19 @@ def create_app(args):
                 status = e.code
                 raise e
               finally:
-                duration = max(default_timer() - start_t, 0)
-                measure_request.labels(request.path, status, ip, ak).observe(duration)
+                request.duration = max(default_timer() - start_t, 0)
+                measure_request.labels(request.path, status, ip, ak).observe(request.duration)
                 g.dec()
           return measure_func
         else:
-          return func
+          @wraps(func)
+          def time_func(*a, **kw):
+            start_t = default_timer()
+            try:
+              return func(*a, **kw)
+            finally:
+              request.duration = max(default_timer() - start_t, 0)
+          return time_func
 
     @bp.errorhandler(400)
     def invalid_api(e):
diff --git a/libretranslate/default_values.py b/libretranslate/default_values.py
index 96a833e..b93cd76 100644
--- a/libretranslate/default_values.py
+++ b/libretranslate/default_values.py
@@ -81,6 +81,11 @@ _default_options_objects = [
         'default_value': -1,
         'value_type': 'int'
     },
+    {
+        'name': 'REQ_TIME_COST',
+        'default_value': -1,
+        'value_type': 'int'
+    },
     {
         'name': 'BATCH_LIMIT',
         'default_value': -1,
diff --git a/libretranslate/main.py b/libretranslate/main.py
index 581e4f5..8337c70 100644
--- a/libretranslate/main.py
+++ b/libretranslate/main.py
@@ -63,6 +63,13 @@ def get_args():
         metavar="<number>",
         help="Set the maximum number of request limit offences that a client can exceed before being banned. (%(default)s)",
     )
+    parser.add_argument(
+        "--req-time-cost",
+        default=DEFARGS['REQ_TIME_COST'],
+        type=int,
+        metavar="<number>",
+        help="Considers a time cost (in seconds) for request limiting purposes. If a request takes 10 seconds and this value is set to 5, the request cost is either 2 or the actual request cost (whichever is greater). (%(default)s)",
+    )
     parser.add_argument(
         "--batch-limit",
         default=DEFARGS['BATCH_LIMIT'],