Fix type-pool cache (#1828)

* Draft fix for type-pool cache

* Enhance and fix cache clearance logic

* Fix periodic clear

* Fix SoftlyReferencingTypePoolCacheTest

* Adding to CHANGELOG

Co-authored-by: Felix Barnsteiner <[email protected]>

Author: eyalkoren

CHANGELOG.asciidoc

@@ -33,7 +33,7 @@ endif::[]
 === Java Agent version 1.x
 
 [[release-notes-1.24.0]]
-==== 1.24.0 - 2021/05/28
+==== 1.24.0 - 2021/05/31
 
 [float]
 ===== Features
@@ -52,6 +52,8 @@ events. With the new `OVERRIDE` option, non-file logs can be ECS-reformatted aut
 * Avoid `IllegalStateException` when multiple `tracestate` headers are used - {pull}1808[#1808]
 * Ensure CLI attach avoids `sudo` only when required and avoid blocking - {pull}1819[#1819]
 * Avoid sending metric-sets without samples, so to adhere to the intake API - {pull}1826[#1826]
+* Fixing our type-pool cache, so that it can't cause OOM (softly-referenced), and it gets cleared when not used for
+a while - {pull}1828[#1828]
 
 [float]
 ===== Refactors

apm-agent-core/src/main/java/co/elastic/apm/agent/bci/bytebuddy/SoftlyReferencingTypePoolCache.java

@@ -37,7 +37,6 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Caches {@link TypeDescription}s which speeds up type matching -
@@ -46,7 +45,12 @@
  * Without a type pool cache those types would have to be re-loaded from the file system if their {@link Class} has not been loaded yet.
  * <p>
  * In order to avoid {@link OutOfMemoryError}s because of this cache,
- * the {@link TypePool.CacheProvider}s are wrapped in a {@link SoftReference}
+ * the {@link TypePool.CacheProvider}s are wrapped in a {@link SoftReference}.
+ * If the underlying cache is being GC'ed, a new one is immediately being created instead, so to both prevent OOME and
+ * enable continuous caching.
+ * Any cache that is not being accessed for some time (configurable through constructor), is being cleared, so that
+ * caches don't occupy heap memory until there is a heap stress. This fits the typical pattern of most type matching
+ * occurring at the beginning of the agent attachment.
  * </p>
  */
 public class SoftlyReferencingTypePoolCache extends AgentBuilder.PoolStrategy.WithTypePoolCache {
@@ -58,6 +62,7 @@ public class SoftlyReferencingTypePoolCache extends AgentBuilder.PoolStrategy.Wi
      */
     private final WeakConcurrentMap<ClassLoader, CacheProviderWrapper> cacheProviders =
         new NullSafeWeakConcurrentMap<ClassLoader, CacheProviderWrapper>(false);
+
     private final ElementMatcher<ClassLoader> ignoredClassLoaders;
 
     public SoftlyReferencingTypePoolCache(final TypePool.Default.ReaderMode readerMode,
@@ -80,16 +85,13 @@ protected TypePool.CacheProvider locate(ClassLoader classLoader) {
             return TypePool.CacheProvider.Simple.withObjectType();
         }
         classLoader = classLoader == null ? getBootstrapMarkerLoader() : classLoader;
-        CacheProviderWrapper cacheProviderRef = cacheProviders.get(classLoader);
-        if (cacheProviderRef == null || cacheProviderRef.get() == null) {
-            cacheProviderRef = new CacheProviderWrapper();
-            cacheProviders.put(classLoader, cacheProviderRef);
+        CacheProviderWrapper cacheProviderWrapper = cacheProviders.get(classLoader);
+        if (cacheProviderWrapper == null) {
+            cacheProviders.put(classLoader, new CacheProviderWrapper());
             // accommodate for race condition
-            cacheProviderRef = cacheProviders.get(classLoader);
+            cacheProviderWrapper = cacheProviders.get(classLoader);
         }
-        final TypePool.CacheProvider cacheProvider = cacheProviderRef.get();
-        // guard against edge case when the soft reference has already been cleared since evaluating the loop condition
-        return cacheProvider != null ? cacheProvider : TypePool.CacheProvider.Simple.withObjectType();
+        return cacheProviderWrapper;
     }
 
     /**
@@ -115,8 +117,9 @@ protected TypePool.CacheProvider locate(ClassLoader classLoader) {
      */
     void clearIfNotAccessedSince(long clearIfNotAccessedSinceMinutes) {
         for (Map.Entry<ClassLoader, CacheProviderWrapper> entry : cacheProviders) {
-            if (System.currentTimeMillis() >= entry.getValue().getLastAccess() + TimeUnit.MINUTES.toMillis(clearIfNotAccessedSinceMinutes)) {
-                cacheProviders.remove(entry.getKey());
+            CacheProviderWrapper cacheWrapper = entry.getValue();
+            if (System.currentTimeMillis() >= cacheWrapper.getLastAccess() + TimeUnit.MINUTES.toMillis(clearIfNotAccessedSinceMinutes)) {
+                cacheWrapper.clear();
             }
         }
     }
@@ -125,21 +128,49 @@ WeakConcurrentMap<ClassLoader, CacheProviderWrapper> getCacheProviders() {
         return cacheProviders;
     }
 
-    private static class CacheProviderWrapper {
-        private final AtomicLong lastAccess = new AtomicLong(System.currentTimeMillis());
-        private final SoftReference<TypePool.CacheProvider> delegate;
+    private static class CacheProviderWrapper implements TypePool.CacheProvider {
+
+        private volatile long lastAccess = System.currentTimeMillis();
+        private volatile SoftReference<TypePool.CacheProvider> delegate;
 
         private CacheProviderWrapper() {
-            this.delegate = new SoftReference<TypePool.CacheProvider>(new TypePool.CacheProvider.Simple());
+            delegate = new SoftReference<TypePool.CacheProvider>(new Simple());
         }
 
         long getLastAccess() {
-            return lastAccess.get();
+            return lastAccess;
+        }
+
+        private TypePool.CacheProvider getDelegate() {
+            TypePool.CacheProvider cacheProvider = delegate.get();
+            if (cacheProvider == null) {
+                synchronized (this) {
+                    cacheProvider = delegate.get();
+                    if (cacheProvider == null) {
+                        cacheProvider = new Simple();
+                        delegate = new SoftReference<TypePool.CacheProvider>(cacheProvider);
+                    }
+                }
+            }
+            return cacheProvider;
         }
 
+        @Override
         @Nullable
-        TypePool.CacheProvider get() {
-            return delegate.get();
+        public TypePool.Resolution find(String name) {
+            lastAccess = System.currentTimeMillis();
+            return getDelegate().find(name);
+        }
+
+        @Override
+        public TypePool.Resolution register(String name, TypePool.Resolution resolution) {
+            lastAccess = System.currentTimeMillis();
+            return getDelegate().register(name, resolution);
+        }
+
+        @Override
+        public void clear() {
+            getDelegate().clear();
         }
     }
 

apm-agent-core/src/test/java/co/elastic/apm/agent/bci/bytebuddy/SoftlyReferencingTypePoolCacheTest.java

@@ -11,9 +11,9 @@
  * the Apache License, Version 2.0 (the "License"); you may
  * not use this file except in compliance with the License.
  * You may obtain a copy of the License at
- * 
+ *
  *   http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
@@ -24,6 +24,7 @@
  */
 package co.elastic.apm.agent.bci.bytebuddy;
 
+import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatchers;
 import net.bytebuddy.pool.TypePool;
 import org.junit.jupiter.api.BeforeEach;
@@ -42,9 +43,9 @@ void setUp() {
 
     @Test
     void testClearEntries() {
-        cache.locate(ClassLoader.getSystemClassLoader());
-        assertThat(cache.getCacheProviders().approximateSize()).isEqualTo(1);
+        cache.locate(ClassLoader.getSystemClassLoader()).register(String.class.getName(), new TypePool.Resolution.Simple(TypeDescription.STRING));
+        assertThat(cache.locate(ClassLoader.getSystemClassLoader()).find(String.class.getName())).isNotNull();
         cache.clearIfNotAccessedSince(0);
-        assertThat(cache.getCacheProviders().approximateSize()).isEqualTo(0);
+        assertThat(cache.locate(ClassLoader.getSystemClassLoader()).find(String.class.getName())).isNull();
     }
 }