nix: complete jwt loadtests

Completes the jwt loadtests, now we have non-cached, cached + worst case
for the cache.

- jwt-hs (edited): now has the cache disabled
- jwt-hs-cache: cache enabled for jwt-hs
- jwt-hs-cache-worst: worst case of the jwt-hs cache
- jwt-rsa (edited): now has the cache disabled
- jwt-rsa-cache: cache enabled for jwt-rsa
- jwt-rsa-cache-worst: worst case of the jwt-rsa cache

Also deletes `nix/tools/generate_targets_rsa.py` and uses a single
python script.

Should prove what's mentioned on
https://github.com/PostgREST/postgrest/pull/4084#issuecomment-2998170423

.github/workflows/test.yaml

@@ -113,7 +113,7 @@ jobs:
   loadtest:
     strategy:
       matrix:
-        kind: ['mixed', 'jwt-hs', 'jwt-rsa']
+        kind: ['mixed', 'jwt-hs', 'jwt-hs-cache', 'jwt-hs-cache-worst', 'jwt-rsa', 'jwt-rsa-cache', 'jwt-rsa-cache-worst']
     name: Loadtest
     runs-on: ubuntu-24.04
     steps:

nix/tools/generate_targets.py

@@ -1,51 +1,53 @@
 # generates a file to be used by the vegeta load testing tool
-# it generates TOTAL_TARGETS amount of requests that will be run
 
-# This is a worst case scenario for the JWT cache:
+# It includes a worst case scenario for the JWT cache:
 # - all requests will have a unique JWT so no cache hits
 # - all jwts have an expiration that will be long enough to be
 #   valid at time of request but short enough that already
 #   validated jwts will expire later during the loadtest run
 # - the above guarantees JWT cache purging will happen
-#
+# - we want this to track resource consumption in the worst case
-# We want this to track resource consumption in the worst case
+
+# And a more normal scenario where non-expiring JWTs are picked
+# from an array
 import time
 import argparse
 import sys
 import random
 import jwt
+import jwcrypto.jwk as jwk
+from typing import Optional
+from pathlib import Path
 
-SECRET = b"reallyreallyreallyreallyverysafe"
 URL = "http://postgrest"
-TOTAL_TARGETS = 200000  # tuned by hand to reduce result variance
+
+secret_key = b"reallyreallyreallyreallyverysafe"
+
+key = jwk.JWK.generate(kty="RSA", size=4096)
+private_key = jwt.algorithms.RSAAlgorithm.from_jwk(key.export_private())
+public_key = key.export_public()
 
 
-def generate_jwt(exp_inc: int) -> str:
+def generate_jwt(now: int, exp_inc: Optional[int], is_hs: bool) -> str:
-    """Generate an HS256 JWT"""
+    """Generate an HS256 or RS256 JWT"""
-    now = int(time.time())
     payload = {
         "sub": f"user_{random.getrandbits(32)}",
         "iat": now,
-        "exp": now + exp_inc,
         "role": "postgrest_test_author",
     }
 
-    return jwt.encode(payload, SECRET, "HS256")
+    if exp_inc is not None:
+        payload["exp"] = now + exp_inc
+
+    k = secret_key if is_hs else private_key
+    alg = "HS256" if is_hs else "RS256"
+    return jwt.encode(payload, k, alg)
 
 
-# We want to ensure 401 Unauthorized responses don't happen during
+def append_targets(lines: list[str], token: str):
-# JWT validation, this can happen when the jwt `exp` is too short.
+    lines.append(f"OPTIONS {URL}/authors_only")
-# At the same time, we want to ensure the `exp` is not too big,
+    lines.append(f"Authorization: Bearer {token}")
-# so expires will occur and postgREST will have to clean cached expired JWTs.
+    lines.append("")  # blank line to separate requests
-def estimate_adequate_jwt_exp_increase(iteration: int) -> int:
-    # estimated time takes to build and run postgrest itself
-    build_run_postgrest_time = 2
-    # estimated time it takes to generate the targets file
-    file_generation_time = TOTAL_TARGETS // (10**-5)
-    # estimated exp time so some JWTs will expire
-    dynamic_exp_inc = iteration // 1000
-
-    return build_run_postgrest_time + file_generation_time + dynamic_exp_inc
 
 
 def main():
@@ -56,16 +58,77 @@ def main():
         "output",
         help="Path to write the generated targets file",
     )
+    parser.add_argument(
+        "--worst",
+        dest="worst",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Generate worst case targets for a JWT cache",
+    )
+    parser.add_argument(
+        "--rsa",
+        dest="jwk_path",
+        metavar="JWK_PATH",
+        type=Path,
+        default=None,
+        help="Path for generating a RSA JWK file to sign tokens with",
+    )
+
     args = parser.parse_args()
 
-    lines = []
+    is_hs = args.jwk_path is None
+
+    nsamples = 1000
+    if is_hs:
+        ntargets = 200000
+    else:
+        # The asymmetric targets take too long to compute so we reduce them
+        ntargets = 50000
+
+    if not is_hs:
+        try:
+            with open(args.jwk_path, "w") as jwk:
+                jwk.write(public_key)
+                print(f"Created {args.jwk_path} file containing the RSA JWK")
+        except IOError as e:
+            print(f"Error writing to {args.jwk_path}: {e}", file=sys.stderr)
+            sys.exit(1)
+
+    print(f"Generating {ntargets} targets...")
+
     start_time = time.time()
 
-    for i in range(TOTAL_TARGETS):
+    now = int(start_time)
-        token = generate_jwt(estimate_adequate_jwt_exp_increase(i))
+
-        lines.append(f"OPTIONS {URL}/authors_only")
+    lines = []
-        lines.append(f"Authorization: Bearer {token}")
+
-        lines.append("")  # blank line to separate requests
+    # We want to ensure 401 Unauthorized responses don't happen during
+    # JWT validation, this can happen when the jwt `exp` is too short.
+    # At the same time, we want to ensure the `exp` is not too big,
+    # so expires will occur and postgREST needs to
+    # clean cached expired JWTs
+    if args.worst:
+        # estimated time takes to build and run postgrest itself
+        build_run_postgrest_time = 2
+        # estimated time it takes to generate the targets file
+        # the division numbers are tuned by hand
+        if is_hs:  # hs generation is much faster
+            gen_time = ntargets // 66666
+        else:  # asymmetric is slower so the time is higher
+            gen_time = ntargets // 220
+
+        # estimated exp time so some JWTs will expire
+        inc = build_run_postgrest_time + gen_time
+
+        for i in range(ntargets):
+            token = generate_jwt(now, inc + i // 1000, is_hs)
+            append_targets(lines, token)
+
+    else:
+        tokens = [generate_jwt(now, None, is_hs) for _ in range(nsamples)]
+        for i in range(ntargets):
+            token = random.choice(tokens)
+            append_targets(lines, token)
 
     try:
         with open(args.output, "w") as f:
@@ -75,7 +138,7 @@ def main():
         sys.exit(1)
 
     elapsed = time.time() - start_time
-    print(f"Created {TOTAL_TARGETS} targets with unique JWTs", end=" ")
+    print(f"Created {ntargets} targets", end=" ")
     print(f"in {args.output} ({elapsed:.2f}s)")
 
 

nix/tools/generate_targets_rsa.py

@@ -1,68 +0,0 @@
-# generates a file to be used by the vegeta load testing tool
-# It runs TOTAL_TARGETS amount of requests with
-# TOTAL_JWTS amount of JWTs randomly spread among them
-import time
-import argparse
-import sys
-import random
-import jwt
-import jwcrypto.jwk as jwk
-
-URL = "http://postgrest"
-TOTAL_JWTS = 1000
-TOTAL_TARGETS = 50000  # tuned by hand to reduce result variance
-
-key = jwk.JWK.generate(kty="RSA", size=4096)
-
-private_key = jwt.algorithms.RSAAlgorithm.from_jwk(key.export_private())
-public_key = key.export_public()
-
-
-def generate_jwt() -> str:
-    """Generate an RS256 JWT"""
-    now = int(time.time())
-    payload = {
-        "sub": f"user_{random.getrandbits(32)}",
-        "iat": now,
-        "role": "postgrest_test_author",
-    }
-
-    return jwt.encode(payload, private_key, "RS256")
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Generate Vegeta targets with unique JWTs"
-    )
-    parser.add_argument(
-        "output",
-        help="Path to write the generated targets file",
-    )
-    parser.add_argument("jwk", help="Path to write the generated JWK")
-    args = parser.parse_args()
-
-    tokens = [generate_jwt() for _ in range(TOTAL_JWTS)]
-    lines = []
-    start_time = time.time()
-
-    for i in range(TOTAL_TARGETS):
-        token = random.choice(tokens)
-        lines.append(f"OPTIONS {URL}/authors_only")
-        lines.append(f"Authorization: Bearer {token}")
-        lines.append("")  # blank line to separate requests
-
-    try:
-        with open(args.jwk, "w") as jwk, open(args.output, "w") as f:
-            jwk.write(public_key)
-            f.write("\n".join(lines))
-    except IOError as e:
-        print(f"Error writing to {args.output}: {e}", file=sys.stderr)
-        sys.exit(1)
-
-    elapsed = time.time() - start_time
-    print(f"Created {TOTAL_TARGETS} targets with unique JWTs", end=" ")
-    print(f"in {args.output} ({elapsed:.2f}s)")
-
-
-if __name__ == "__main__":
-    main()

nix/tools/loadtest.nix

@@ -42,10 +42,8 @@ let
           "ARG_OPTIONAL_SINGLE([output], [o], [Filename to dump json output to], [./loadtest/result.bin])"
           "ARG_OPTIONAL_SINGLE([testdir], [t], [Directory to load tests and fixtures from], [./test/load])"
           "ARG_OPTIONAL_SINGLE([kind], [k], [Kind of loadtest], [mixed])"
-          "ARG_OPTIONAL_SINGLE([jwtcache], [], [JWT Cache on/off], [on])"
+          "ARG_TYPE_GROUP_SET([KIND], [KIND], [kind], [mixed,jwt-hs,jwt-hs-cache,jwt-hs-cache-worst,jwt-rsa,jwt-rsa-cache,jwt-rsa-cache-worst])"
-          "ARG_TYPE_GROUP_SET([KIND], [KIND], [kind], [mixed,jwt-hs,jwt-rsa])"
           "ARG_OPTIONAL_SINGLE([monitor], [m], [Monitoring file], [./loadtest/result.csv])"
-          "ARG_TYPE_GROUP_SET([JWTCACHE], [JWTCACHE], [jwtcache], [on,off])"
           "ARG_LEFTOVERS([additional vegeta arguments])"
         ];
         workingDir = "/";
@@ -68,49 +66,54 @@ let
 
         case "$_arg_kind" in
           jwt-hs)
-
             ${genTargetsHS} "$_arg_testdir"/gen_targets.http
+            export PGRST_JWT_CACHE_MAX_LIFETIME="0"
+            ;;
 
-            if [ "$_arg_jwtcache" = "off" ]; then
+          jwt-hs-cache)
-              export PGRST_JWT_CACHE_MAX_LIFETIME="0"
+            ${genTargetsHS} "$_arg_testdir"/gen_targets.http
-            fi
+            ;;
 
-            # shellcheck disable=SC2145
+          jwt-hs-cache-worst)
-            ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \
+            ${genTargetsHS} --worst "$_arg_testdir"/gen_targets.http
-            ${withTools.withPgrst} -m "$_arg_monitor" \
-            sh -c "cd \"$_arg_testdir\" && \
-            ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\""
             ;;
 
           jwt-rsa)
-
+            ${genTargetsHS} --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http
-            ${genTargetsRSA} "$_arg_testdir"/gen_targets.http "$_arg_testdir"/gen_jwk.json
+            export PGRST_JWT_CACHE_MAX_LIFETIME="0"
-
             export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json"
+            ;;
 
-            if [ "$_arg_jwtcache" = "off" ]; then
+          jwt-rsa-cache)
-              export PGRST_JWT_CACHE_MAX_LIFETIME="0"
+            ${genTargetsHS} --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http
-            fi
+            export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json"
+            ;;
 
-            # shellcheck disable=SC2145
+          jwt-rsa-cache-worst)
-            ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \
+            ${genTargetsHS} --worst --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http
-            ${withTools.withPgrst} -m "$_arg_monitor" \
+            export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json"
-            sh -c "cd \"$_arg_testdir\" && \
-            ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\""
             ;;
 
           *)
-
-            # shellcheck disable=SC2145
-            ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \
-            ${withTools.withSlowPg} \
-            ${withTools.withPgrst} -m "$_arg_monitor" \
-            ${withTools.withSlowPgrst} \
-            sh -c "cd \"$_arg_testdir\" && \
-            ${runner} -targets targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\""
             ;;
         esac
 
+        if [ "$_arg_kind" == "mixed" ]; then
+          # shellcheck disable=SC2145
+          ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \
+          ${withTools.withSlowPg} \
+          ${withTools.withPgrst} -m "$_arg_monitor" \
+          ${withTools.withSlowPgrst} \
+          sh -c "cd \"$_arg_testdir\" && \
+          ${runner} -targets targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\""
+        else
+          # shellcheck disable=SC2145
+          ${withTools.withPg} -f "$_arg_testdir"/fixtures.sql \
+          ${withTools.withPgrst} -m "$_arg_monitor" \
+          sh -c "cd \"$_arg_testdir\" && \
+          ${runner} -lazy -targets gen_targets.http -output \"$abs_output\" \"''${_arg_leftovers[@]}\""
+        fi
+
         ${vegeta}/bin/vegeta report -type=text "$_arg_output"
       '';
 
@@ -252,13 +255,6 @@ let
       }
       (builtins.readFile ./generate_targets.py);
 
-  genTargetsRSA =
-    writers.writePython3 "postgrest-gen-loadtest-targets-rsa"
-      {
-        libraries = [ python3Packages.pyjwt python3Packages.jwcrypto ];
-      }
-      (builtins.readFile ./generate_targets_rsa.py);
-
   mergeMonitorResults =
     writers.writePython3 "postgrest-merge-monitor-results"
       {