diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 24321a569..a98400bdd 100644 --- a/.github/workflows/test.yaml +++ b/.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: diff --git a/nix/tools/generate_targets.py b/nix/tools/generate_targets.py index 79e869ef8..762a2c21a 100644 --- a/nix/tools/generate_targets.py +++ b/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: - """Generate an HS256 JWT""" - now = int(time.time()) +def generate_jwt(now: int, exp_inc: Optional[int], is_hs: bool) -> str: + """Generate an HS256 or RS256 JWT""" 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 -# 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 will have to clean cached expired JWTs. -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 append_targets(lines: list[str], token: str): + lines.append(f"OPTIONS {URL}/authors_only") + lines.append(f"Authorization: Bearer {token}") + lines.append("") # blank line to separate requests 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): - token = generate_jwt(estimate_adequate_jwt_exp_increase(i)) - lines.append(f"OPTIONS {URL}/authors_only") - lines.append(f"Authorization: Bearer {token}") - lines.append("") # blank line to separate requests + now = int(start_time) + + lines = [] + + # 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)") diff --git a/nix/tools/generate_targets_rsa.py b/nix/tools/generate_targets_rsa.py deleted file mode 100644 index 3de7b209b..000000000 --- a/nix/tools/generate_targets_rsa.py +++ /dev/null @@ -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() diff --git a/nix/tools/loadtest.nix b/nix/tools/loadtest.nix index 7ab7a97c6..b70350c3d 100644 --- a/nix/tools/loadtest.nix +++ b/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-rsa])" + "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_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 - export PGRST_JWT_CACHE_MAX_LIFETIME="0" - fi + jwt-hs-cache) + ${genTargetsHS} "$_arg_testdir"/gen_targets.http + ;; - # 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[@]}\"" + jwt-hs-cache-worst) + ${genTargetsHS} --worst "$_arg_testdir"/gen_targets.http ;; jwt-rsa) - - ${genTargetsRSA} "$_arg_testdir"/gen_targets.http "$_arg_testdir"/gen_jwk.json - + ${genTargetsHS} --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http + export PGRST_JWT_CACHE_MAX_LIFETIME="0" export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json" + ;; - if [ "$_arg_jwtcache" = "off" ]; then - export PGRST_JWT_CACHE_MAX_LIFETIME="0" - fi + jwt-rsa-cache) + ${genTargetsHS} --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http + export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json" + ;; - # 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[@]}\"" + jwt-rsa-cache-worst) + ${genTargetsHS} --worst --rsa="$_arg_testdir"/gen_jwk.json "$_arg_testdir"/gen_targets.http + export PGRST_JWT_SECRET="@$_arg_testdir/gen_jwk.json" ;; *) - - # 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" {