Associate value with class loader (#10051)

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderMap.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.util;
+
+import io.opentelemetry.instrumentation.api.internal.cache.Cache;
+import java.lang.ref.WeakReference;
+import java.lang.reflect.Field;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.description.modifier.Ownership;
+import net.bytebuddy.description.modifier.Visibility;
+import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
+
+class ClassLoaderMap {
+  private static final Cache<ClassLoader, WeakReference<Map<Object, Object>>> data = Cache.weak();
+
+  public static Object get(ClassLoader classLoader, Object key) {
+    return getClassLoaderData(classLoader, false).get(key);
+  }
+
+  public static void put(ClassLoader classLoader, Object key, Object value) {
+    getClassLoaderData(classLoader, true).put(key, value);
+  }
+
+  private static Map<Object, Object> getClassLoaderData(
+      ClassLoader classLoader, boolean initialize) {
+    classLoader = maskNullClassLoader(classLoader);
+    WeakReference<Map<Object, Object>> weakReference = data.get(classLoader);
+    Map<Object, Object> map = weakReference != null ? weakReference.get() : null;
+    if (map == null) {
+      // skip setting up the map if get was called
+      if (!initialize) {
+        return Collections.emptyMap();
+      }
+      map = createMap(classLoader);
+      data.put(classLoader, new WeakReference<>(map));
+    }
+    return map;
+  }
+
+  @SuppressWarnings("unchecked")
+  private static Map<Object, Object> createMap(ClassLoader classLoader) {
+    // generate a class with a single static field named "data" and define it in the given class
+    // loader
+    Class<?> clazz =
+        new ByteBuddy()
+            .subclass(Object.class)
+            .name(
+                "io.opentelemetry.javaagent.ClassLoaderData$$"
+                    + Integer.toHexString(System.identityHashCode(classLoader)))
+            .defineField("data", Object.class, Ownership.STATIC, Visibility.PUBLIC)
+            .make()
+            .load(classLoader, ClassLoadingStrategy.Default.INJECTION.allowExistingTypes())
+            .getLoaded();
+    Map<Object, Object> map;
+    try {
+      Field field = clazz.getField("data");
+      synchronized (classLoader) {
+        map = (Map<Object, Object>) field.get(classLoader);
+        if (map == null) {
+          map = new ConcurrentHashMap<>();
+          field.set(null, map);
+        }
+      }
+    } catch (Exception exception) {
+      throw new IllegalStateException(exception);
+    }
+    return map;
+  }
+
+  private static final ClassLoader BOOT_LOADER = new ClassLoader(null) {};
+
+  private static ClassLoader maskNullClassLoader(ClassLoader classLoader) {
+    return classLoader == null ? BOOT_LOADER : classLoader;
+  }
+
+  private ClassLoaderMap() {}
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValue.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.util;
+
+/**
+ * Associate value with a class loader. Added value will behave as if it was stored in a field in
+ * the class loader object, meaning that the value can be garbage collected once the class loader is
+ * garbage collected and referencing the class loader from the value will not prevent garbage
+ * collector from collecting the class loader.
+ */
+public final class ClassLoaderValue<T> {
+
+  @SuppressWarnings("unchecked")
+  public T get(ClassLoader classLoader) {
+    return (T) ClassLoaderMap.get(classLoader, this);
+  }
+
+  public void put(ClassLoader classLoader, T value) {
+    ClassLoaderMap.put(classLoader, this, value);
+  }
+}

javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValueTest.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.util;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.instrumentation.test.utils.GcUtils;
+import java.lang.ref.WeakReference;
+import org.junit.jupiter.api.Test;
+
+class ClassLoaderValueTest {
+
+  @Test
+  void testValue() {
+    testClassLoader(this.getClass().getClassLoader());
+    testClassLoader(null);
+  }
+
+  void testClassLoader(ClassLoader classLoader) {
+    ClassLoaderValue<String> classLoaderValue = new ClassLoaderValue<>();
+    classLoaderValue.put(classLoader, "value");
+    assertThat(classLoaderValue.get(classLoader)).isEqualTo("value");
+  }
+
+  @Test
+  void testGc() throws InterruptedException {
+    ClassLoader testClassLoader = new ClassLoader() {};
+    ClassLoaderValue<Value> classLoaderValue = new ClassLoaderValue<>();
+    Value value = new Value();
+    classLoaderValue.put(testClassLoader, value);
+    WeakReference<Value> valueWeakReference = new WeakReference<>(value);
+    WeakReference<ClassLoader> classLoaderWeakReference = new WeakReference<>(testClassLoader);
+
+    assertThat(classLoaderWeakReference.get()).isNotNull();
+    assertThat(valueWeakReference.get()).isNotNull();
+
+    value = null;
+    testClassLoader = null;
+
+    GcUtils.awaitGc(classLoaderWeakReference);
+    GcUtils.awaitGc(valueWeakReference);
+
+    assertThat(classLoaderWeakReference.get()).isNull();
+    assertThat(valueWeakReference.get()).isNull();
+  }
+
+  private static class Value {}
+}