Refactor Servlet API and JDBC plugins to use a MethodHandles-based advice dispatching approach

apm-agent-benchmarks/src/main/java/co/elastic/apm/agent/benchmark/MockHttpServletRequest.java

@@ -11,9 +11,9 @@
  * 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

apm-agent-benchmarks/src/main/java/co/elastic/apm/agent/benchmark/MockHttpServletResponse.java

@@ -11,9 +11,9 @@
  * 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

apm-agent-core/src/main/java/co/elastic/apm/agent/bci/ElasticApmAgent.java

@@ -389,6 +389,7 @@ public static synchronized void reset() {
             transformer.reset(instrumentation, RedefinitionStrategy.RETRANSFORMATION);
         }
         dynamicClassFileTransformers.clear();
+        HelperClassManager.ForDispatcher.clear();
         instrumentation = null;
     }
 

apm-agent-core/src/main/java/co/elastic/apm/agent/bci/ElasticApmInstrumentation.java

@@ -11,9 +11,9 @@
  * 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
@@ -37,6 +37,8 @@
 import javax.annotation.Nullable;
 import java.security.ProtectionDomain;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
 
 import static net.bytebuddy.matcher.ElementMatchers.any;
 
@@ -169,6 +171,14 @@ public Advice.OffsetMapping.Factory<?> getOffsetMapping() {
         return null;
     }
 
-    public void onTypeMatch(TypeDescription typeDescription, ClassLoader classLoader, ProtectionDomain protectionDomain, @Nullable Class<?> classBeingRedefined) {
+    public List<String> helpers() throws Exception {
+        return Collections.emptyList();
+    }
+
+    public void onTypeMatch(TypeDescription typeDescription, ClassLoader classLoader, ProtectionDomain protectionDomain, @Nullable Class<?> classBeingRedefined) throws Exception {
+        List<String> helpers = helpers();
+        if (!helpers.isEmpty()) {
+            HelperClassManager.ForDispatcher.inject(classLoader, protectionDomain, helpers);
+        }
     }
 }

apm-agent-core/src/main/java/co/elastic/apm/agent/bci/HelperClassManager.java

@@ -24,27 +24,42 @@
  */
 package co.elastic.apm.agent.bci;
 
+import co.elastic.apm.agent.bootstrap.MethodHandleDispatcher;
 import co.elastic.apm.agent.impl.ElasticApmTracer;
 import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.dynamic.loading.ByteArrayClassLoader;
 import net.bytebuddy.dynamic.loading.ClassInjector;
+import net.bytebuddy.dynamic.loading.MultipleParentClassLoader;
 import net.bytebuddy.dynamic.loading.PackageDefinitionStrategy;
+import net.bytebuddy.pool.TypePool;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import javax.annotation.Nullable;
 import java.io.IOException;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
 import java.lang.ref.WeakReference;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
 import java.security.ProtectionDomain;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.WeakHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
 
 /**
  * This class helps to overcome the fact that the agent classes can't access the classes they want to instrument.
@@ -142,7 +157,9 @@ public T getForClassLoaderOfClass(Class<?> classOfTargetClassLoader) {
      * </p>
      *
      * @param <T> the type of the helper class interface
+     * @deprecated In favor of {@link ForDispatcher}
      */
+    @Deprecated
     public static class ForSingleClassLoader<T> extends HelperClassManager<T> {
 
         @Nullable
@@ -199,7 +216,9 @@ public T doGetForClassLoaderOfClass(Class<?> classOfTargetClassLoader) {
      * startup time and when new applications are deployed). These maps shouldn't grow big as they have an entry per class loader.
      *
      * @param <T>
+     * @deprecated In favor of {@link ForDispatcher}
      */
+    @Deprecated
     public static class ForAnyClassLoader<T> extends HelperClassManager<T> {
 
         // doesn't need to be concurrent - invoked only from a synchronized context
@@ -271,7 +290,148 @@ private synchronized T loadAndReferenceHelper(Class<?> classOfTargetClassLoader)
         }
     }
 
-    static Class injectClass(@Nullable ClassLoader targetClassLoader, @Nullable ProtectionDomain pd, String className, boolean isBootstrapClass) throws IOException, ClassNotFoundException {
+    public static class ForDispatcher {
+
+        private static final Map<ClassLoader, Set<Collection<String>>> alreadyInjected = new WeakHashMap<ClassLoader, Set<Collection<String>>>();
+
+        /**
+         * Creates an isolated CL that has two parents: the target class loader and the agent CL.
+         * The agent class loader is currently the bootstrap CL but in the future it will be an isolated CL that is a child of the bootstrap CL.
+         * <p>
+         * After the helper CL is created, it registers {@link RegisterMethodHandle}-annotated methods in {@link MethodHandleDispatcher}.
+         * These method handles are then called from within advices ({@link net.bytebuddy.asm.Advice.OnMethodEnter}/{@link net.bytebuddy.asm.Advice.OnMethodExit}).
+         * This lets them call the helpers that are loaded from the helper CL.
+         * The helpers have full access to both the agent API and the types visible to the target CL (such as the servlet API).
+         * </p>
+         * <p>
+         * See {@link MethodHandleDispatcher} for a diagram.
+         * </p>
+         */
+        public synchronized static void inject(@Nullable ClassLoader targetClassLoader, @Nullable ProtectionDomain protectionDomain, List<String> classesToInject) throws Exception {
+            classesToInject = new ArrayList<>(classesToInject);
+            classesToInject.add(MethodHandleRegisterer.class.getName());
+
+            Set<Collection<String>> injectedClasses = getOrCreateInjectedClasses(targetClassLoader);
+            if (injectedClasses.contains(classesToInject)) {
+                return;
+            }
+            injectedClasses.add(classesToInject);
+            logger.debug("Creating helper class loader for {} containing {}", targetClassLoader, classesToInject);
+
+            ClassLoader parent = getHelperClassLoaderParent(targetClassLoader);
+            Map<String, byte[]> typeDefinitions = getTypeDefinitions(classesToInject);
+            // child first semantics are important here as the helper CL contains classes that are also present in the agent CL
+            ClassLoader helperCL = new ByteArrayClassLoader.ChildFirst(parent, true, typeDefinitions);
+
+            registerMethodHandles(classesToInject, helperCL, targetClassLoader, protectionDomain);
+        }
+
+        private static Set<Collection<String>> getOrCreateInjectedClasses(@Nullable ClassLoader targetClassLoader) {
+            Set<Collection<String>> injectedClasses = alreadyInjected.get(targetClassLoader);
+            if (injectedClasses == null) {
+                injectedClasses = new HashSet<>();
+                alreadyInjected.put(targetClassLoader, injectedClasses);
+            }
+            return injectedClasses;
+        }
+
+        private static ConcurrentMap<String, MethodHandle> ensureDispatcherForClassLoaderCreated(@Nullable ClassLoader targetClassLoader, @Nullable ProtectionDomain protectionDomain) throws IOException, ClassNotFoundException, IllegalAccessException, NoSuchFieldException {
+            ConcurrentMap<String, MethodHandle> dispatcherForClassLoader = MethodHandleDispatcher.getDispatcherForClassLoader(targetClassLoader);
+            ConcurrentMap<String, Method> reflectionDispatcherForClassLoader = MethodHandleDispatcher.getReflectionDispatcherForClassLoader(targetClassLoader);
+            if (dispatcherForClassLoader == null) {
+                // there's always a dispatcher for the bootstrap CL
+                assert targetClassLoader != null;
+                Class<?> dispatcher = injectClass(targetClassLoader, protectionDomain, MethodHandleDispatcherHolder.class.getName(), true);
+                dispatcherForClassLoader = (ConcurrentMap<String, MethodHandle>) dispatcher.getField("registry").get(null);
+                reflectionDispatcherForClassLoader = (ConcurrentMap<String, Method>) dispatcher.getField("reflectionRegistry").get(null);
+                MethodHandleDispatcher.setDispatcherForClassLoader(targetClassLoader, dispatcherForClassLoader, reflectionDispatcherForClassLoader);
+            }
+            return dispatcherForClassLoader;
+        }
+
+        @Nullable
+        private static ClassLoader getHelperClassLoaderParent(@Nullable ClassLoader targetClassLoader) {
+            ClassLoader agentClassLoader = HelperClassManager.class.getClassLoader();
+            if (agentClassLoader != null && agentClassLoader != ClassLoader.getSystemClassLoader()) {
+                // future world: when the agent is loaded from an isolated class loader
+                // the helper class loader has both, the agent class loader and the target class loader as the parent
+                return new MultipleParentClassLoader(Arrays.asList(agentClassLoader, targetClassLoader));
+            }
+           return targetClassLoader;
+        }
+
+        /**
+         * Gets all {@link RegisterMethodHandle} annotated methods and registers them in {@link MethodHandleDispatcher}
+         */
+        private static void registerMethodHandles(List<String> classesToInject, ClassLoader helperCL, @Nullable ClassLoader targetClassLoader, @Nullable ProtectionDomain protectionDomain) throws Exception {
+            ensureDispatcherForClassLoaderCreated(targetClassLoader, protectionDomain);
+            ConcurrentMap<String, MethodHandle> dispatcherForClassLoader = MethodHandleDispatcher.getDispatcherForClassLoader(targetClassLoader);
+            ConcurrentMap<String, Method> reflectionDispatcherForClassLoader = MethodHandleDispatcher.getReflectionDispatcherForClassLoader(targetClassLoader);
+
+            Class<MethodHandleRegisterer> methodHandleRegistererClass = (Class<MethodHandleRegisterer>) Class.forName(MethodHandleRegisterer.class.getName(), true, helperCL);
+            TypePool typePool = TypePool.Default.of(getAgentClassFileLocator());
+            for (String name : classesToInject) {
+                // check with Byte Buddy matchers, acting on the byte code as opposed to reflection first
+                // to avoid NoClassDefFoundErrors when the classes refer to optional types
+                boolean isHelperClass = !typePool.describe(name)
+                    .resolve()
+                    .getDeclaredMethods()
+                    .filter(isAnnotatedWith(named(RegisterMethodHandle.class.getName())))
+                    .isEmpty();
+                if (isHelperClass) {
+
+                    Class<?> clazz = Class.forName(name, true, helperCL);
+                    if (clazz.getClassLoader() != helperCL) {
+                        throw new IllegalStateException("Helper classes not loaded from helper class loader, instead loaded from " + clazz.getClassLoader());
+                    }
+                    try {
+                        // call in from the context of the created helper class loader so that helperClass.getMethods() doesn't throw NoClassDefFoundError
+                        Method registerStaticMethodHandles = methodHandleRegistererClass.getMethod("registerStaticMethodHandles", ConcurrentMap.class, ConcurrentMap.class, Class.class);
+                        registerStaticMethodHandles.invoke(null, dispatcherForClassLoader, reflectionDispatcherForClassLoader, clazz);
+                    } catch (Exception e) {
+                        logger.error("Exception while trying to register method handles for {}", name);
+                        if (e instanceof InvocationTargetException) {
+                            Throwable targetException = ((InvocationTargetException) e).getTargetException();
+                            if (targetException instanceof Exception) {
+                                throw (Exception) targetException;
+                            }
+                        }
+                        throw e;
+                    }
+                }
+            }
+        }
+
+        /**
+         * This class is loaded form the helper class loader so that {@link Class#getDeclaredMethods()} doesn't throw a
+         * {@link NoClassDefFoundError}.
+         * This would otherwise happen as the methods may reference types not visible to the agent class loader (such as the servlet API).
+         */
+        public static class MethodHandleRegisterer {
+
+            public static void registerStaticMethodHandles(ConcurrentMap<String, MethodHandle> dispatcher, ConcurrentMap<String, Method> reflectionDispatcher, Class<?> helperClass) throws ReflectiveOperationException {
+                for (Method method : helperClass.getDeclaredMethods()) {
+                    if (Modifier.isStatic(method.getModifiers()) && method.getAnnotation(RegisterMethodHandle.class) != null) {
+                        MethodHandle methodHandle = MethodHandles.lookup().unreflect(method);
+                        String key = helperClass.getName() + "#" + method.getName();
+                        // intern() to speed up map lookups (short-circuits String::equals via reference equality check)
+                        MethodHandle previousValue = dispatcher.put(key.intern(), methodHandle);
+                        reflectionDispatcher.put(key.intern(), method);
+                        if (previousValue != null) {
+                            throw new IllegalArgumentException("There is already a mapping for '" + key + "'");
+                        }
+                    }
+                }
+            }
+        }
+
+        public synchronized static void clear() {
+            alreadyInjected.clear();
+            MethodHandleDispatcher.clear();
+        }
+    }
+
+    static Class<?> injectClass(@Nullable ClassLoader targetClassLoader, @Nullable ProtectionDomain pd, String className, boolean isBootstrapClass) throws IOException, ClassNotFoundException {
         if (targetClassLoader == null) {
             if (isBootstrapClass) {
                 return Class.forName(className, false, null);
@@ -335,15 +495,23 @@ private static <T> Class<T> loadHelperClass(@Nullable ClassLoader targetClassLoa
 
     private static Map<String, byte[]> getTypeDefinitions(List<String> helperClassNames) throws IOException {
         Map<String, byte[]> typeDefinitions = new HashMap<>();
+        ClassFileLocator agentClassFileLocator = getAgentClassFileLocator();
         for (final String helperName : helperClassNames) {
-            final byte[] classBytes = getAgentClassBytes(helperName);
+            final byte[] classBytes = agentClassFileLocator.locate(helperName).resolve();
             typeDefinitions.put(helperName, classBytes);
         }
         return typeDefinitions;
     }
 
     private static byte[] getAgentClassBytes(String className) throws IOException {
-        final ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(ClassLoader.getSystemClassLoader());
-        return locator.locate(className).resolve();
+        return getAgentClassFileLocator().locate(className).resolve();
+    }
+
+    private static ClassFileLocator getAgentClassFileLocator() {
+        ClassLoader agentClassLoader = HelperClassManager.class.getClassLoader();
+        if (agentClassLoader == null) {
+            agentClassLoader = ClassLoader.getSystemClassLoader();
+        }
+        return ClassFileLocator.ForClassLoader.of(agentClassLoader);
     }
 }

apm-agent-core/src/main/java/co/elastic/apm/agent/bci/MethodHandleDispatcherHolder.java

@@ -0,0 +1,48 @@
+/*-
+ * #%L
+ * Elastic APM Java agent
+ * %%
+ * Copyright (C) 2018 - 2020 Elastic and contributors
+ * %%
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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.
+ * #L%
+ */
+package co.elastic.apm.agent.bci;
+
+import co.elastic.apm.agent.bootstrap.MethodHandleDispatcher;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.reflect.Method;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+/**
+ * A class that holds a static reference to method handler dispatcher for a specific class loader.
+ * See also {@link MethodHandleDispatcher#dispatcherByClassLoader}
+ * Used to be loaded by the parent of the class loader that loads the helper
+ * itself, thus making the helper instance non-GC-eligible as long as the parent class loader is alive.
+ * NOTE: THIS CLASS SHOULD NEVER BE INSTANTIATED NOR REFERENCED EXPLICITLY, IT SHOULD ONLY BE USED THROUGH REFLECTION
+ */
+public class MethodHandleDispatcherHolder {
+    public static final ConcurrentMap<String, MethodHandle> registry = new ConcurrentHashMap<String, MethodHandle>();
+    public static final ConcurrentMap<String, Method> reflectionRegistry = new ConcurrentHashMap<String, Method>();
+
+    // should never be instanciated
+    private MethodHandleDispatcherHolder() {
+    }
+}

apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/InputStreamFactoryHelperImpl.java → apm-agent-core/src/main/java/co/elastic/apm/agent/bci/RegisterMethodHandle.java

@@ -22,16 +22,18 @@
  * under the License.
  * #L%
  */
-package co.elastic.apm.agent.servlet.helper;
+package co.elastic.apm.agent.bci;
 
-import co.elastic.apm.agent.impl.context.Request;
-import co.elastic.apm.agent.servlet.RequestStreamRecordingInstrumentation;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
 
-import javax.servlet.ServletInputStream;
-
-public class InputStreamFactoryHelperImpl implements RequestStreamRecordingInstrumentation.InputStreamWrapperFactory {
-    @Override
-    public ServletInputStream wrap(Request request, ServletInputStream servletInputStream) {
-        return new RecordingServletInputStreamWrapper(request, servletInputStream);
-    }
+/**
+ * Static methods annotated with this will be registered in {@link co.elastic.apm.agent.bootstrap.MethodHandleDispatcher}.
+ * This enables advices to call them by looking them up via {@link co.elastic.apm.agent.bootstrap.MethodHandleDispatcher#getMethodHandle(Class, String)}
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.METHOD)
+public @interface RegisterMethodHandle {
 }