content: package com.alibaba.ttl.threadpool.agent;
  
  import com.alibaba.ttl.threadpool.agent.logging.Logger;
  import com.alibaba.ttl.threadpool.agent.transformlet.ClassInfo;
  import com.alibaba.ttl.threadpool.agent.transformlet.TtlTransformlet;
  import edu.umd.cs.findbugs.annotations.NonNull;
  import javassist.CannotCompileException;
  import javassist.NotFoundException;
  
  import java.io.BufferedReader;
  import java.io.IOException;
  import java.io.InputStream;
  import java.io.InputStreamReader;
  import java.lang.reflect.InvocationTargetException;
  import java.net.URL;
  import java.nio.charset.StandardCharsets;
  import java.util.*;
  
  import static com.alibaba.ttl.threadpool.agent.transformlet.helper.TtlTransformletHelper.getLocationUrlOfClass;
  import static java.nio.charset.StandardCharsets.UTF_8;
  
  /**
   * @author Jerry Lee (oldratlee at gmail dot com)
   * @since 3.0.0
   */
  final class TtlExtensionTransformletManager {
      private static final Logger logger = Logger.getLogger(TtlExtensionTransformletManager.class);
  
      private static final String TTL_AGENT_EXTENSION_TRANSFORMLET_FILE = "META-INF/ttl.agent.transformlets";
  
      public TtlExtensionTransformletManager() {
      }
  
      public String extensionTransformletDoTransform(@NonNull final ClassInfo classInfo) throws NotFoundException, CannotCompileException, IOException {
          final Map<String, TtlTransformlet> transformlets = classLoader2ExtensionTransformletsIncludeParentCL.get(classInfo.getClassLoader());
          if (transformlets == null) return null;
  
          for (Map.Entry<String, TtlTransformlet> entry : transformlets.entrySet()) {
              final String className = entry.getKey();
              final TtlTransformlet transformlet = entry.getValue();
  
              transformlet.doTransform(classInfo);
              if (classInfo.isModified()) {
                  return className;
              }
          }
  
          return null;
      }
  
      // NOTE: use WeakHashMap as a Set collection, value is always null.
      private final WeakHashMap<ClassLoader, ?> collectedClassLoaderHistory = new WeakHashMap<>(512);
  
      // Map: ExtensionTransformlet ClassLoader -> ExtensionTransformlet ClassName -> ExtensionTransformlet instance(not include from parent classloader)
      private final WeakHashMap<ClassLoader, Map<String, TtlTransformlet>> classLoader2ExtensionTransformlets =
              new WeakHashMap<>(512);
  
      // Map: ExtensionTransformlet ClassLoader -> ExtensionTransformlet ClassName -> ExtensionTransformlet instance(include from parent classloader)
      private final WeakHashMap<ClassLoader, Map<String, TtlTransformlet>> classLoader2ExtensionTransformletsIncludeParentCL =
              new WeakHashMap<>(512);
  
      public void collectExtensionTransformlet(@NonNull final ClassInfo classInfo) throws IOException {
          final ClassLoader classLoader = classInfo.getClassLoader();
          // classloader may null be if the bootstrap loader,
          // which classloader must contains NO Ttl Agent Extension Transformlet, so just safe skip
          if (classLoader == null) return;
  
          // this classLoader is collected, so skip collection
          if (collectedClassLoaderHistory.containsKey(classLoader)) return;
          collectedClassLoaderHistory.put(classLoader, null);
  
          logger.info("[TtlExtensionTransformletCollector] collecting TTL Extension Transformlets from classloader " + classLoader);
  
          final LinkedHashSet<String> extensionTransformletClassNames = readExtensionTransformletClassNames(classLoader);
  
          final String foundMsgHead = "[TtlExtensionTransformletCollector] found TTL Extension Transformlet class ";
          final String failLoadMsgHead = "[TtlExtensionTransformletCollector] fail to load TTL Extension Transformlet ";
          final Map<ClassLoader, Set<TtlTransformlet>> loadedTransformlet =
              loadExtensionInstances(classLoader, extensionTransformletClassNames, TtlTransformlet.class, foundMsgHead, failLoadMsgHead);
  
          mergeToClassLoader2ExtensionTransformlet(classLoader2ExtensionTransformlets, loadedTransformlet);
  
          updateClassLoader2ExtensionTransformletsIncludeParentCL(
              classLoader2ExtensionTransformlets, classLoader2ExtensionTransformletsIncludeParentCL);
      }
  
      // extension transformlet configuration file URL location string -> URL contained extension transformlet class names
      private final Map<String, LinkedHashSet<String>> redExtensionTransformletFileHistory = new HashMap<>();
  
      private LinkedHashSet<String> readExtensionTransformletClassNames(ClassLoader classLoader) throws IOException {
          final Enumeration<URL> extensionFiles = classLoader.getResources(TTL_AGENT_EXTENSION_TRANSFORMLET_FILE);
  
          final Pair<LinkedHashSet<String>, Set<String>> pair = readLinesFromExtensionFiles(extensionFiles, redExtensionTransformletFileHistory);
          final LinkedHashSet<String> extensionTransformletClassNames = pair.first;
          final Set<String> stringUrls = pair.second;
          if (!stringUrls.isEmpty())
              logger.info("[TtlExtensionTransformletCollector] found TTL Extension Transformlet configuration files from classloader "
                  + classLoader + " : " + stringUrls);
  
          return extensionTransformletClassNames;
      }
  
      private static void mergeToClassLoader2ExtensionTransformlet(
          Map<ClassLoader, Map<String, TtlTransformlet>> destination, Map<ClassLoader, Set<TtlTransformlet>> loadedTransformlets
      ) {
          for (Map.Entry<ClassLoader, Set<TtlTransformlet>> entry : loadedTransformlets.entrySet()) {
              final ClassLoader classLoader = entry.getKey();
              final Set<TtlTransformlet> transformlets = entry.getValue();
  
              Map<String, TtlTransformlet> className2Transformlets = destination.computeIfAbsent(classLoader, k -> new HashMap<>());
  
              for (TtlTransformlet t : transformlets) {
                  final String className = t.getClass().getName();
                  if (className2Transformlets.containsKey(className)) continue;
  
                  className2Transformlets.put(className, t);
                  logger.info("[TtlExtensionTransformletCollector] add TTL Extension Transformlet " + className + " success");
              }
          }
      }
  
      static void updateClassLoader2ExtensionTransformletsIncludeParentCL(
          Map<ClassLoader, Map<String, TtlTransformlet>> classLoader2ExtensionTransformlets,
          Map<ClassLoader, Map<String, TtlTransformlet>> classLoader2ExtensionTransformletsIncludeParentCL
      ) {
          for (Map.Entry<ClassLoader, Map<String, TtlTransformlet>> entry : classLoader2ExtensionTransformlets.entrySet()) {
              final ClassLoader classLoader = entry.getKey();
              final Map<String, TtlTransformlet> merged = childClassLoaderFirstMergeTransformlets(classLoader2ExtensionTransformlets, classLoader);
              classLoader2ExtensionTransformletsIncludeParentCL.put(classLoader, merged);
          }
      }
  
      static Map<String, TtlTransformlet> childClassLoaderFirstMergeTransformlets(
          Map<ClassLoader, Map<String, TtlTransformlet>> classLoader2Transformlet, ClassLoader classLoader
      ) {
          Map<String, TtlTransformlet> ret = new HashMap<>();
  
          final ArrayDeque<ClassLoader> chain = new ArrayDeque<>();
          chain.add(classLoader);
          while (classLoader.getParent() != null) {
              classLoader = classLoader.getParent();
  
              chain.addFirst(classLoader);
          }
  
          for (ClassLoader loader : chain) {
              final Map<String, TtlTransformlet> m = classLoader2Transformlet.get(loader);
              if (m == null) continue;
  
              ret.putAll(m);
          }
  
          return ret;
      }
  
      // ======== Extension load util methods ========
  
      static <T> Map<ClassLoader, Set<T>> loadExtensionInstances(
          ClassLoader classLoader, LinkedHashSet<String> instanceClassNames, Class<T> superType,
          String foundMsgHead, String failLoadMsgHead
      ) {
          Map<ClassLoader, Set<T>> ret = new HashMap<>();
  
          for (final String className : instanceClassNames) {
              try {
                  final Class<?> clazz = classLoader.loadClass(className);
                  if (!superType.isAssignableFrom(clazz)) {
                      final String msg = foundMsgHead + className
                          + " from classloader " + classLoader
                          + " at location " + getLocationUrlOfClass(clazz)
                          + ", but NOT subtype of " + superType.getName() + ", ignored!";
                      logger.error(msg);
                      continue;
                  }
  
                  Object instance = clazz.getDeclaredConstructor().newInstance();
  
                  final ClassLoader actualClassLoader = instance.getClass().getClassLoader();
                  Set<T> set = ret.computeIfAbsent(actualClassLoader, k -> new HashSet<>());
                  set.add(superType.cast(instance));
  
                  final String msg = foundMsgHead + className
                      + ", and loaded from classloader " + classLoader
                      + " at location " + getLocationUrlOfClass(clazz);
                  logger.info(msg);
              } catch (ClassNotFoundException e) {
                  final String msg = failLoadMsgHead + className + " from classloader " + classLoader + ", cause: " + e.toString();
                  logger.warn(msg, e);
              } catch (IllegalAccessException | InstantiationException | NoSuchMethodException |
                       InvocationTargetException e) {
                  final String msg = failLoadMsgHead + className + " from classloader " + classLoader + ", cause: " + e.toString();
                  logger.error(msg, e);
              }
          }
  
          return ret;
      }
  
      // return: read lines from URL, url strings
      @NonNull
      static Pair<LinkedHashSet<String>, Set<String>> readLinesFromExtensionFiles(
          /* input */ @NonNull Enumeration<URL> extensionFiles,
          /* input/output, map url string -> content lines */ @NonNull Map<String, LinkedHashSet<String>> redExtensionFilesHistory
      ) {
          final LinkedHashSet<String> mergedLines = new LinkedHashSet<>();
          final Set<String> stringUrls = new HashSet<>();
  
          while (extensionFiles.hasMoreElements()) {
              final URL url = extensionFiles.nextElement();
  
              final String urlString = url.toString();
              stringUrls.add(urlString);
  
              LinkedHashSet<String> lines;
              if (redExtensionFilesHistory.containsKey(urlString)) {
                  lines = redExtensionFilesHistory.get(urlString);
              } else {
                  lines = readLines(url);
  
                  redExtensionFilesHistory.put(urlString, lines);
              }
  
              mergedLines.addAll(lines);
          }
  
          return new Pair<>(mergedLines, stringUrls);
      }
  
      /**
       * this method is modified based on {@link java.util.ServiceLoader}
       */
      @SuppressWarnings("StatementWithEmptyBody")
      static LinkedHashSet<String> readLines(URL extensionFile) {
          InputStream inputStream = null;
          BufferedReader reader = null;
  
          LinkedHashSet<String> names = new LinkedHashSet<>();
          try {
              inputStream = extensionFile.openStream();
              reader = new BufferedReader(new InputStreamReader(inputStream, UTF_8));
              int lineNum = 1;
              while ((lineNum = parseLine(extensionFile, reader, lineNum, names)) >= 0) ;
          } catch (IOException x) {
              logger.error("Error reading configuration file " + extensionFile, x);
          } finally {
              try {
                  if (reader != null) reader.close();
                  if (inputStream != null) inputStream.close();
              } catch (IOException y) {
                  logger.warn("Error closing configuration file " + extensionFile, y);
              }
          }
  
          return names;
      }
  
  
      /**
       * this method is modified based on {@link java.util.ServiceLoader}
       */
      private static int parseLine(URL url, BufferedReader reader, int lineNum, LinkedHashSet<String> names) throws IOException {
          String line = reader.readLine();
          if (line == null) {
              return -1;
          }
  
          // remove comments that start with `#`
          int ci = line.indexOf('#');
          if (ci >= 0) line = line.substring(0, ci);
  
          line = line.trim();
  
          int n = line.length();
          if (n != 0) {
              if ((line.indexOf(' ') >= 0) || (line.indexOf('\t') >= 0)) {
                  logger.error("Illegal syntax " + line + "in configuration file" + url + ", contains space or tab; ignore this line!");
                  return lineNum + 1;
              }
  
              int cp = line.codePointAt(0);
              if (!Character.isJavaIdentifierStart(cp)) {
                  logger.error("Illegal extension class name " + line + " in configuration file " + url + "; ignore this line!");
                  return lineNum + 1;
              }
              for (int i = Character.charCount(cp); i < n; i += Character.charCount(cp)) {
                  cp = line.codePointAt(i);
                  if (!Character.isJavaIdentifierPart(cp) && (cp != '.')) {
                      logger.error("Illegal extension class name: " + line + " in configuration file " + url + "; ignore this line!");
                      return lineNum + 1;
                  }
              }
  
              names.add(line);
          }
  
          return lineNum + 1;
      }
  
      static class Pair<T, U> {
          T first;
          U second;
  
          public Pair(T first, U second) {
              this.first = first;
              this.second = second;
          }
      }
  }
  