Improve performance: JWT decoder cache - benchmark [do not merge]

jwt-server/spring/jwt-spring-common/src/test/java/org/entur/jwt/spring/decode/cache/DecodedJwtCacheJwtDecoderBenchmarkTest.java

@@ -0,0 +1,205 @@
+package org.entur.jwt.spring.decode.cache;
+
+import com.nimbusds.jose.JWSAlgorithm;
+import com.nimbusds.jose.JWSHeader;
+import com.nimbusds.jose.crypto.RSASSASigner;
+import com.nimbusds.jose.jwk.JWKSet;
+import com.nimbusds.jose.jwk.RSAKey;
+import com.nimbusds.jose.jwk.source.CachingJWKSetSource;
+import com.nimbusds.jwt.JWTClaimsSet;
+import com.nimbusds.jwt.SignedJWT;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.springframework.security.oauth2.core.OAuth2TokenValidatorResult;
+import org.springframework.security.oauth2.jwt.JwtDecoder;
+import org.springframework.security.oauth2.jwt.NimbusJwtDecoder;
+
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.interfaces.RSAPublicKey;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Date;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Back-to-back-request benchmark of {@link DecodedJwtCacheJwtDecoder}, run manually (not
+ * part of the regular build) to get a ballpark of the speedup the decoded-JWT cache gives
+ * over always fully re-verifying the JWT signature (i.e. behaviour without this cache,
+ * as on {@code master}).
+ * <p>
+ * Three scenarios are simulated, each issuing {@link #REQUESTS} decode calls back to back
+ * on a single thread, after a warm-up phase to let the JIT settle:
+ * <ul>
+ *     <li>{@code baselineNoCache} - every request is fully verified, no cache involved.
+ *     This approximates {@code master} for every request, and also the decoded-JWT cache's
+ *     worst case (a cache miss).</li>
+ *     <li>{@code singleTokenReused} - a single token is decoded over and over, i.e. a
+ *     100% cache-hit rate. This approximates a single long-lived client-credentials token
+ *     reused for many outgoing service-to-service requests.</li>
+ *     <li>{@code tokenPoolReused} - a small pool of distinct tokens (simulating several
+ *     concurrent clients/tenants) is cycled through repeatedly, i.e. a more realistic
+ *     mixed workload with a high but not 100% cache-hit rate.</li>
+ * </ul>
+ * To run: remove/comment out {@code @Disabled} and execute this test class directly.
+ */
+@Tag("performance")
+@Disabled("Manual benchmark - not meant to run as part of the regular build")
+public class DecodedJwtCacheJwtDecoderBenchmarkTest {
+
+    private static final int WARMUP_REQUESTS = 20_000;
+    private static final int REQUESTS = 200_000;
+    private static final int TOKEN_POOL_SIZE = 50;
+
+    @Test
+    public void baselineNoCache() throws Exception {
+        RSAKey signingKey = generateRsaKey("kid-0");
+        JwtDecoder delegate = nimbusDecoder(signingKey);
+        String token = signedToken(signingKey);
+
+        Stats stats = measure("baseline (no cache, every request fully verified)",
+                REQUESTS, WARMUP_REQUESTS, i -> delegate.decode(token));
+
+        assertTrue(stats.opsPerSecond > 0);
+    }
+
+    @Test
+    public void singleTokenReused() throws Exception {
+        RSAKey signingKey = generateRsaKey("kid-0");
+        JwtDecoder delegate = nimbusDecoder(signingKey);
+        String token = signedToken(signingKey);
+
+        DecodedJwtCacheJwtDecoder cachedDecoder = cachedDecoder(delegate, signingKey);
+        cachedDecoder.decode(token); // populate cache (first call is necessarily a miss)
+
+        Stats baseline = measure("baseline (no cache, every request fully verified)",
+                REQUESTS, WARMUP_REQUESTS, i -> delegate.decode(token));
+
+        Stats cached = measure("single token reused (100% cache-hit rate)",
+                REQUESTS, WARMUP_REQUESTS, i -> cachedDecoder.decode(token));
+
+        printSpeedup(baseline, cached);
+
+        assertTrue(cached.opsPerSecond > baseline.opsPerSecond,
+                "Cache-hit decoding is expected to be substantially faster than full re-verification");
+    }
+
+    @Test
+    public void tokenPoolReused() throws Exception {
+        RSAKey signingKey = generateRsaKey("kid-0");
+        JwtDecoder delegate = nimbusDecoder(signingKey);
+
+        List<String> tokens = new ArrayList<>(TOKEN_POOL_SIZE);
+        for (int i = 0; i < TOKEN_POOL_SIZE; i++) {
+            tokens.add(signedToken(signingKey));
+        }
+
+        DecodedJwtCacheJwtDecoder cachedDecoder = cachedDecoder(delegate, signingKey);
+        for (String token : tokens) {
+            cachedDecoder.decode(token); // populate cache for each token in the pool
+        }
+
+        Stats baseline = measure("baseline (no cache, every request fully verified)",
+                REQUESTS, WARMUP_REQUESTS, i -> delegate.decode(tokens.get(i % tokens.size())));
+
+        Stats cached = measure("token pool of " + TOKEN_POOL_SIZE + " reused round-robin (simulated multi-client traffic)",
+                REQUESTS, WARMUP_REQUESTS, i -> cachedDecoder.decode(tokens.get(i % tokens.size())));
+
+        printSpeedup(baseline, cached);
+
+        assertTrue(cached.opsPerSecond > baseline.opsPerSecond,
+                "Cache-hit decoding with a reused token pool is expected to be substantially faster than full re-verification");
+    }
+
+    // --- helpers -------------------------------------------------------------------
+
+    private interface DecodeCall {
+        void decode(int i) throws Exception;
+    }
+
+    private record Stats(String label, double avgLatencyMicros, double opsPerSecond) {
+    }
+
+    /**
+     * Runs {@code call} back to back {@code requests} times (after a warm-up phase of
+     * {@code warmup} calls to let the JIT settle), then reports average per-call latency
+     * and throughput.
+     */
+    private Stats measure(String label, int requests, int warmup, DecodeCall call) throws Exception {
+        for (int i = 0; i < warmup; i++) {
+            call.decode(i);
+        }
+
+        long start = System.nanoTime();
+        for (int i = 0; i < requests; i++) {
+            call.decode(i);
+        }
+        long durationNanos = System.nanoTime() - start;
+
+        double avgLatencyMicros = (durationNanos / 1000.0) / requests;
+        double opsPerSecond = requests / (durationNanos / 1_000_000_000.0);
+
+        System.out.printf("[%s] %.3f us/op, %.0f ops/sec (%d requests back to back)%n",
+                label, avgLatencyMicros, opsPerSecond, requests);
+
+        return new Stats(label, avgLatencyMicros, opsPerSecond);
+    }
+
+    private void printSpeedup(Stats baseline, Stats cached) {
+        System.out.printf("Speedup vs baseline: %.1fx (latency), %.1fx (throughput)%n%n",
+                baseline.avgLatencyMicros / cached.avgLatencyMicros,
+                cached.opsPerSecond / baseline.opsPerSecond);
+    }
+
+    private JwtDecoder nimbusDecoder(RSAKey signingKey) throws Exception {
+        return NimbusJwtDecoder.withPublicKey((RSAPublicKey) signingKey.toRSAPublicKey()).build();
+    }
+
+    private String signedToken(RSAKey signingKey) throws Exception {
+        JWSHeader header = new JWSHeader.Builder(JWSAlgorithm.RS256).keyID(signingKey.getKeyID()).build();
+        JWTClaimsSet claims = new JWTClaimsSet.Builder()
+                .subject("test-subject")
+                .issuer("https://issuer.example.com")
+                .audience("test-audience")
+                .expirationTime(Date.from(Instant.now().plusSeconds(3600)))
+                .issueTime(new Date())
+                .build();
+
+        SignedJWT signedJWT = new SignedJWT(header, claims);
+        signedJWT.sign(new RSASSASigner(signingKey));
+        return signedJWT.serialize();
+    }
+
+    /**
+     * Wraps {@code delegate} with a {@link DecodedJwtCacheJwtDecoder}, and feeds it a
+     * simulated JWKS refresh so it treats {@code signingKey}'s key id as an active
+     * signing key (required before it will actually cache any decoded JWT).
+     */
+    @SuppressWarnings("unchecked")
+    private DecodedJwtCacheJwtDecoder cachedDecoder(JwtDecoder delegate, RSAKey signingKey) {
+        DecodedJwtCacheJwtDecoder cachedDecoder = new DecodedJwtCacheJwtDecoder(
+                delegate, jwt -> OAuth2TokenValidatorResult.success(), 0L, -1);
+
+        JWKSet jwkSet = new JWKSet(List.of(signingKey.toPublicJWK()));
+        CachingJWKSetSource.RefreshCompletedEvent<?> event = mock(CachingJWKSetSource.RefreshCompletedEvent.class);
+        when(event.getJWKSet()).thenReturn(jwkSet);
+        new DecodedJwtCacheJwkEventListener(cachedDecoder).notify(event);
+
+        return cachedDecoder;
+    }
+
+    private RSAKey generateRsaKey(String kid) throws Exception {
+        KeyPairGenerator gen = KeyPairGenerator.getInstance("RSA");
+        gen.initialize(2048);
+        KeyPair pair = gen.generateKeyPair();
+        return new RSAKey.Builder((RSAPublicKey) pair.getPublic())
+                .privateKey(pair.getPrivate())
+                .keyID(kid)
+                .build();
+    }
+}

jwt-server/spring/jwt-spring-grpc-native/src/test/java/org/entur/jwt/spring/grpc/perf/DecodedJwtCacheGrpcBenchmarkCachedTest.java

@@ -0,0 +1,191 @@
+package org.entur.jwt.spring.grpc.perf;
+
+import com.nimbusds.jose.jwk.JWK;
+import com.nimbusds.jose.jwk.JWKSelector;
+import com.nimbusds.jose.jwk.source.JWKSource;
+import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
+import org.entur.jwt.junit5.AccessToken;
+import org.entur.jwt.junit5.AuthorizationServer;
+import org.entur.jwt.spring.JwkSourceMap;
+import org.entur.jwt.spring.decode.cache.DecodedJwtCacheJwtDecoder;
+import org.entur.jwt.spring.grpc.netty.GrpcJwtDecoderHolder;
+import org.entur.jwt.spring.grpc.AbstractGrpcTest;
+import org.entur.jwt.spring.grpc.JwtCallCredentials;
+import org.entur.jwt.spring.grpc.test.GreetingResponse;
+import org.entur.jwt.spring.grpc.test.GreetingServiceGrpc;
+import org.entur.jwt.spring.grpc.test.GreetingServiceGrpc.GreetingServiceBlockingStub;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.MethodOrderer;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestMethodOrder;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.security.oauth2.jwt.JwtDecoder;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.TestPropertySource;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Back-to-back gRPC call benchmark against a real, fully started Spring Boot context
+ * (embedded gRPC server on a real TCP loopback port, real security interceptor chain,
+ * real JWKS mock issuer), with the decoded-JWT cache <b>enabled</b> (this branch's
+ * optimization). Repeated tokens hit the decoded-JWT cache instead of being fully
+ * re-verified, still going through the whole gRPC stack end to end over the network
+ * interface (not an in-process channel).
+ * <p>
+ * See {@link DecodedJwtCacheGrpcBenchmarkNoCacheTest} for the equivalent benchmark with
+ * the cache disabled, for direct comparison.
+ * <p>
+ * Test methods run in a fixed, explicit order ({@link Order}) and each scenario has a
+ * large, independent warm-up phase, so that results are not skewed by JIT/GC/channel
+ * warm-up carried over from a previously run method sharing the same Spring context/JVM.
+ * <p>
+ * Disabled by default - not part of the regular build. Remove {@code @Disabled} to run
+ * manually.
+ */
+@Tag("performance")
+@Disabled("Manual benchmark - not meant to run as part of the regular build")
+@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
+@AuthorizationServer("a")
+@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
+@TestPropertySource(properties = {
+        "entur.jwt.tenants.a.decoder-cache.enabled=true",
+        "entur.jwt.jwk.cache.preemptive.eager.enabled=true",
+        // suppress per-request log noise so it doesn't distort request timings/output
+        "logging.level.org.entur=WARN",
+        "logging.level.org.springframework=WARN",
+        "logging.level.io.grpc=WARN",
+})
+@DirtiesContext
+public class DecodedJwtCacheGrpcBenchmarkCachedTest extends AbstractGrpcTest {
+
+    // large enough that each scenario independently reaches JIT/GC/channel steady state,
+    // regardless of execution order or warm-up carried over from other methods
+    private static final int WARMUP_REQUESTS = 20_000;
+    private static final int REQUESTS = 10_000;
+
+    // 10 parallel clients, 2 tokens each (20 tokens total) - reduced from abt-core's
+    // observed peak (08:00/16:00) concurrent client count, per Entur Compass, to
+    // investigate throughput scaling with fewer concurrent clients
+    private static final int CLIENTS = 10;
+    private static final int TOKENS_PER_CLIENT = 2;
+
+    @Autowired
+    private GrpcJwtDecoderHolder grpcJwtDecoderHolder;
+
+    @Autowired
+    private JwkSourceMap jwkSourceMap;
+
+    @BeforeEach
+    public void readinessProbe() throws Exception {
+        GrpcBenchmarkSupport.assertJacocoAgentDisabled();
+
+        // make sure JWKs are loaded before timing anything
+        Map<String, JWKSource> jwkSources = jwkSourceMap.getJwkSources();
+        for (Map.Entry<String, JWKSource> entry : jwkSources.entrySet()) {
+            JWKSelector selector = mock(JWKSelector.class);
+            when(selector.select(any())).thenReturn(List.of(mock(JWK.class)));
+            entry.getValue().get(selector, null);
+        }
+
+        JwtDecoder jwtDecoder = grpcJwtDecoderHolder.getJwtDecoder();
+        assertThat(jwtDecoder).isInstanceOf(DecodedJwtCacheJwtDecoder.class);
+        ((DecodedJwtCacheJwtDecoder) jwtDecoder).clear();
+    }
+
+    @Test
+    @Order(1)
+    public void singleTokenReused(@AccessToken(by = "a", audience = "https://my.audience") String token) {
+        ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost", 9090).usePlaintext().build();
+        try {
+            GreetingServiceBlockingStub stub = GreetingServiceGrpc.newBlockingStub(channel)
+                    .withCallCredentials(new JwtCallCredentials(token));
+            GrpcBenchmarkSupport.GrpcCall[] calls = {
+                    i -> assertProtectedCallSucceeds(stub)
+            };
+
+            GrpcBenchmarkSupport.Stats stats = GrpcBenchmarkSupport.measureParallel(
+                    "decoded-JWT cache enabled, single token reused (100% cache-hit rate)",
+                    REQUESTS, WARMUP_REQUESTS, calls);
+
+            assertTrue(stats.opsPerSecond() > 0);
+        } finally {
+            channel.shutdown();
+        }
+    }
+
+    @Test
+    @Order(2)
+    public void tokenPoolReused(
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "1") String token1,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "2") String token2,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "3") String token3,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "4") String token4,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "5") String token5,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "6") String token6,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "7") String token7,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "8") String token8,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "9") String token9,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "10") String token10,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "11") String token11,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "12") String token12,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "13") String token13,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "14") String token14,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "15") String token15,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "16") String token16,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "17") String token17,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "18") String token18,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "19") String token19,
+            @AccessToken(by = "a", audience = "https://my.audience", scope = "20") String token20) {
+
+        String[] tokens = {token1, token2, token3, token4, token5, token6, token7, token8, token9, token10, token11, token12, token13, token14, token15, token16, token17, token18, token19, token20};
+
+        ManagedChannel[] channels = new ManagedChannel[CLIENTS];
+        GrpcBenchmarkSupport.GrpcCall[] calls = new GrpcBenchmarkSupport.GrpcCall[CLIENTS];
+        for (int c = 0; c < CLIENTS; c++) {
+            ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost", 9090).usePlaintext().build();
+            channels[c] = channel;
+            String[] clientTokens = new String[TOKENS_PER_CLIENT];
+            for (int t = 0; t < TOKENS_PER_CLIENT; t++) {
+                clientTokens[t] = tokens[c * TOKENS_PER_CLIENT + t];
+            }
+            GreetingServiceBlockingStub[] clientStubs = new GreetingServiceBlockingStub[clientTokens.length];
+            for (int t = 0; t < clientTokens.length; t++) {
+                clientStubs[t] = GreetingServiceGrpc.newBlockingStub(channel)
+                        .withCallCredentials(new JwtCallCredentials(clientTokens[t]));
+            }
+            calls[c] = i -> assertProtectedCallSucceeds(clientStubs[i % clientStubs.length]);
+        }
+
+        try {
+            GrpcBenchmarkSupport.Stats stats = GrpcBenchmarkSupport.measureParallel(
+                    "decoded-JWT cache enabled, " + CLIENTS + " parallel clients, " + TOKENS_PER_CLIENT
+                            + " tokens each (" + tokens.length + " tokens total)",
+                    REQUESTS, WARMUP_REQUESTS, calls);
+
+            assertTrue(stats.opsPerSecond() > 0);
+        } finally {
+            for (ManagedChannel channel : channels) {
+                channel.shutdown();
+            }
+        }
+    }
+
+    private void assertProtectedCallSucceeds(GreetingServiceBlockingStub stub) {
+        GreetingResponse response = stub.protectedWithPartnerTenant(greetingRequest);
+
+        assertThat(response.getMessage()).isEqualTo("Hello protected tenant");
+    }
+}