Skip to content

Commit

Permalink
Instrumentation priority (#1200)
Browse files Browse the repository at this point in the history
  • Loading branch information
meiao authored May 1, 2023
1 parent a083ca3 commit 2142bb5
Show file tree
Hide file tree
Showing 4 changed files with 217 additions and 13 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
Expand All @@ -47,6 +48,7 @@
*/
public class PackageValidationResult {
private static final Pattern WEAVE_PACKAGE_PATTERN = Pattern.compile("^com.newrelic.weave..*");
public static final Comparator<PackageValidationResult> CONFIG_COMPARATOR = Comparator.comparing(PackageValidationResult::weavePackageConfig);
private final Queue<WeaveViolation> violations = Queues.newConcurrentLinkedQueue();
private final Map<String, ClassNode> utilClasses = new ConcurrentHashMap<>();
private final Map<String, ClassNode> allAnnotationClasses = new ConcurrentHashMap<>();
Expand Down Expand Up @@ -569,4 +571,8 @@ public Map<String, byte[]> computeUtilityClassBytes(ClassCache cache) {
public WeavePackage getWeavePackage() {
return this.weavePackage;
}

private static WeavePackageConfig weavePackageConfig(PackageValidationResult p) {
return p.getWeavePackage().getConfig();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,21 @@

package com.newrelic.weave.weavepackage;

import org.objectweb.asm.tree.ClassNode;

import java.io.IOException;
import java.lang.NumberFormatException;
import java.lang.instrument.Instrumentation;
import java.net.URL;
import java.util.jar.Attributes;
import java.util.jar.JarInputStream;
import java.util.regex.Pattern;
import java.util.regex.Matcher;

import org.objectweb.asm.tree.ClassNode;
import java.util.regex.Pattern;

/**
* Holds metadata and config options for a {@link WeavePackage}. Use {@link #builder()} to instantiate.
*/
public class WeavePackageConfig {
public class WeavePackageConfig implements Comparable<WeavePackageConfig>{

/**
* Creates a WeavePackageConfig using the builder pattern.
*/
Expand All @@ -35,6 +35,7 @@ public static class Builder {
private float version = 1.0f;
private String source = null;
private boolean enabled = true;
private long priority = 0L;
private boolean custom = false;
private ClassNode errorHandleClassNode = ErrorTrapHandler.NO_ERROR_TRAP_HANDLER;
private WeavePreprocessor preprocessor = WeavePreprocessor.NO_PREPROCESSOR;
Expand Down Expand Up @@ -94,7 +95,7 @@ public Builder vendorId(String vendorId) {
}

/**
* Set the of the weave package.
* Set the version of the weave package.
*
* @param version the version of the weave package
* @return builder with updated state
Expand Down Expand Up @@ -129,9 +130,32 @@ public Builder enabled(boolean enabled) {
}

/**
* Set whether or not the package is "custom".
* Sets the priority of this package. The higher the priority, the earlier it will be weaved.<br/>
* A package that is weaved later wraps any instrumentation that was weaved earlier.<br/>
* <pre>
* Priority -10 {
* Priority 0 {
* Priority 10 {
* Original code
* }
* }
* }
*
* </pre>
* The default priority (and the most common for the agent instrumentation) is 0.<br/>
* So to have your code be wrapped by the agent instrumentation, set a positive priority.
* @param priority the priority of the instrumentation package
* @return builder with updated state
*/
public Builder priority(long priority) {
this.priority = priority;
return this;
}

/**
* Set whether the package is "custom".
*
* @param custom whether or not the package is "custom"
* @param custom whether the package is "custom"
* @return builder with updated state
*/
public Builder custom(boolean custom) {
Expand Down Expand Up @@ -249,7 +273,10 @@ public Builder jarInputStream(JarInputStream jarStream) throws Exception {
enabled = Boolean.parseBoolean(enabledS);
}

return this.name(name).alias(alias).vendorId(vendorId).version(version).enabled(enabled);
String priorityS = mainAttributes.getValue("Priority");
long priority = priorityS == null ? 0 : Long.parseLong(priorityS);

return this.name(name).alias(alias).vendorId(vendorId).version(version).enabled(enabled).priority(priority);
}

/**
Expand All @@ -262,7 +289,7 @@ public WeavePackageConfig build() {
if (null == name) {
throw new RuntimeException("WeavePackageConfig must have an Implementation-Name");
}
return new WeavePackageConfig(name, alias, vendorId, version, enabled, source, custom, instrumentation,
return new WeavePackageConfig(name, alias, vendorId, version, enabled, priority, source, custom, instrumentation,
errorHandleClassNode, preprocessor, postprocessor, extensionClassTemplate);
}
}
Expand Down Expand Up @@ -293,15 +320,17 @@ public static Builder builder() {
private final WeavePreprocessor preprocessor;
private final WeavePostprocessor postprocessor;
private final ClassNode extensionClassTemplate;
private final long priority;

private WeavePackageConfig(String name, String alias, String vendorId, float version, boolean enabled,
String source, boolean custom, Instrumentation instrumentation, ClassNode errorTrapClassNode,
long priority, String source, boolean custom, Instrumentation instrumentation, ClassNode errorTrapClassNode,
WeavePreprocessor preprocessor, WeavePostprocessor postprocessor, ClassNode extensionClassTemplate) {
this.name = name;
this.alias = alias;
this.vendorId = vendorId;
this.version = version;
this.enabled = enabled;
this.priority = priority;
this.source = source;
this.custom = custom;

Expand Down Expand Up @@ -356,6 +385,13 @@ public boolean isEnabled() {
return enabled;
}

/**
* The priority of this package.
*/
public long getPriority() {
return priority;
}

/**
* Whether or not the {@link WeavePackage} is "custom", i.e. loaded by the agent from the /extensions directory.
*/
Expand Down Expand Up @@ -403,6 +439,26 @@ public ClassNode getExtensionTemplate() {

@Override
public String toString() {
return "WeavePackageConfig [name=" + name + ", version=" + version + ", enabled=" + enabled + "]";
return "WeavePackageConfig [name=" + name + ", version=" + version + ", enabled=" + enabled + ", priority=" + priority + "]";
}


@Override
public int compareTo(WeavePackageConfig that) {
// higher priority should come first
if (this.priority != that.priority) {
return (this.priority > that.priority ? -1 : 1);
}

if (this.name == null && that.name == null) {
return 0;
} else if (this.name == null) {
return -1;
} else if (that.name == null) {
return 1;
} else {
// reverse sorting for name
return that.name.compareTo(this.name);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
Expand Down Expand Up @@ -339,7 +340,10 @@ public byte[] weave(ClassLoader classloader, ClassCache cache, String className,

ClassNode composite = WeaveUtils.convertToClassNode(targetBytes);
PackageWeaveResult finalResult = null;
for (PackageValidationResult weavePackageResult : matchedPackageResults) {

List<PackageValidationResult> sortedMatchedPackages = new ArrayList<>(matchedPackageResults);
sortedMatchedPackages.sort(PackageValidationResult.CONFIG_COMPARATOR);
for (PackageValidationResult weavePackageResult : sortedMatchedPackages) {
PackageWeaveResult result = weavePackageResult.weave(className, superNames, interfaceNames, composite,
cache, skipMethods);
if (null != weaveListener) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@
package com.newrelic.weave.weavepackage;

import com.newrelic.api.agent.weaver.MatchType;
import com.newrelic.api.agent.weaver.Weave;
import com.newrelic.api.agent.weaver.Weaver;
import com.newrelic.weave.WeaveTestUtils;
import org.junit.Test;
import org.junit.runner.RunWith;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;

public class InstrumentationPriorityTest {

/**
* Tests three modules applying to the same method, but with different priorities.
*/
@Test
public void testWeavePriority() throws IOException {
WeavePackageManager wpm = new WeavePackageManager();
WeavePackage weavePackage1 = createWeavePackage("instrumentation.order-one",
"com.newrelic.weave.weavepackage.InstrumentationPriorityTest$OrderOne", -1L);
wpm.register(weavePackage1);
WeavePackage weavePackage2 = createWeavePackage("instrumentation.order-two",
"com.newrelic.weave.weavepackage.InstrumentationPriorityTest$OrderTwo", 0L);
wpm.register(weavePackage2);
WeavePackage weavePackage3 = createWeavePackage("instrumentation.order-three",
"com.newrelic.weave.weavepackage.InstrumentationPriorityTest$OrderThree", 1L);
wpm.register(weavePackage3);

String internalName = "com/newrelic/weave/weavepackage/InstrumentationPriorityTest$OriginalClass";
String className = "com.newrelic.weave.weavepackage.InstrumentationPriorityTest$OriginalClass";
byte[] compositeBytes = WeaveTestUtils.getClassBytes(className);
// /*-
compositeBytes = wpm.weave(Thread.currentThread().getContextClassLoader(), internalName, compositeBytes,
Collections.emptyMap());
for (PackageValidationResult res :
wpm.validPackages.getIfPresent(Thread.currentThread().getContextClassLoader()).values()) {
WeaveTestUtils.expectViolations(res);
}

assertNotNull(compositeBytes);
WeaveTestUtils.addToContextClassloader(className, compositeBytes);
InstrumentationPriorityTest.OriginalClass oc = new InstrumentationPriorityTest.OriginalClass();
assertEquals("123original321", oc.aMethodToWeave());
}

/**
* Tests three modules applying to the same method, but with the same priorities,
* should use alphabetical order.
*/
@Test
public void testWeaveMatchingPriority() throws IOException {
WeavePackageManager wpm = new WeavePackageManager();
WeavePackage weavePackage1 = createWeavePackage("instrumentation.order-one",
"com.newrelic.weave.weavepackage.InstrumentationPriorityTest$OrderOne", 0L);
wpm.register(weavePackage1);
WeavePackage weavePackage2 = createWeavePackage("instrumentation.order-two",
"com.newrelic.weave.weavepackage.InstrumentationPriorityTest$OrderTwo", 0L);
wpm.register(weavePackage2);
WeavePackage weavePackage3 = createWeavePackage("instrumentation.order-three",
"com.newrelic.weave.weavepackage.InstrumentationPriorityTest$OrderThree", 0L);
wpm.register(weavePackage3);

String internalName = "com/newrelic/weave/weavepackage/InstrumentationPriorityTest$OriginalClass";
String className = "com.newrelic.weave.weavepackage.InstrumentationPriorityTest$OriginalClass2";
byte[] compositeBytes = WeaveTestUtils.getClassBytes(className);
assertNotNull(compositeBytes);

// /*-
compositeBytes = wpm.weave(Thread.currentThread().getContextClassLoader(), internalName, compositeBytes,
Collections.emptyMap());
for (PackageValidationResult res :
wpm.validPackages.getIfPresent(Thread.currentThread().getContextClassLoader()).values()) {
WeaveTestUtils.expectViolations(res);
}

assertNotNull(compositeBytes);
WeaveTestUtils.addToContextClassloader(className, compositeBytes);
InstrumentationPriorityTest.OriginalClass2 oc = new InstrumentationPriorityTest.OriginalClass2();
assertEquals("132original231", oc.aMethodToWeave());
}

private WeavePackage createWeavePackage(String weavePackageName, String instrumentationClass, long priority) throws IOException {
List<byte[]> weaveBytes = new ArrayList<>();
weaveBytes.add(WeaveTestUtils.getClassBytes(instrumentationClass));
WeavePackageConfig config = WeavePackageConfig.builder()
.name(weavePackageName)
.priority(priority)
.source("com.newrelic.weave.weavepackage.testclasses")
.build();
return new WeavePackage(config, weaveBytes);
}

private interface TargetInterface {
String aMethodToWeave();
}

private static class OriginalClass implements TargetInterface {
public String aMethodToWeave() {
return "original";
}
}

// this second class is needed because the classloader was not cleaned between tests
private static class OriginalClass2 implements TargetInterface {
public String aMethodToWeave() {
return "original";
}
}

@Weave(type = MatchType.Interface, originalName = "com.newrelic.weave.weavepackage.InstrumentationPriorityTest$TargetInterface")
private static class OrderOne {
public String aMethodToWeave() {
return "1" + Weaver.callOriginal() + "1";
}
}

@Weave(type = MatchType.Interface, originalName = "com.newrelic.weave.weavepackage.InstrumentationPriorityTest$TargetInterface")
private static class OrderTwo {

public String aMethodToWeave() {
return "2" + Weaver.callOriginal() + "2";
}
}

@Weave(type = MatchType.Interface, originalName = "com.newrelic.weave.weavepackage.InstrumentationPriorityTest$TargetInterface")
private static class OrderThree {
public String aMethodToWeave() {
return "3" + Weaver.callOriginal() + "3";
}
}
}

0 comments on commit 2142bb5

Please sign in to comment.