Detect GC leaks of scopes in StrictContextStorage. (#2164)

* Detect GC leaks of scopes in StrictContextStorage.

* More

* Finish

* Force GC more aggressively

* Cleanup

* Vendor code directly

* Copy test too

* Try waiting more

* ep

* oops

* Remove from build.gradle

* Drift

* Log on multiple

* Cleaner ourselves.

* EP

* Move into if

* Revert accidental

Author: Anuraag Agrawal

api/build.gradle

@@ -19,7 +19,3 @@ dependencies {
     testImplementation libraries.jqf,
             libraries.guava_testlib
 }
-
-javadoc {
-    exclude 'io/opentelemetry/internal/**'
-}

build.gradle

@@ -200,6 +200,10 @@ configure(opentelemetryProjects) {
         withSourcesJar()
     }
 
+    javadoc {
+        exclude 'io/opentelemetry/internal/**'
+    }
+
     tasks {
         def testJava8 = register('testJava8', Test) {
             javaLauncher = javaToolchains.launcherFor {

context/src/main/java/io/opentelemetry/context/internal/shaded/AbstractWeakConcurrentMap.java

@@ -0,0 +1,374 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Includes work from:
+/*
+ * Copyright Rafael Winterhalter
+ *
+ * Licensed 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.
+ */
+
+// Suppress warnings since this is vendored as-is.
+// CHECKSTYLE:OFF
+
+package io.opentelemetry.context.internal.shaded;
+
+import java.lang.ref.Reference;
+import java.lang.ref.ReferenceQueue;
+import java.lang.ref.WeakReference;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+/**
+ * A thread-safe map with weak keys. Entries are based on a key's system hash code and keys are
+ * considered equal only by reference equality. This class offers an abstract-base implementation
+ * that allows to override methods. This class does not implement the {@link Map} interface because
+ * this implementation is incompatible with the map contract. While iterating over a map's entries,
+ * any key that has not passed iteration is referenced non-weakly.
+ *
+ * <p>This class has been copied as is from
+ * https://github.com/raphw/weak-lock-free/blob/ad0e5e0c04d4a31f9485bf12b89afbc9d75473b3/src/main/java/com/blogspot/mydailyjava/weaklockfree/WeakConcurrentMap.java
+ * This is used in multiple artifacts in OpenTelemetry and while it is in our internal API,
+ * generally backwards compatible changes should not be made to avoid a situation where different
+ * versions of OpenTelemetry artifacts become incompatible with each other.
+ */
+// Suppress warnings since this is vendored as-is.
+@SuppressWarnings({"MissingSummary", "EqualsBrokenForNull", "FieldMissingNullable"})
+public abstract class AbstractWeakConcurrentMap<K, V, L> extends ReferenceQueue<K>
+    implements Runnable, Iterable<Map.Entry<K, V>> {
+
+  final ConcurrentMap<WeakKey<K>, V> target;
+
+  protected AbstractWeakConcurrentMap() {
+    this(new ConcurrentHashMap<WeakKey<K>, V>());
+  }
+
+  /** @param target ConcurrentMap implementation that this class wraps. */
+  protected AbstractWeakConcurrentMap(ConcurrentMap<WeakKey<K>, V> target) {
+    this.target = target;
+  }
+
+  /**
+   * Override with care as it can cause lookup failures if done incorrectly. The result must have
+   * the same {@link Object#hashCode()} as the input and be {@link Object#equals(Object) equal to} a
+   * weak reference of the key. When overriding this, also override {@link #resetLookupKey}.
+   */
+  protected abstract L getLookupKey(K key);
+
+  /** Resets any reusable state in the {@linkplain #getLookupKey lookup key}. */
+  protected abstract void resetLookupKey(L lookupKey);
+
+  /**
+   * @param key The key of the entry.
+   * @return The value of the entry or the default value if it did not exist.
+   */
+  public V get(K key) {
+    if (key == null) throw new NullPointerException();
+    V value;
+    L lookupKey = getLookupKey(key);
+    try {
+      value = target.get(lookupKey);
+    } finally {
+      resetLookupKey(lookupKey);
+    }
+    if (value == null) {
+      value = defaultValue(key);
+      if (value != null) {
+        V previousValue = target.putIfAbsent(new WeakKey<K>(key, this), value);
+        if (previousValue != null) {
+          value = previousValue;
+        }
+      }
+    }
+    return value;
+  }
+
+  /**
+   * @param key The key of the entry.
+   * @return The value of the entry or null if it did not exist.
+   */
+  public V getIfPresent(K key) {
+    if (key == null) throw new NullPointerException();
+    L lookupKey = getLookupKey(key);
+    try {
+      return target.get(lookupKey);
+    } finally {
+      resetLookupKey(lookupKey);
+    }
+  }
+
+  /**
+   * @param key The key of the entry.
+   * @return {@code true} if the key already defines a value.
+   */
+  public boolean containsKey(K key) {
+    if (key == null) throw new NullPointerException();
+    L lookupKey = getLookupKey(key);
+    try {
+      return target.containsKey(lookupKey);
+    } finally {
+      resetLookupKey(lookupKey);
+    }
+  }
+
+  /**
+   * @param key The key of the entry.
+   * @param value The value of the entry.
+   * @return The previous entry or {@code null} if it does not exist.
+   */
+  public V put(K key, V value) {
+    if (key == null || value == null) throw new NullPointerException();
+    return target.put(new WeakKey<K>(key, this), value);
+  }
+
+  /**
+   * @param key The key of the entry.
+   * @param value The value of the entry.
+   * @return The previous entry or {@code null} if it does not exist.
+   */
+  public V putIfAbsent(K key, V value) {
+    if (key == null || value == null) throw new NullPointerException();
+    V previous;
+    L lookupKey = getLookupKey(key);
+    try {
+      previous = target.get(lookupKey);
+    } finally {
+      resetLookupKey(lookupKey);
+    }
+    return previous == null ? target.putIfAbsent(new WeakKey<K>(key, this), value) : previous;
+  }
+
+  /**
+   * @param key The key of the entry.
+   * @param value The value of the entry.
+   * @return The previous entry or {@code null} if it does not exist.
+   */
+  public V putIfProbablyAbsent(K key, V value) {
+    if (key == null || value == null) throw new NullPointerException();
+    return target.putIfAbsent(new WeakKey<K>(key, this), value);
+  }
+
+  /**
+   * @param key The key of the entry.
+   * @return The removed entry or {@code null} if it does not exist.
+   */
+  public V remove(K key) {
+    if (key == null) throw new NullPointerException();
+    L lookupKey = getLookupKey(key);
+    try {
+      return target.remove(lookupKey);
+    } finally {
+      resetLookupKey(lookupKey);
+    }
+  }
+
+  /** Clears the entire map. */
+  public void clear() {
+    target.clear();
+  }
+
+  /**
+   * Creates a default value. There is no guarantee that the requested value will be set as a once
+   * it is created in case that another thread requests a value for a key concurrently.
+   *
+   * @param key The key for which to create a default value.
+   * @return The default value for a key without value or {@code null} for not defining a default
+   *     value.
+   */
+  protected V defaultValue(K key) {
+    return null;
+  }
+
+  /** Cleans all unused references. */
+  public void expungeStaleEntries() {
+    Reference<?> reference;
+    while ((reference = poll()) != null) {
+      target.remove(reference);
+    }
+  }
+
+  /**
+   * Returns the approximate size of this map where the returned number is at least as big as the
+   * actual number of entries.
+   *
+   * @return The minimum size of this map.
+   */
+  public int approximateSize() {
+    return target.size();
+  }
+
+  @Override
+  public void run() {
+    try {
+      while (!Thread.interrupted()) {
+        target.remove(remove());
+      }
+    } catch (InterruptedException ignored) {
+      // do nothing
+    }
+  }
+
+  @Override
+  public Iterator<Map.Entry<K, V>> iterator() {
+    return new EntryIterator(target.entrySet().iterator());
+  }
+
+  @Override
+  public String toString() {
+    return target.toString();
+  }
+
+  /*
+   * Why this works:
+   * ---------------
+   *
+   * Note that this map only supports reference equality for keys and uses system hash codes. Also, for the
+   * WeakKey instances to function correctly, we are voluntarily breaking the Java API contract for
+   * hashCode/equals of these instances.
+   *
+   * System hash codes are immutable and can therefore be computed prematurely and are stored explicitly
+   * within the WeakKey instances. This way, we always know the correct hash code of a key and always
+   * end up in the correct bucket of our target map. This remains true even after the weakly referenced
+   * key is collected.
+   *
+   * If we are looking up the value of the current key via WeakConcurrentMap::get or any other public
+   * API method, we know that any value associated with this key must still be in the map as the mere
+   * existence of this key makes it ineligible for garbage collection. Therefore, looking up a value
+   * using another WeakKey wrapper guarantees a correct result.
+   *
+   * If we are looking up the map entry of a WeakKey after polling it from the reference queue, we know
+   * that the actual key was already collected and calling WeakKey::get returns null for both the polled
+   * instance and the instance within the map. Since we explicitly stored the identity hash code for the
+   * referenced value, it is however trivial to identify the correct bucket. From this bucket, the first
+   * weak key with a null reference is removed. Due to hash collision, we do not know if this entry
+   * represents the weak key. However, we do know that the reference queue polls at least as many weak
+   * keys as there are stale map entries within the target map. If no key is ever removed from the map
+   * explicitly, the reference queue eventually polls exactly as many weak keys as there are stale entries.
+   *
+   * Therefore, we can guarantee that there is no memory leak.
+   *
+   * It is the responsibility of the actual map implementation to implement a lookup key that is used for
+   * lookups. The lookup key must supply the same semantics as the weak key with regards to hash code.
+   * The weak key invokes the latent key's equality method upon evaluation.
+   */
+
+  public static final class WeakKey<K> extends WeakReference<K> {
+
+    private final int hashCode;
+
+    WeakKey(K key, ReferenceQueue<? super K> queue) {
+      super(key, queue);
+      hashCode = System.identityHashCode(key);
+    }
+
+    @Override
+    public int hashCode() {
+      return hashCode;
+    }
+
+    @Override
+    public boolean equals(Object other) {
+      if (other instanceof WeakKey<?>) {
+        return ((WeakKey<?>) other).get() == get();
+      } else {
+        return other.equals(this);
+      }
+    }
+
+    @Override
+    public String toString() {
+      return String.valueOf(get());
+    }
+  }
+
+  private class EntryIterator implements Iterator<Map.Entry<K, V>> {
+
+    private final Iterator<Map.Entry<WeakKey<K>, V>> iterator;
+
+    private Map.Entry<WeakKey<K>, V> nextEntry;
+
+    private K nextKey;
+
+    private EntryIterator(Iterator<Map.Entry<WeakKey<K>, V>> iterator) {
+      this.iterator = iterator;
+      findNext();
+    }
+
+    private void findNext() {
+      while (iterator.hasNext()) {
+        nextEntry = iterator.next();
+        nextKey = nextEntry.getKey().get();
+        if (nextKey != null) {
+          return;
+        }
+      }
+      nextEntry = null;
+      nextKey = null;
+    }
+
+    @Override
+    public boolean hasNext() {
+      return nextKey != null;
+    }
+
+    @Override
+    public Map.Entry<K, V> next() {
+      if (nextKey == null) {
+        throw new NoSuchElementException();
+      }
+      try {
+        return new SimpleEntry(nextKey, nextEntry);
+      } finally {
+        findNext();
+      }
+    }
+
+    @Override
+    public void remove() {
+      throw new UnsupportedOperationException();
+    }
+  }
+
+  private class SimpleEntry implements Map.Entry<K, V> {
+
+    private final K key;
+
+    final Map.Entry<WeakKey<K>, V> entry;
+
+    private SimpleEntry(K key, Map.Entry<WeakKey<K>, V> entry) {
+      this.key = key;
+      this.entry = entry;
+    }
+
+    @Override
+    public K getKey() {
+      return key;
+    }
+
+    @Override
+    public V getValue() {
+      return entry.getValue();
+    }
+
+    @Override
+    public V setValue(V value) {
+      if (value == null) throw new NullPointerException();
+      return entry.setValue(value);
+    }
+  }
+}