From 1caf251db08fd8cdfdfe98913ab706129bc15582 Mon Sep 17 00:00:00 2001 From: Lauri Tulmin Date: Fri, 15 Dec 2023 01:04:25 +0200 Subject: [PATCH] Associate value with class loader (#10051) --- .../tooling/util/ClassLoaderMap.java | 83 +++++++++++++++++++ .../tooling/util/ClassLoaderValue.java | 24 ++++++ .../tooling/util/ClassLoaderValueTest.java | 51 ++++++++++++ 3 files changed, 158 insertions(+) create mode 100644 javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderMap.java create mode 100644 javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValue.java create mode 100644 javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValueTest.java diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderMap.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderMap.java new file mode 100644 index 000000000000..9da9c3ac6fa3 --- /dev/null +++ b/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>> 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 getClassLoaderData( + ClassLoader classLoader, boolean initialize) { + classLoader = maskNullClassLoader(classLoader); + WeakReference> weakReference = data.get(classLoader); + Map 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 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 map; + try { + Field field = clazz.getField("data"); + synchronized (classLoader) { + map = (Map) 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() {} +} diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValue.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValue.java new file mode 100644 index 000000000000..04df8da95615 --- /dev/null +++ b/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 { + + @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); + } +} diff --git a/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValueTest.java b/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValueTest.java new file mode 100644 index 000000000000..799248f8db8d --- /dev/null +++ b/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 classLoaderValue = new ClassLoaderValue<>(); + classLoaderValue.put(classLoader, "value"); + assertThat(classLoaderValue.get(classLoader)).isEqualTo("value"); + } + + @Test + void testGc() throws InterruptedException { + ClassLoader testClassLoader = new ClassLoader() {}; + ClassLoaderValue classLoaderValue = new ClassLoaderValue<>(); + Value value = new Value(); + classLoaderValue.put(testClassLoader, value); + WeakReference valueWeakReference = new WeakReference<>(value); + WeakReference 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 {} +}