Cache sql sanitized extraction (#2094)

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/JdbcUtils.java

@@ -7,6 +7,7 @@ package io.opentelemetry.javaagent.instrumentation.jdbc;
 
 import static io.opentelemetry.javaagent.instrumentation.api.db.QueryNormalizationConfig.isQueryNormalizationEnabled;
 
+import io.opentelemetry.javaagent.instrumentation.api.BoundedCache;
 import io.opentelemetry.javaagent.instrumentation.api.db.SqlSanitizer;
 import io.opentelemetry.javaagent.instrumentation.api.db.SqlStatementInfo;
 import java.lang.reflect.Field;
@@ -22,6 +23,8 @@ public abstract class JdbcUtils {
   private static final boolean NORMALIZATION_ENABLED = isQueryNormalizationEnabled("jdbc");
 
   private static Field c3poField = null;
+  private static final BoundedCache<String, SqlStatementInfo> sqlToStatementInfoCache =
+      BoundedCache.build(1000);
 
   /** Returns the unwrapped connection or null if exception was thrown. */
   public static Connection connectionFromStatement(Statement statement) {
@@ -71,6 +74,11 @@ public abstract class JdbcUtils {
     if (!NORMALIZATION_ENABLED) {
       return new SqlStatementInfo(sql, null, null);
     }
-    return SqlSanitizer.sanitize(sql);
+    return sqlToStatementInfoCache.get(
+        sql,
+        k -> {
+          log.trace("SQL statement cache miss");
+          return SqlSanitizer.sanitize(sql);
+        });
   }
 }

javaagent-api/src/main/java/io/opentelemetry/javaagent/instrumentation/api/BoundedCache.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.api;
+
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+
+/** An LRU cache that has a fixed maximum size. */
+public interface BoundedCache<K, V> {
+
+  V get(K key, Function<? super K, ? extends V> mappingFunction);
+
+  static <K, V> BoundedCache<K, V> build(long maxSize) {
+    return Provider.get().build(maxSize);
+  }
+
+  interface Builder {
+    <K, V> BoundedCache<K, V> build(long maxSize);
+  }
+
+  class Provider {
+    /*
+     The default implementation just delegates to the lookup function and should not normally be used.
+     It will be replaced at startup by the AgentInstaller.
+    */
+    private static final Builder NEVER_ACTUALLY_CACHES =
+        new Builder() {
+          @Override
+          public <K, V> BoundedCache<K, V> build(long maxSize) {
+            return (key, mappingFunction) -> mappingFunction.apply(key);
+          }
+        };
+    private static final AtomicReference<Builder> builderRef =
+        new AtomicReference<>(NEVER_ACTUALLY_CACHES);
+
+    private Provider() {}
+
+    public static boolean registerIfAbsent(Builder builder) {
+      return builderRef.compareAndSet(NEVER_ACTUALLY_CACHES, builder);
+    }
+
+    // Method exists for testing only
+    static void reset() {
+      builderRef.set(NEVER_ACTUALLY_CACHES);
+    }
+
+    public static Builder get() {
+      return builderRef.get();
+    }
+  }
+}

javaagent-api/src/test/java/io/opentelemetry/javaagent/instrumentation/api/BoundedCacheTest.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.api;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+
+class BoundedCacheTest {
+  private AtomicInteger cacheMisses;
+
+  @BeforeEach
+  void setup() {
+    cacheMisses = new AtomicInteger(0);
+  }
+
+  @AfterEach
+  void reset() {
+    BoundedCache.Provider.reset();
+  }
+
+  String mockLookupFunction(String s) {
+    cacheMisses.incrementAndGet();
+    return s.toUpperCase();
+  }
+
+  @Test
+  @Order(1)
+  void testCanUseBeforeRegister() {
+    BoundedCache<String, String> cache = BoundedCache.build(3);
+    String result1 = cache.get("foo", this::mockLookupFunction);
+    String result2 = cache.get("bAr", this::mockLookupFunction);
+    assertThat(result1).isEqualTo("FOO");
+    assertThat(result2).isEqualTo("BAR");
+    assertThat(cacheMisses.get()).isEqualTo(2);
+  }
+
+  @Test
+  @Order(2)
+  void testRegisterUsesInstance() {
+    Map<String, String> map = new HashMap<>();
+    BoundedCache.Builder builder = buildMapBackedBuilder(map);
+    BoundedCache.Provider.registerIfAbsent(builder);
+    BoundedCache<String, String> cache = BoundedCache.build(3);
+    String result1 = cache.get("foo", this::mockLookupFunction);
+    String result2 = cache.get("fOo", this::mockLookupFunction);
+    String result3 = cache.get("foo", this::mockLookupFunction);
+    assertThat(result1).isEqualTo("FOO");
+    assertThat(result2).isEqualTo("FOO");
+    assertThat(result3).isEqualTo("FOO");
+    assertThat(cacheMisses.get()).isEqualTo(2); // once for "foo" once for "fOo"
+    assertThat(map.size()).isEqualTo(2);
+    assertThat(map).containsKey("foo");
+    assertThat(map).containsKey("fOo");
+  }
+
+  @Test
+  @Order(3)
+  void testRegisterMultipleFails() {
+    Map<String, String> map = new HashMap<>();
+    BoundedCache.Builder builder = buildMapBackedBuilder(map);
+    assertThat(BoundedCache.Provider.registerIfAbsent(builder)).isTrue();
+    assertThat(BoundedCache.Provider.registerIfAbsent(builder)).isFalse();
+  }
+
+  @NotNull
+  private BoundedCache.Builder buildMapBackedBuilder(Map map) {
+    return new BoundedCache.Builder() {
+      @Override
+      public <K, V> BoundedCache<K, V> build(long maxSize) {
+        return new MapBackedCache<K, V>((Map<K, V>) map);
+      }
+    };
+  }
+
+  private static class MapBackedCache<K, V> implements BoundedCache<K, V> {
+    private final Map<K, V> map;
+
+    public MapBackedCache(Map<K, V> map) {
+      this.map = map;
+    }
+
+    @Override
+    public V get(K key, Function<? super K, ? extends V> mappingFunction) {
+      V v = map.get(key);
+      if (v == null) {
+        v = mappingFunction.apply(key);
+        map.put(key, v);
+      }
+      return v;
+    }
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java

@@ -69,6 +69,8 @@ public class AgentInstaller {
     BootstrapPackagePrefixesHolder.setBoostrapPackagePrefixes(loadBootstrapPackagePrefixes());
     // WeakMap is used by other classes below, so we need to register the provider first.
     AgentTooling.registerWeakMapProvider();
+    // Instrumentation can use a bounded cache, so register here.
+    AgentTooling.registerBoundedCacheProvider();
     // this needs to be done as early as possible - before the first Config.get() call
     ConfigInitializer.initialize();
   }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentTooling.java

@@ -5,8 +5,11 @@
 
 package io.opentelemetry.javaagent.tooling;
 
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
 import io.opentelemetry.javaagent.bootstrap.WeakCache;
 import io.opentelemetry.javaagent.bootstrap.WeakCache.Provider;
+import io.opentelemetry.javaagent.instrumentation.api.BoundedCache;
 import io.opentelemetry.javaagent.instrumentation.api.WeakMap;
 import io.opentelemetry.javaagent.tooling.bytebuddy.AgentCachingPoolStrategy;
 import io.opentelemetry.javaagent.tooling.bytebuddy.AgentLocationStrategy;
@@ -33,6 +36,21 @@ public class AgentTooling {
     }
   }
 
+  /**
+   * Instances of BoundCache are backed by a guava instance that lives in the agent classloader and
+   * is bridged to user/instrumentation classloader through the BoundedCache.Provider interface.
+   */
+  static void registerBoundedCacheProvider() {
+    BoundedCache.Provider.registerIfAbsent(
+        new BoundedCache.Builder() {
+          @Override
+          public <K, V> BoundedCache<K, V> build(long maxSize) {
+            Cache<K, V> cache = CacheBuilder.newBuilder().maximumSize(maxSize).build();
+            return new GuavaBoundedCache<>(cache);
+          }
+        });
+  }
+
   private static <K, V> Provider loadWeakCacheProvider() {
     Iterator<Provider> providers =
         ServiceLoader.load(Provider.class, AgentInstaller.class.getClassLoader()).iterator();

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/GuavaBoundedCache.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling;
+
+import com.google.common.cache.Cache;
+import io.opentelemetry.javaagent.instrumentation.api.BoundedCache;
+import java.util.concurrent.ExecutionException;
+import java.util.function.Function;
+
+class GuavaBoundedCache<K, V> implements BoundedCache<K, V> {
+
+  private final Cache<K, V> delegate;
+
+  public GuavaBoundedCache(Cache<K, V> delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public V get(K key, Function<? super K, ? extends V> mappingFunction) {
+    try {
+      return delegate.get(key, () -> mappingFunction.apply(key));
+    } catch (ExecutionException e) {
+      throw new IllegalStateException("Unexpected cache exception", e);
+    }
+  }
+}

javaagent-tooling/src/test/groovy/io/opentelemetry/javaagent/tooling/GuavaBoundedCacheTest.groovy

@@ -0,0 +1,30 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling
+
+import com.google.common.cache.CacheBuilder
+import spock.lang.Specification
+
+class GuavaBoundedCacheTest extends Specification {
+
+  def "test cache"() {
+
+    when:
+
+    def delegate = CacheBuilder.newBuilder().maximumSize(3).build()
+    def cache = new GuavaBoundedCache<>(delegate)
+    def fn = { x -> x.toUpperCase() }
+
+    then:
+
+    cache.get("foo", fn) == "FOO"
+    cache.get("bar", fn) == "BAR"
+    cache.get("baz", fn) == "BAZ"
+    cache.get("fizz", fn) == "FIZZ"
+    cache.get("buzz", fn) == "BUZZ"
+    delegate.size() == 3
+  }
+}