C2-2727: improve loadAuthenticationSubject performance (#2595)

stargate · Jun 5, 2023 · 3442766 · 3442766
1 parent bd30983
commit 3442766
Show file tree

Hide file tree

Showing 6 changed files with 227 additions and 47 deletions.
diff --git a/coordinator/authnz/src/main/java/io/stargate/auth/AuthenticationSubject.java b/coordinator/authnz/src/main/java/io/stargate/auth/AuthenticationSubject.java
@@ -16,61 +16,51 @@
 package io.stargate.auth;
 
 import io.stargate.db.AuthenticatedUser;
+import io.stargate.db.ImmutableAuthenticatedUser;
 import java.util.Collections;
 import java.util.Map;
 import javax.annotation.Nullable;
 import org.immutables.value.Value;
 
 @Value.Immutable
+@Value.Style(builtinContainerAttributes = false)
 public interface AuthenticationSubject {
 
   @Nullable
+  @Value.Parameter
   String token();
 
+  @Value.Parameter
   String roleName();
 
+  @Value.Parameter
   boolean isFromExternalAuth();
 
+  @Value.Parameter
   Map<String, String> customProperties();
 
   default AuthenticatedUser asUser() {
-    return AuthenticatedUser.of(roleName(), token(), isFromExternalAuth(), customProperties());
+    return ImmutableAuthenticatedUser.of(
+        roleName(), token(), isFromExternalAuth(), customProperties());
   }
 
   static AuthenticationSubject of(
       String token, String roleName, boolean fromExternalAuth, Map<String, String> properties) {
-    return ImmutableAuthenticationSubject.builder()
-        .token(token)
-        .roleName(roleName)
-        .isFromExternalAuth(fromExternalAuth)
-        .customProperties(properties)
-        .build();
+    return ImmutableAuthenticationSubject.of(
+        token, roleName, fromExternalAuth, Collections.unmodifiableMap(properties));
   }
 
   static AuthenticationSubject of(String token, String roleName, boolean fromExternalAuth) {
-    return ImmutableAuthenticationSubject.builder()
-        .token(token)
-        .roleName(roleName)
-        .isFromExternalAuth(fromExternalAuth)
-        .customProperties(Collections.emptyMap())
-        .build();
+    return ImmutableAuthenticationSubject.of(
+        token, roleName, fromExternalAuth, Collections.emptyMap());
   }
 
   static AuthenticationSubject of(String token, String roleName) {
-    return ImmutableAuthenticationSubject.builder()
-        .token(token)
-        .roleName(roleName)
-        .isFromExternalAuth(false)
-        .customProperties(Collections.emptyMap())
-        .build();
+    return ImmutableAuthenticationSubject.of(token, roleName, false, Collections.emptyMap());
   }
 
   static AuthenticationSubject of(AuthenticatedUser user) {
-    return ImmutableAuthenticationSubject.builder()
-        .token(user.token())
-        .roleName(user.name())
-        .isFromExternalAuth(user.isFromExternalAuth())
-        .customProperties(user.customProperties())
-        .build();
+    return ImmutableAuthenticationSubject.of(
+        user.token(), user.name(), user.isFromExternalAuth(), user.customProperties());
   }
 }
diff --git a/coordinator/microbench/README.md b/coordinator/microbench/README.md
@@ -0,0 +1,68 @@
+# Stargate Micro-benchmarking
+
+This module contains [JMH](https://github.com/openjdk/jmh) benchmarks for Stargate.
+It uses the [jmh-maven-plugin](https://github.com/metlos/jmh-maven-plugin) to run the benchmarks, primarily from the command line.
+
+## Running benchmarks
+
+> **NOTE:** It's advised that before running the benchmarks you build the project with `./mvnw clean install -DskipTests` to ensure that the latest changes are included. 
+
+To run all benchmarks, ensure you are in the `microbench` directory and run:
+
+```bash
+../mvnw jmh:benchmark
+```
+
+To run a specific benchmark class:
+
+```bash
+../mvnw jmh:benchmark -Djmh.benchmarks=MyBenchmark
+```
+
+The JHM framework allows specification of different parameters for a benchmark run.
+Check how the parameter passing works in the [jmh-maven-plugin](https://github.com/metlos/jmh-maven-plugin#passing-parameters) documentation.
+
+## Using profilers
+
+JMH comes with a set of available profilers that can be used to analyze the benchmark runs.
+You can get the list of available profilers by running:
+
+```bash
+../mvnw jmh:benchmark -Djmh.lprof
+```
+
+To use a profiler, you need to specify it in the `-Djmh.prof=<profiler>` parameter.
+For example, running with `gc` profiler:
+
+```bash
+../mvnw jmh:benchmark -Djmh.prof=gc
+```
+
+Very good tutorial on using profilers with JMH can be found [here](https://gist.github.com/markrmiller/a04f5c734fad879f688123bc312c21af#jmh-profilers).
+
+### Using perf profilers
+
+The `linux-tools-common` bring a set of useful `perf` profilers.
+In order to properly setup the local environment, you need to do the following:
+
+```bash
+sudo apt-get install linux-tools-common linux-tools-generic linux-tools-`uname -r`
+sudo sysctl -w kernel.perf_event_paranoid=-1
+sudo sysctl --system
+```
+
+Listing the profilers should show new available profiles now.
+
+### Using async-profiler
+
+The support for [async-profiler](https://github.com/async-profiler/async-profiler) is also available.
+You'll need to download the profiler and unpack it somewhere on your system.
+Then, you can use one of the following methods to setup the library on your system:
+
+> Ensure asyncProfiler library is on LD_LIBRARY_PATH (Linux), DYLD_LIBRARY_PATH (Mac OS), or -Djava.library.path. Alternatively, point to explicit library location with -prof async:libPath=<path>.
+
+To pass options to the async-profiler, you can use the following specification:
+
+```bash
+../mvnw jmh:benchmark -Djmh.prof=async:libPath=[PATH_TO]/libasyncProfiler.so\;output=flamegraph\;dir=profile-results\;event=alloc
+```
diff --git a/coordinator/microbench/pom.xml b/coordinator/microbench/pom.xml
@@ -0,0 +1,51 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+  <!-- Parent -->
+  <parent>
+    <groupId>io.stargate</groupId>
+    <artifactId>stargate</artifactId>
+    <version>2.0.14-SNAPSHOT</version>
+  </parent>
+  <!-- Artifact props -->
+  <artifactId>microbench</artifactId>
+  <name>Stargate - Coordinator - Micro-benchmarking</name>
+  <properties>
+    <jmh.version>1.36</jmh.version>
+  </properties>
+  <dependencies>
+    <dependency>
+      <groupId>io.stargate.db</groupId>
+      <artifactId>persistence-api</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.stargate.auth</groupId>
+      <artifactId>authnz</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.openjdk.jmh</groupId>
+      <artifactId>jmh-core</artifactId>
+      <version>${jmh.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.commons</groupId>
+      <artifactId>commons-lang3</artifactId>
+      <version>3.12.0</version>
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+  <build>
+    <plugins>
+      <plugin>
+        <groupId>pw.krejci</groupId>
+        <artifactId>jmh-maven-plugin</artifactId>
+        <version>0.2.2</version>
+      </plugin>
+    </plugins>
+  </build>
+</project>
diff --git a/...nator/microbench/src/test/java/io/stargate/jmh/auth/AuthenticationSerializationBench.java b/...nator/microbench/src/test/java/io/stargate/jmh/auth/AuthenticationSerializationBench.java
@@ -0,0 +1,76 @@
+package io.stargate.jmh.auth;
+
+import io.stargate.auth.AuthenticationSubject;
+import io.stargate.db.AuthenticatedUser;
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import org.apache.commons.lang3.RandomStringUtils;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+/**
+ * Benchmarks for {@link AuthenticationSubject} serialization.
+ *
+ * <p>Run with: <code>
+ * ../mvnw jmh:benchmark -Djmh.benchmarks=AuthenticationSerializationBench -Djmh.prof=gc</code>,
+ * gave these results:
+ *
+ * <p><code>
+ * Benchmark                                                                       (propertyCount)   Mode  Cnt     Score     Error   Units
+ * AuthenticationSerializationBench.loadAuthenticationSubject                                    2  thrpt    5     3.384 ±   0.230  ops/us
+ * AuthenticationSerializationBench.loadAuthenticationSubject:·gc.alloc.rate                     2  thrpt    5  1807.395 ± 123.001  MB/sec
+ * AuthenticationSerializationBench.loadAuthenticationSubject:·gc.alloc.rate.norm                2  thrpt    5   560.000 ±   0.001    B/op
+ * AuthenticationSerializationBench.loadAuthenticationSubject:·gc.count                          2  thrpt    5   338.000            counts
+ * AuthenticationSerializationBench.loadAuthenticationSubject:·gc.time                           2  thrpt    5   228.000                ms
+ * AuthenticationSerializationBench.loadAuthenticationSubject                                    4  thrpt    5     2.137 ±   0.029  ops/us
+ * AuthenticationSerializationBench.loadAuthenticationSubject:·gc.alloc.rate                     4  thrpt    5  1728.442 ±  23.497  MB/sec
+ * AuthenticationSerializationBench.loadAuthenticationSubject:·gc.alloc.rate.norm                4  thrpt    5   848.000 ±   0.001    B/op
+ * AuthenticationSerializationBench.loadAuthenticationSubject:·gc.count                          4  thrpt    5   333.000            counts
+ * AuthenticationSerializationBench.loadAuthenticationSubject:·gc.time                           4  thrpt    5   227.000                ms
+ * </code>
+ */
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@State(Scope.Benchmark)
+@Warmup(iterations = 3, time = 5)
+@Measurement(iterations = 5, time = 5)
+@Fork(1)
+public class AuthenticationSerializationBench {
+
+  private Map<String, ByteBuffer> payload;
+
+  @Param({"2", "4"})
+  int propertyCount;
+
+  @Setup(Level.Trial)
+  public void setup() {
+    Map<String, String> properties = new HashMap<>();
+    for (int i = 0; i < propertyCount; i++) {
+      properties.put(
+          RandomStringUtils.randomAlphanumeric(10), RandomStringUtils.randomAlphanumeric(10));
+    }
+
+    AuthenticatedUser user = AuthenticatedUser.of("role", "token", true, properties);
+    payload = AuthenticatedUser.Serializer.serialize(user);
+  }
+
+  @Benchmark
+  public void loadAuthenticationSubject(Blackhole bh) {
+    AuthenticatedUser user = AuthenticatedUser.Serializer.load(payload);
+    AuthenticationSubject subject = AuthenticationSubject.of(user);
+    bh.consume(subject);
+  }
+}
diff --git a/coordinator/persistence-api/src/main/java/io/stargate/db/AuthenticatedUser.java b/coordinator/persistence-api/src/main/java/io/stargate/db/AuthenticatedUser.java
@@ -20,51 +20,44 @@
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Map.Entry;
 import javax.annotation.Nullable;
 import org.immutables.value.Value;
 
 @Value.Immutable
+@Value.Style(builtinContainerAttributes = false)
 public interface AuthenticatedUser extends Serializable {
 
+  @Value.Parameter
   String name();
 
   @Nullable
+  @Value.Parameter
   String token();
 
+  @Value.Parameter
   boolean isFromExternalAuth();
 
+  @Value.Parameter
   Map<String, String> customProperties();
 
   static AuthenticatedUser of(String userName) {
-    return ImmutableAuthenticatedUser.builder()
-        .name(userName)
-        .isFromExternalAuth(false)
-        .customProperties(Collections.emptyMap())
-        .build();
+    return ImmutableAuthenticatedUser.of(userName, null, false, Collections.emptyMap());
   }
 
   static AuthenticatedUser of(String userName, String token) {
-    return ImmutableAuthenticatedUser.builder()
-        .name(userName)
-        .token(token)
-        .isFromExternalAuth(false)
-        .customProperties(Collections.emptyMap())
-        .build();
+    return ImmutableAuthenticatedUser.of(userName, token, false, Collections.emptyMap());
   }
 
   static AuthenticatedUser of(
       String userName,
       String token,
       boolean useTransitionalAuth,
       Map<String, String> customProperties) {
-    return ImmutableAuthenticatedUser.builder()
-        .name(userName)
-        .token(token)
-        .isFromExternalAuth(useTransitionalAuth)
-        .customProperties(customProperties)
-        .build();
+    return ImmutableAuthenticatedUser.of(
+        userName, token, useTransitionalAuth, Collections.unmodifiableMap(customProperties));
   }
 
   class Serializer {
@@ -106,13 +99,13 @@ public static Map<String, ByteBuffer> serialize(AuthenticatedUser user) {
     public static AuthenticatedUser load(Map<String, ByteBuffer> customPayload) {
       ByteBuffer token = customPayload.get(TOKEN);
       ByteBuffer roleName = customPayload.get(ROLE);
-      ByteBuffer isFromExternalAuth = customPayload.get(EXTERNAL);
+      boolean isFromExternalAuth = customPayload.containsKey(EXTERNAL);
 
       if (token == null || roleName == null) {
         throw new IllegalStateException("token and roleName must be provided");
       }
 
-      ImmutableMap.Builder<String, String> map = ImmutableMap.builder();
+      Map<String, String> map = new HashMap<>(customPayload.size() - 2, 1f);
       for (Entry<String, ByteBuffer> e : customPayload.entrySet()) {
         String key = e.getKey();
         if (key.startsWith(CUSTOM_PAYLOAD_NAME_PREFIX)) {
@@ -122,11 +115,11 @@ public static AuthenticatedUser load(Map<String, ByteBuffer> customPayload) {
         }
       }
 
-      return AuthenticatedUser.of(
+      return ImmutableAuthenticatedUser.of(
           StandardCharsets.UTF_8.decode(roleName).toString(),
           StandardCharsets.UTF_8.decode(token).toString(),
-          (isFromExternalAuth != null),
-          map.build());
+          isFromExternalAuth,
+          Collections.unmodifiableMap(map));
     }
   }
 }
diff --git a/coordinator/pom.xml b/coordinator/pom.xml
@@ -215,6 +215,7 @@
         <module>grpc-proto</module>
         <module>bridge-proto</module>
         <module>grpc</module>
+        <module>microbench</module>
         <module>metrics-jersey</module>
         <module>grpc-examples</module>
         <!-- testing last -->
@@ -250,6 +251,7 @@
         <module>grpc-proto</module>
         <module>bridge-proto</module>
         <module>grpc</module>
+        <module>microbench</module>
         <module>metrics-jersey</module>
         <module>grpc-examples</module>
         <!-- testing last -->