Benchmark

package com.atlassian.asap.performance;

import com.atlassian.asap.api.Jwt;
import com.atlassian.asap.api.JwtBuilder;
import com.atlassian.asap.api.SigningAlgorithm;
import com.atlassian.asap.api.client.http.AuthorizationHeaderGenerator;
import com.atlassian.asap.api.exception.AuthenticationFailedException;
import com.atlassian.asap.api.exception.CannotRetrieveKeyException;
import com.atlassian.asap.api.exception.InvalidTokenException;
import com.atlassian.asap.api.server.http.RequestAuthenticator;
import com.atlassian.asap.core.client.http.AuthorizationHeaderGeneratorImpl;
import com.atlassian.asap.core.server.http.RequestAuthenticatorImpl;
import com.atlassian.asap.core.validator.JwtValidator;
import com.atlassian.asap.core.validator.JwtValidatorImpl;
import com.google.common.base.Stopwatch;
import org.apache.commons.math3.stat.descriptive.SummaryStatistics;

import java.net.URI;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static com.atlassian.asap.core.client.SimpleClientRunner.AUDIENCE_SYSPROP;
import static com.atlassian.asap.core.client.SimpleClientRunner.ISSUER_SYSPROP;
import static com.atlassian.asap.core.client.SimpleClientRunner.KEYID_SYSPROP;
import static com.atlassian.asap.core.client.SimpleClientRunner.PRIVATE_KEY_SYSPROP;

/**
 * A main program that runs basic performance checking of the ASAP library, both the client side (signing
 * and serialising tokens) and the server side (parsing and verifying signatures).
 */
public class Benchmark {
    private static final SigningAlgorithm ALGORITHM = SigningAlgorithm.valueOf(System.getProperty("asap.client.algorithm", "RS256"));

    /**
     * A high value for repetitions is recommended, as cryptographic operations are highly influenced by runtime optimizations.
     * A high repetition value allows sufficient time for the effect of such optimizations to plateau
     */
    private static final long REPETITIONS = Long.getLong("asap.benchmark.repetitions", 200000L);

    private static final int NUM_THREADS = Integer.getInteger("asap.benchmark.threads", 1);

    private static final String ISSUER = System.getProperty(ISSUER_SYSPROP, "issuer1");
    private static final String KEY_ID = System.getProperty(KEYID_SYSPROP, "issuer1/rsa-key-for-tests");
    private static final String AUDIENCE = System.getProperty(AUDIENCE_SYSPROP, "audience");

    private static final URI PRIVATE_KEY_BASE_URL = URI.create(System.getProperty(PRIVATE_KEY_SYSPROP,
            "classpath:/privatekeys/"));
    private static final String PUBLIC_KEY_BASE_URL = System.getProperty("asap.public.key.repo.url",
            "https://s3-ap-southeast-2.amazonaws.com/keymaker.syd.dev.atlassian.io/");

    /**
     * Main function to run the benchmark test.
     *
     * @param args command line arguments
     * @throws Exception if something did not go as planned
     */
    public static void main(String[] args) throws Exception {
        Jwt jwtPrototype = JwtBuilder.newJwt()
                .algorithm(ALGORITHM)
                .issuer(ISSUER)
                .audience(AUDIENCE)
                .keyId(KEY_ID)
                .expirationTime(Instant.now().plus(59, ChronoUnit.MINUTES)) // allow enough time for benchmark
                .build();

        System.out.println("============== Benchmark Start ===============");
        clientTest(REPETITIONS, jwtPrototype, PRIVATE_KEY_BASE_URL);
        serverValidTest(REPETITIONS, jwtPrototype, PUBLIC_KEY_BASE_URL, PRIVATE_KEY_BASE_URL);
        System.out.println("============== Benchmark End ================");

    }

    private static void clientTest(long repetitions, Jwt jwtPrototype, URI privateKeyBaseUrl)
            throws InterruptedException {
        System.out.printf("Client test using key %s and %s algorithm with %d repetitions and %d threads...%n", KEY_ID, ALGORITHM, repetitions, NUM_THREADS);
        ExecutorService executorService = Executors.newFixedThreadPool(NUM_THREADS);

        AuthorizationHeaderGenerator authorizationHeaderGenerator =
                AuthorizationHeaderGeneratorImpl.createDefault(privateKeyBaseUrl);
        Runnable runnableJob = () -> {
            try {
                SummaryStatistics stats = new SummaryStatistics();
                for (long i = 0; i < repetitions; i++) {
                    long startTime = System.nanoTime();
                    // for a more realistic test, we create a new Jwt object each time
                    Jwt jwt = JwtBuilder.copyJwt(jwtPrototype).build();
                    authorizationHeaderGenerator.generateAuthorizationHeader(jwt);
                    stats.addValue(System.nanoTime() - startTime);   // calculate stats with resolution of nanos
                }
                System.out.printf("Client-side token generation min/avg/max/stddev = %f/%f/%f/%f ms%n",
                        nanosToMillis(stats.getMin()), nanosToMillis(stats.getMean()),
                        nanosToMillis(stats.getMax()), nanosToMillis(stats.getStandardDeviation()));
            } catch (InvalidTokenException | CannotRetrieveKeyException e) {
                e.printStackTrace();
                System.exit(1);
            }
        };

        Stopwatch stopwatch = Stopwatch.createUnstarted();
        stopwatch.start();
        for (int job = 0; job < NUM_THREADS; job++) {
            executorService.execute(runnableJob);
        }
        executorService.shutdown();
        executorService.awaitTermination(1, TimeUnit.HOURS);
        stopwatch.stop();

        System.out.printf("Elapsed time: %d ms. Throughput: %f requests/sec%n",
                stopwatch.elapsed(TimeUnit.MILLISECONDS),
                1000.0 * repetitions * NUM_THREADS / stopwatch.elapsed(TimeUnit.MILLISECONDS));

    }

    private static void serverValidTest(long repetitions, Jwt jwtPrototype, String publicKeyBaseUrl, URI privateKeyBaseUrl)
            throws CannotRetrieveKeyException, InvalidTokenException, AuthenticationFailedException {
        System.out.printf("Server test using key %s and %s algorithm with %d repetitions...%n", KEY_ID, ALGORITHM, repetitions);
        String firstAudience = jwtPrototype.getClaims().getAudience().iterator().next();
        JwtValidator jwtValidator = JwtValidatorImpl.createDefault(firstAudience, publicKeyBaseUrl);
        RequestAuthenticator requestAuthenticator = new RequestAuthenticatorImpl(jwtValidator);
        Stopwatch stopwatch = Stopwatch.createUnstarted();
        AuthorizationHeaderGenerator authorizationHeaderGenerator =
                AuthorizationHeaderGeneratorImpl.createDefault(privateKeyBaseUrl);
        SummaryStatistics stats = new SummaryStatistics();

        // generate an authorisation header before starting the timer
        final String authorizationHeader = authorizationHeaderGenerator.generateAuthorizationHeader(jwtPrototype);

        stopwatch.start();
        for (long i = 0; i < repetitions; i++) {
            long startTime = System.nanoTime();
            requestAuthenticator.authenticateRequest(authorizationHeader);
            stats.addValue(TimeUnit.MILLISECONDS.convert(System.nanoTime() - startTime, TimeUnit.NANOSECONDS));

        }
        stopwatch.stop();

        System.out.printf("Elapsed time: %d ms. Throughput: %f requests/sec%n",
                stopwatch.elapsed(TimeUnit.MILLISECONDS),
                1000.0 * repetitions / stopwatch.elapsed(TimeUnit.MILLISECONDS));
        // the latency between server receiving a verification request and then completing the verification with a public key retrieved from key server.
        System.out.printf("Server-side Token Verification Latency min/avg/max/stddev = %f/%f/%f/%f ms%n",
                stats.getMin(), stats.getMean(), stats.getMax(), stats.getStandardDeviation());
    }

    private static double nanosToMillis(double nanos) {
        return nanos / 1_000_000;
    }
}