io.opentelemetry.javaagent.tooling.util.ClassLoaderMap Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of opentelemetry-javaagent-tooling Show documentation

Instrumentation of Java libraries using OpenTelemetry.

There is a newer version: 2.10.0-alpha

Show newest version

/*
 * 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 java.util.function.Supplier;
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);
  }

  public static Object computeIfAbsent(
      ClassLoader classLoader, Object key, Supplier value) {
    return getClassLoaderData(classLoader, true).computeIfAbsent(key, unused -> value.get());
  }

  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() {}
}