null-safe weak set/map + detached thread local (#1597)

* add null-safe weak map
* add null-safe weak set
* use null-safe weak maps/sets when possible
* prevent npe with threadlocal

Author: SylvainJuge

CHANGELOG.asciidoc

@@ -45,6 +45,7 @@ endif::[]
 * Fix improper request body capturing - {pull}1579[#1579]
 * avoid `NullPointerException` due to null return values instrumentation advices - {pull}1601[#1601]
 * Update async-profiler to 1.8.3 {pull}1602[1602]
+* Use null-safe data structures to avoid `NullPointerException` {pull}1597[1597]
 
 [float]
 ===== Refactors

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

@@ -26,6 +26,7 @@
 
 import co.elastic.apm.agent.bci.classloading.IndyPluginClassLoader;
 import co.elastic.apm.agent.impl.ElasticApmTracer;
+import co.elastic.apm.agent.sdk.weakmap.NullSafeWeakConcurrentMap;
 import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.ClassFileLocator;
@@ -223,7 +224,7 @@ public static class ForAnyClassLoader<T> extends HelperClassManager<T> {
         private ForAnyClassLoader(ElasticApmTracer tracer, String implementation, String... additionalHelpers) {
             super(tracer, implementation, additionalHelpers);
             // deliberately doesn't use WeakMapSupplier as this class manages the cleanup manually
-            clId2helperMap = new WeakConcurrentMap<>(false);
+            clId2helperMap = new NullSafeWeakConcurrentMap<>(false);
         }
 
         /**

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

@@ -24,6 +24,7 @@
  */
 package co.elastic.apm.agent.bci.bytebuddy;
 
+import co.elastic.apm.agent.sdk.weakmap.NullSafeWeakConcurrentMap;
 import co.elastic.apm.agent.util.ExecutorUtils;
 import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 import net.bytebuddy.agent.builder.AgentBuilder;
@@ -56,7 +57,7 @@ public class SoftlyReferencingTypePoolCache extends AgentBuilder.PoolStrategy.Wi
      * deliberately doesn't use WeakMapSupplier as this class manages the cleanup manually
      */
     private final WeakConcurrentMap<ClassLoader, CacheProviderWrapper> cacheProviders =
-        new WeakConcurrentMap<ClassLoader, CacheProviderWrapper>(false);
+        new NullSafeWeakConcurrentMap<ClassLoader, CacheProviderWrapper>(false);
     private final ElementMatcher<ClassLoader> ignoredClassLoaders;
 
     public SoftlyReferencingTypePoolCache(final TypePool.Default.ReaderMode readerMode,

apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/state/GlobalThreadLocal.java

@@ -24,6 +24,7 @@
  */
 package co.elastic.apm.agent.sdk.state;
 
+import co.elastic.apm.agent.sdk.weakmap.NullCheck;
 import com.blogspot.mydailyjava.weaklockfree.DetachedThreadLocal;
 
 import javax.annotation.Nullable;
@@ -82,6 +83,14 @@ public T get(T defaultValue) {
         return defaultValue;
     }
 
+    @Override
+    public void set(@Nullable T value) {
+        if (NullCheck.isNullKey(value)) {
+            return;
+        }
+        super.set(value);
+    }
+
     @Override
     @Nullable
     protected T initialValue(Thread thread) {

apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/weakmap/NullCheck.java

@@ -0,0 +1,38 @@
+package co.elastic.apm.agent.sdk.weakmap;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.Nullable;
+
+public class NullCheck {
+
+    private static final Logger logger = LoggerFactory.getLogger(NullCheck.class);
+
+    /**
+     * checks if key or value is {@literal null}
+     *
+     * @param v key or value
+     * @return {@literal true} if key is non-null, {@literal false} if null
+     */
+    private static <T> boolean isNull(@Nullable T v, boolean isKey) {
+        if (null != v) {
+            return false;
+        }
+        String msg = String.format("trying to use null %s", isKey ? "key" : "value");
+        if (logger.isDebugEnabled()) {
+            logger.debug(msg, new RuntimeException(msg));
+        } else {
+            logger.warn(msg);
+        }
+        return true;
+    }
+
+    public static <T> boolean isNullKey(@Nullable T key){
+        return isNull(key, true);
+    }
+
+    public static <T> boolean isNullValue(@Nullable T value) {
+        return isNull(value, false);
+    }
+}

apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/weakmap/NullSafeWeakConcurrentMap.java

@@ -0,0 +1,109 @@
+/*-
+ * #%L
+ * Elastic APM Java agent
+ * %%
+ * Copyright (C) 2018 - 2021 Elastic and contributors
+ * %%
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * 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
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ * #L%
+ */
+package co.elastic.apm.agent.sdk.weakmap;
+
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
+
+import javax.annotation.Nullable;
+
+import static co.elastic.apm.agent.sdk.weakmap.NullCheck.isNullKey;
+import static co.elastic.apm.agent.sdk.weakmap.NullCheck.isNullValue;
+
+/**
+ * {@link WeakConcurrentMap} implementation that prevents throwing {@link NullPointerException} and helps debugging if needed
+ *
+ * @param <K> key type
+ * @param <V> value type
+ */
+public class NullSafeWeakConcurrentMap<K, V> extends WeakConcurrentMap<K, V> {
+
+    public NullSafeWeakConcurrentMap(boolean cleanerThread) {
+        super(cleanerThread);
+    }
+
+    @Nullable
+    @Override
+    public V get(K key) {
+        if (isNullKey(key)) {
+            // super implementation silently adds entries from default value when there is none
+            // in the case of 'null', we won't return the default value nor create a map entry with it.
+            return null;
+        }
+        return super.get(key);
+    }
+
+    @Nullable
+    @Override
+    public V getIfPresent(K key) {
+        if (isNullKey(key)) {
+            return null;
+        }
+        return super.getIfPresent(key);
+    }
+
+    @Override
+    public boolean containsKey(K key) {
+        if (isNullKey(key)) {
+            return false;
+        }
+        return super.containsKey(key);
+    }
+
+    @Nullable
+    @Override
+    public V put(K key, V value) {
+        if (isNullKey(key) || isNullValue(value)) {
+            return null;
+        }
+        return super.put(key, value);
+    }
+
+    @Nullable
+    @Override
+    public V putIfAbsent(K key, V value) {
+        if (isNullKey(key) || isNullValue(value)) {
+            return null;
+        }
+        return super.putIfAbsent(key, value);
+    }
+
+    @Nullable
+    @Override
+    public V putIfProbablyAbsent(K key, V value) {
+        if (isNullKey(key) || isNullValue(value)) {
+            return null;
+        }
+        return super.putIfProbablyAbsent(key, value);
+    }
+
+    @Nullable
+    @Override
+    public V remove(K key) {
+        if (isNullKey(key)) {
+            return null;
+        }
+        return super.remove(key);
+    }
+}

apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/weakmap/NullSafeWeakConcurrentSet.java

@@ -0,0 +1,41 @@
+package co.elastic.apm.agent.sdk.weakmap;
+
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentSet;
+
+import static co.elastic.apm.agent.sdk.weakmap.NullCheck.isNullValue;
+
+/**
+ * {@link WeakConcurrentSet} implementation that prevents throwing {@link NullPointerException} and helps debugging if needed
+ *
+ * @param <V> value type
+ */
+public class NullSafeWeakConcurrentSet<V> extends WeakConcurrentSet<V> {
+
+    public NullSafeWeakConcurrentSet(Cleaner cleaner) {
+        super(cleaner);
+    }
+
+    @Override
+    public boolean add(V value) {
+        if(isNullValue(value)){
+            return false;
+        }
+        return super.add(value);
+    }
+
+    @Override
+    public boolean contains(V value) {
+        if(isNullValue(value)){
+            return false;
+        }
+        return super.contains(value);
+    }
+
+    @Override
+    public boolean remove(V value) {
+        if(isNullValue(value)){
+            return false;
+        }
+        return super.remove(value);
+    }
+}

apm-agent-plugin-sdk/src/main/java/co/elastic/apm/agent/sdk/weakmap/WeakMapSupplier.java

@@ -36,13 +36,13 @@ public class WeakMapSupplier {
     private static final WeakConcurrentSet<WeakConcurrentSet<?>> registeredSets = new WeakConcurrentSet<>(WeakConcurrentSet.Cleaner.INLINE);
 
     public static <K, V> WeakConcurrentMap<K, V> createMap() {
-        WeakConcurrentMap<K, V> result = new WeakConcurrentMap<>(false);
+        WeakConcurrentMap<K, V> result = new NullSafeWeakConcurrentMap<>(false);
         registeredMaps.add(result);
         return result;
     }
 
     public static <V> WeakConcurrentSet<V> createSet() {
-        WeakConcurrentSet<V> weakSet = new WeakConcurrentSet<V>(WeakConcurrentSet.Cleaner.MANUAL);
+        WeakConcurrentSet<V> weakSet = new NullSafeWeakConcurrentSet<>(WeakConcurrentSet.Cleaner.MANUAL);
         registeredSets.add(weakSet);
         return weakSet;
     }

apm-agent-plugin-sdk/src/test/java/co/elastic/apm/agent/sdk/state/GlobalThreadLocalTest.java

@@ -0,0 +1,13 @@
+package co.elastic.apm.agent.sdk.state;
+
+import org.junit.jupiter.api.Test;
+
+public class GlobalThreadLocalTest {
+
+    GlobalThreadLocal<String> threadLocal = GlobalThreadLocal.get(GlobalThreadLocalTest.class, "test");
+
+    @Test
+    void setNullValueShouldNotThrow() {
+        threadLocal.set(null);
+    }
+}