XML extensions to mirror class annotations trace lambda and trace by return type behaviour

newrelic-agent/src/main/java/com/newrelic/agent/extension/beans/Extension.java

@@ -454,6 +454,10 @@ public static class Pointcut {
             protected boolean excludeFromTransactionTrace;
             protected boolean leaf;
             protected boolean ignoreTransaction;
+            protected boolean traceLambda;
+            protected String pattern;
+            protected boolean includeNonstatic;
+            protected List<String> traceReturnTypeDescriptors;
             protected String transactionType;
 
             /**
@@ -643,6 +647,85 @@ public void setIgnoreTransaction(Boolean value) {
                 ignoreTransaction = value;
             }
 
+            /**
+             * Gets the value of the traceLambda property.
+             *
+             * @return possible object is {@link Boolean }
+             */
+            public boolean isTraceLambda() {
+                return traceLambda;
+            }
+
+            /**
+             * Sets the value of the traceLambda property.
+             *
+             * @param value allowed object is {@link Boolean }
+             */
+            public void setTraceLambda(Boolean value) {
+                traceLambda = value;
+            }
+
+            /**
+             * Gets the value of the traceReturnTypeDescriptors property.
+             *
+             * <p>
+             * This accessor traceReturnTypeDescriptors returns a reference to the live list,
+             * not a snapshot. Therefore any modification you make to the
+             * returned list will be present inside the object.
+             * This is why there is not a <CODE>set</CODE> traceReturnTypeDescriptors for the traceReturnTypeDescriptors property.
+             *
+             * <p>
+             * For example, to add a new item, do as follows:
+             * <pre>
+             *    getTraceReturnTypeDescriptors().add(newItem);
+             * </pre>
+             *
+             * <p>
+             * Objects of the following type(s) are allowed in the list {@link java.lang.String }
+             */
+            public List<String> getTraceReturnTypeDescriptors() {
+                if (traceReturnTypeDescriptors == null) {
+                    traceReturnTypeDescriptors = new ArrayList<>();
+                }
+                return traceReturnTypeDescriptors;
+            }
+
+            /**
+             * Gets the value of the pattern property.
+             *
+             * @return possible object is {@link String }
+             */
+            public String getPattern() {
+                return pattern;
+            }
+
+            /**
+             * Sets the value of the pattern property.
+             *
+             * @param value allowed object is {@link String }
+             */
+            public void setPattern(String value) {
+                pattern = value;
+            }
+
+            /**
+             * Gets the value of the includeNonstatic property.
+             *
+             * @return possible object is {@link Boolean }
+             */
+            public Boolean getIncludeNonstatic() {
+                return includeNonstatic;
+            }
+
+            /**
+             * Sets the value of the includeNonstatic property.
+             *
+             * @param value allowed object is {@link Boolean }
+             */
+            public void setIncludeNonstatic(Boolean value) {
+                includeNonstatic = value;
+            }
+
             /**
              * Gets the value of the transactionType property.
              *

newrelic-agent/src/main/java/com/newrelic/agent/extension/dom/ExtensionDomParser.java

@@ -233,6 +233,7 @@ public static Extension parseDocument(String extensionXML, boolean setSchema) th
                             pointcut.setTransactionType(getAttribute("transactionType", instrumentationChildNode, null));
 
                             List<Extension.Instrumentation.Pointcut.Method> methods = pointcut.getMethod();
+                            List<String> traceReturnTypeDescriptors = pointcut.getTraceReturnTypeDescriptors();
 
                             NodeList pointcutChildNodes = instrumentationChildNode.getChildNodes();
                             for (int p = 0; p < pointcutChildNodes.getLength(); p++) {
@@ -246,6 +247,12 @@ public static Extension parseDocument(String extensionXML, boolean setSchema) th
                                     pointcut.setInterfaceName(node.getTextContent());
                                 } else if (node.getNodeName().equals("methodAnnotation") || node.getNodeName().endsWith(":methodAnnotation")) {
                                     pointcut.setMethodAnnotation(node.getTextContent());
+                                } else if (node.getNodeName().equals("traceLambda") || node.getNodeName().endsWith(":traceLambda")) {
+                                    pointcut.setTraceLambda(Boolean.valueOf(node.getTextContent()));
+                                    pointcut.setPattern(getAttribute("pattern", node, "^\\$?(lambda|anonfun)\\$(?<name>.*)"));
+                                    pointcut.setIncludeNonstatic(Boolean.valueOf(getAttribute("includeNonstatic", node, "false")));
+                                } else if (node.getNodeName().equals("traceByReturnType") || node.getNodeName().endsWith(":traceByReturnType")) {
+                                    traceReturnTypeDescriptors.add(node.getTextContent());
                                 } else if (node.getNodeName().equals("method") || node.getNodeName().endsWith(":method")) {
                                     NodeList methodChildNodes = node.getChildNodes();
                                     Extension.Instrumentation.Pointcut.Method method = new Extension.Instrumentation.Pointcut.Method();

newrelic-agent/src/main/java/com/newrelic/agent/extension/util/ExtensionConversionUtility.java

@@ -22,7 +22,9 @@
 import com.newrelic.agent.instrumentation.classmatchers.InterfaceMatcher;
 import com.newrelic.agent.instrumentation.custom.ExtensionClassAndMethodMatcher;
 import com.newrelic.agent.instrumentation.methodmatchers.AnnotationMethodMatcher;
+import com.newrelic.agent.instrumentation.methodmatchers.LambdaMethodMatcher;
 import com.newrelic.agent.instrumentation.methodmatchers.MethodMatcher;
+import com.newrelic.agent.instrumentation.methodmatchers.ReturnTypeMethodMatcher;
 import com.newrelic.agent.instrumentation.tracing.ParameterAttributeName;
 import org.objectweb.asm.Type;
 
@@ -278,10 +280,15 @@ private static List<ParameterAttributeName> getParameterAttributeNames(List<Meth
     private static MethodMatcher createMethodMatcher(final Pointcut cut, final String pExtName,
             final Map<String, MethodMapper> classesToMethods) throws XmlException {
         List<Method> methods = cut.getMethod();
+        List<String> traceReturnTypeDescriptors = cut.getTraceReturnTypeDescriptors();
         if (methods != null && !methods.isEmpty()) {
             return MethodMatcherUtility.createMethodMatcher(getClassName(cut), methods, classesToMethods, pExtName);
         } else if (cut.getMethodAnnotation() != null) {
             return new AnnotationMethodMatcher(Type.getObjectType(cut.getMethodAnnotation().replace('.', '/')));
+        } else if (cut.isTraceLambda()) {
+            return new LambdaMethodMatcher(cut.getPattern(), cut.getIncludeNonstatic());
+        } else if (traceReturnTypeDescriptors != null && !traceReturnTypeDescriptors.isEmpty()) {
+            return new ReturnTypeMethodMatcher(traceReturnTypeDescriptors);
         } else {
             throw new XmlException(MessageFormat.format(XmlParsingMessages.NO_METHOD, pExtName));
         }

newrelic-agent/src/main/java/com/newrelic/agent/instrumentation/context/InstrumentationContextManager.java

@@ -111,6 +111,27 @@ public InstrumentationContextManager(Instrumentation instrumentation) {
           Agent.LOG.log(Level.FINEST, "scala_future_trace instrumentation is disabled because it is not explicitly enabled");
         }
 
+        if (agentConfig.getValue("instrumentation.scala_future_trace.enabled", false)) {
+          Agent.LOG.log(Level.FINEST, "scala_future_trace instrumentation is enabled");
+          matchVisitors.put(new TraceByReturnTypeMatchVisitor(), NO_OP_TRANSFORMER);
+        } else {
+          Agent.LOG.log(Level.FINEST, "scala_future_trace instrumentation is disabled because it is not explicitly enabled");
+        }
+
+        if (agentConfig.getValue("instrumentation.trace_lambda.enabled", false)) {
+            Agent.LOG.log(Level.FINEST, "trace_lambda instrumentation is enabled");
+            matchVisitors.put(new TraceLambdaVisitor(), NO_OP_TRANSFORMER);
+        } else {
+            Agent.LOG.log(Level.FINEST, "trace_lambda instrumentation is disabled because it is not explicitly enabled");
+        }
+
+        if (agentConfig.getValue("instrumentation.scala_future_trace.enabled", false)) {
+          Agent.LOG.log(Level.FINEST, "scala_future_trace instrumentation is enabled");
+          matchVisitors.put(new TraceByReturnTypeMatchVisitor(), NO_OP_TRANSFORMER);
+        } else {
+          Agent.LOG.log(Level.FINEST, "scala_future_trace instrumentation is disabled because it is not explicitly enabled");
+        }
+
         Config instrumentationConfig = agentConfig.getClassTransformerConfig().getInstrumentationConfig(
                 "com.newrelic.instrumentation.ejb-3.0");
         if (instrumentationConfig.getProperty("enabled", false)) {

newrelic-agent/src/main/java/com/newrelic/agent/instrumentation/methodmatchers/LambdaMethodMatcher.java

@@ -0,0 +1,55 @@
+/*
+ *
+ *  * Copyright 2020 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.instrumentation.methodmatchers;
+
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.commons.Method;
+
+import java.util.Set;
+import java.util.regex.Pattern;
+
+/**
+ * A method matcher that matches lambda methods.
+ */
+public final class LambdaMethodMatcher implements MethodMatcher {
+
+    private final Pattern lambdaMethodPattern;
+    private final Boolean includeNonstatic;
+
+    public LambdaMethodMatcher(String pattern, boolean includeNonstatic) {
+        super();
+        this.lambdaMethodPattern = Pattern.compile(pattern);
+        this.includeNonstatic = includeNonstatic;
+    }
+
+    @Override
+    public boolean matches(int access, String name, String desc, Set<String> annotations) {
+        return (includeNonstatic || access == Opcodes.ACC_STATIC) && lambdaMethodPattern.matcher(name).matches();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+        if (obj == null) {
+            return false;
+        }
+        return getClass() == obj.getClass();
+    }
+
+    @Override
+    public int hashCode() {
+        return super.hashCode();
+    }
+
+    @Override
+    public Method[] getExactMethods() {
+        return null;
+    }
+}

newrelic-agent/src/main/java/com/newrelic/agent/instrumentation/methodmatchers/ReturnTypeMethodMatcher.java

@@ -0,0 +1,64 @@
+/*
+ *
+ *  * Copyright 2020 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.instrumentation.methodmatchers;
+
+import org.objectweb.asm.commons.Method;
+
+import java.util.List;
+import java.util.Set;
+
+/**
+ * A method matcher that matches methods based on return type.
+ */
+public final class ReturnTypeMethodMatcher implements MethodMatcher {
+
+    private final List<String> traceReturnTypeDescriptors;
+
+    public ReturnTypeMethodMatcher(List<String> traceReturnTypeDescriptors) {
+        super();
+        this.traceReturnTypeDescriptors = traceReturnTypeDescriptors;
+    }
+
+    @Override
+    public boolean matches(int access, String name, String desc, Set<String> annotations) {
+
+        return isTracedMethod(desc);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+        if (obj == null) {
+            return false;
+        }
+        return getClass() == obj.getClass();
+    }
+
+    @Override
+    public int hashCode() {
+        return super.hashCode();
+    }
+
+    @Override
+    public Method[] getExactMethods() {
+        return null;
+    }
+
+    private boolean isTracedMethod(String descriptor) {
+        boolean matchFound = false;
+        for(String returnTypeDescriptor: traceReturnTypeDescriptors) {
+            if(descriptor.endsWith(returnTypeDescriptor)) {
+                matchFound = true;
+                break;
+            }
+        }
+        return matchFound;
+    }
+}

newrelic-agent/src/main/resources/META-INF/extensions/extension.xsd

@@ -81,7 +81,8 @@
                                                         </xs:annotation>
                                                     </xs:element>
                                                 </xs:choice>
-                                                <xs:element name="method" minOccurs="1" maxOccurs="unbounded">
+                                                <xs:choice minOccurs="1">
+                                                    <xs:element name="method" maxOccurs="unbounded">
                                                     <xs:complexType>
 														<xs:choice minOccurs="1">
 			                                                <xs:element name="returnType" type="xs:string" minOccurs="0">
@@ -136,6 +137,18 @@
 
                                                     </xs:complexType>
                                                 </xs:element>
+                                                    <xs:element name="traceLambda">
+                                                        <xs:complexType>
+                                                            <xs:simpleContent>
+                                                                <xs:extension base="xs:boolean">
+                                                                    <xs:attribute type="xs:string" name="pattern"/>
+                                                                    <xs:attribute type="xs:boolean" name="includeNonstatic"/>
+                                                                </xs:extension>
+                                                            </xs:simpleContent>
+                                                        </xs:complexType>
+                                                    </xs:element>
+                                                    <xs:element maxOccurs="unbounded" name="traceByReturnType"/>
+                                                </xs:choice>
                                             </xs:sequence>
                                         </xs:choice>
                                     </xs:sequence>

newrelic-agent/src/main/resources/extension-example.xml

@@ -121,5 +121,14 @@
 			<methodAnnotation>com.example.myAnnotation</methodAnnotation>
 		</pointcut>
 
+		<pointcut transactionStartPoint="true">
+			<className>com.sample.SampleString</className>
+			<traceLambda>true</traceLambda>
+		</pointcut>
+
+		<pointcut transactionStartPoint="true">
+			<className>com.sample.SampleString</className>
+			<traceByReturnType>Lcom/sample/SampleString;</traceByReturnType>
+		</pointcut>
 	</instrumentation>
 </extension>

newrelic-agent/src/main/resources/extension.xsd

@@ -81,7 +81,8 @@
                                                         </xs:annotation>
                                                     </xs:element>
                                                 </xs:choice>
-                                                <xs:element name="method" minOccurs="1" maxOccurs="unbounded">
+                                                <xs:choice minOccurs="1">
+                                                    <xs:element maxOccurs="unbounded" name="method">
                                                     <xs:complexType>
                                                         <xs:choice minOccurs="1">
                                                             <xs:element name="returnType" type="xs:string" minOccurs="0">
@@ -129,6 +130,18 @@
 
                                                     </xs:complexType>
                                                 </xs:element>
+                                                    <xs:element name="traceLambda">
+                                                        <xs:complexType>
+                                                            <xs:simpleContent>
+                                                                <xs:extension base="xs:boolean">
+                                                                    <xs:attribute type="xs:string" name="pattern"/>
+                                                                    <xs:attribute type="xs:boolean" name="includeNonstatic"/>
+                                                                </xs:extension>
+                                                            </xs:simpleContent>
+                                                        </xs:complexType>
+                                                    </xs:element>
+                                                    <xs:element maxOccurs="unbounded" name="traceByReturnType"/>
+                                                </xs:choice>
                                             </xs:sequence>
                                         </xs:choice>
                                     </xs:sequence>

newrelic-agent/src/test/java/com/newrelic/agent/instrumentation/methodmatchers/MethodMatcherTests.java

@@ -339,4 +339,22 @@ public void testNotMethodMatcher() {
         Assert.assertEquals(matcher.hashCode(), new NotMethodMatcher(allMatcher).hashCode());
         Assert.assertFalse(matcher.hashCode() == nestedMatcher.hashCode());
     }
+
+    @Test
+    public void testLambdaMethodMatcher() {
+        Assert.assertFalse(new LambdaMethodMatcher("", false).matches(Opcodes.ACC_STATIC, "methodName", null, com.google.common.collect.ImmutableSet.<String> of()));
+        Assert.assertFalse(new LambdaMethodMatcher("methodName", false).matches(Opcodes.ACC_PUBLIC, "methodName", null, com.google.common.collect.ImmutableSet.<String> of()));
+        Assert.assertTrue(new LambdaMethodMatcher("methodName", false).matches(Opcodes.ACC_STATIC, "methodName", null, com.google.common.collect.ImmutableSet.<String> of()));
+        Assert.assertFalse(new LambdaMethodMatcher("", true).matches(Opcodes.ACC_PUBLIC, "methodName", null, com.google.common.collect.ImmutableSet.<String> of()));
+        Assert.assertTrue(new LambdaMethodMatcher("methodName", true).matches(Opcodes.ACC_PUBLIC, "methodName", null, com.google.common.collect.ImmutableSet.<String> of()));
+    }
+
+    @Test
+    public void testReturnTypeMethodMatcher() {
+        Assert.assertFalse(new ReturnTypeMethodMatcher(Collections.<String>emptyList()).matches(Opcodes.ACC_PUBLIC, null, Type.getDescriptor(String.class), com.google.common.collect.ImmutableSet.<String> of()));
+        Assert.assertFalse(new ReturnTypeMethodMatcher(Collections.singletonList(Type.getDescriptor(Boolean.class))).matches(Opcodes.ACC_PUBLIC, null, Type.getDescriptor(String.class), com.google.common.collect.ImmutableSet.<String> of()));
+        Assert.assertTrue(new ReturnTypeMethodMatcher(Collections.singletonList(Type.getDescriptor(String.class))).matches(Opcodes.ACC_PUBLIC, null, Type.getDescriptor(String.class), com.google.common.collect.ImmutableSet.<String> of()));
+        Assert.assertFalse(new ReturnTypeMethodMatcher(Arrays.asList(Type.getDescriptor(Boolean.class), Type.getDescriptor(Short.class))).matches(Opcodes.ACC_PUBLIC, null, Type.getDescriptor(String.class), com.google.common.collect.ImmutableSet.<String> of()));
+        Assert.assertTrue(new ReturnTypeMethodMatcher(Arrays.asList(Type.getDescriptor(Boolean.class), Type.getDescriptor(String.class))).matches(Opcodes.ACC_PUBLIC, null, Type.getDescriptor(String.class), com.google.common.collect.ImmutableSet.<String> of()));
+    }
 }