Skip to content

Commit

Permalink
Added possibility to preregister types with a class reloading strategy.
Browse files Browse the repository at this point in the history
  • Loading branch information
Rafael Winterhalter committed Dec 9, 2015
1 parent 61eb2bc commit b564fc1
Show file tree
Hide file tree
Showing 2 changed files with 141 additions and 49 deletions.
142 changes: 99 additions & 43 deletions byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/loading/ClassReloadingStrategy.java
View file
Expand Up @@ -6,10 +6,7 @@
import java.io.File;
import java.lang.instrument.*;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
Expand Down Expand Up @@ -59,36 +56,20 @@ public class ClassReloadingStrategy implements ClassLoadingStrategy {
private final BootstrapInjection bootstrapInjection;

/**
* Creates a class reloading strategy for the given instrumentation. The given instrumentation must either
* support {@link java.lang.instrument.Instrumentation#isRedefineClassesSupported()} or
* {@link java.lang.instrument.Instrumentation#isRetransformClassesSupported()}. If both modes are supported,
* classes will be transformed using a class redefinition.
*
* @param instrumentation The instrumentation to be used by this reloading strategy.
* The preregistered types of this instance.
*/
public ClassReloadingStrategy(Instrumentation instrumentation) {
this.instrumentation = instrumentation;
if (instrumentation.isRedefineClassesSupported()) {
engine = Engine.REDEFINITION;
} else if (instrumentation.isRetransformClassesSupported()) {
engine = Engine.RETRANSFORMATION;
} else {
throw new IllegalArgumentException("Instrumentation does not support manipulation of loaded classes: " + instrumentation);
}
bootstrapInjection = BootstrapInjection.Disabled.INSTANCE;
}
private final Map<String, Class<?>> preregisteredTypes;

/**
* Creates a class reloading strategy for the given instrumentation using an explicit transformation strategy which
* is represented by an {@link net.bytebuddy.dynamic.loading.ClassReloadingStrategy.Engine}.
* is represented by an {@link net.bytebuddy.dynamic.loading.ClassReloadingStrategy.Engine}. The given instrumentation
* must support the engine's transformation type.
*
* @param instrumentation The instrumentation to be used by this reloading strategy.
* @param engine An engine which performs the actual redefinition of a {@link java.lang.Class}.
*/
public ClassReloadingStrategy(Instrumentation instrumentation, Engine engine) {
this.instrumentation = instrumentation;
this.engine = engine;
bootstrapInjection = BootstrapInjection.Disabled.INSTANCE;
this(instrumentation, engine.validate(instrumentation), BootstrapInjection.Disabled.INSTANCE, Collections.<String, Class<?>>emptyMap());
}

/**
Expand All @@ -97,11 +78,36 @@ public ClassReloadingStrategy(Instrumentation instrumentation, Engine engine) {
* @param instrumentation The instrumentation to be used by this reloading strategy.
* @param engine An engine which performs the actual redefinition of a {@link java.lang.Class}.
* @param bootstrapInjection The bootstrap class loader injection strategy to use.
* @param preregisteredTypes The preregistered types of this instance.
*/
protected ClassReloadingStrategy(Instrumentation instrumentation, Engine engine, BootstrapInjection bootstrapInjection) {
protected ClassReloadingStrategy(Instrumentation instrumentation,
Engine engine,
BootstrapInjection bootstrapInjection,
Map<String, Class<?>> preregisteredTypes) {
this.instrumentation = instrumentation;
this.engine = engine;
this.bootstrapInjection = bootstrapInjection;
this.preregisteredTypes = preregisteredTypes;
}

/**
* Creates a class reloading strategy for the given instrumentation. The given instrumentation must either
* support {@link java.lang.instrument.Instrumentation#isRedefineClassesSupported()} or
* {@link java.lang.instrument.Instrumentation#isRetransformClassesSupported()}. If both modes are supported,
* classes will be transformed using a class redefinition.
*
* @param instrumentation The instrumentation to be used by this reloading strategy.
*/
public static ClassReloadingStrategy of(Instrumentation instrumentation) {
Engine engine;
if (instrumentation.isRedefineClassesSupported()) {
engine = Engine.REDEFINITION;
} else if (instrumentation.isRetransformClassesSupported()) {
engine = Engine.RETRANSFORMATION;
} else {
throw new IllegalArgumentException("Instrumentation does not support manipulation of loaded classes: " + instrumentation);
}
return new ClassReloadingStrategy(instrumentation, engine);
}

/**
Expand All @@ -128,7 +134,7 @@ public static ClassReloadingStrategy fromInstalledAgent() {
if (instrumentation == null) {
throw new IllegalStateException("The Byte Buddy agent is not installed");
}
return new ClassReloadingStrategy(instrumentation);
return ClassReloadingStrategy.of(instrumentation);
} catch (RuntimeException exception) {
throw exception;
} catch (Exception exception) {
Expand All @@ -138,15 +144,15 @@ public static ClassReloadingStrategy fromInstalledAgent() {

@Override
public Map<TypeDescription, Class<?>> load(ClassLoader classLoader, Map<TypeDescription, byte[]> types) {
Map<String, Class<?>> loadedTypes = new HashMap<String, Class<?>>();
Map<String, Class<?>> availableTypes = new HashMap<String, Class<?>>(preregisteredTypes);
for (Class<?> type : instrumentation.getInitiatedClasses(classLoader)) {
loadedTypes.put(type.getName(), type);
availableTypes.put(type.getName(), type);
}
Map<Class<?>, ClassDefinition> classDefinitions = new ConcurrentHashMap<Class<?>, ClassDefinition>();
Map<TypeDescription, Class<?>> loadedClasses = new HashMap<TypeDescription, Class<?>>();
Map<TypeDescription, byte[]> unloadedClasses = new LinkedHashMap<TypeDescription, byte[]>();
for (Map.Entry<TypeDescription, byte[]> entry : types.entrySet()) {
Class<?> type = loadedTypes.get(entry.getKey().getName());
Class<?> type = availableTypes.get(entry.getKey().getName());
if (type != null) {
classDefinitions.put(type, new ClassDefinition(type, entry.getValue()));
loadedClasses.put(entry.getKey(), type);
Expand Down Expand Up @@ -199,20 +205,41 @@ public ClassReloadingStrategy reset(Class<?>... type) {
* @return A class reloading strategy with bootstrap injection enabled.
*/
public ClassReloadingStrategy enableBootstrapInjection(File folder) {
return new ClassReloadingStrategy(instrumentation, engine, new BootstrapInjection.Enabled(folder));
return new ClassReloadingStrategy(instrumentation, engine, new BootstrapInjection.Enabled(folder), preregisteredTypes);
}

/**
* Registers a type to be explicitly available without explicit lookup.
*
* @param type The loaded types that are explicitly available.
* @return This class reloading strategy with the given types being explicitly available.
*/
public ClassReloadingStrategy preregistered(Class<?>... type) {
Map<String, Class<?>> preregisteredTypes = new HashMap<String, Class<?>>(this.preregisteredTypes);
for (Class<?> aType : type) {
preregisteredTypes.put(aType.getName(), aType);
}
return new ClassReloadingStrategy(instrumentation, engine, bootstrapInjection, preregisteredTypes);
}

@Override
public boolean equals(Object other) {
return this == other || !(other == null || getClass() != other.getClass())
&& engine == ((ClassReloadingStrategy) other).engine
&& instrumentation.equals(((ClassReloadingStrategy) other).instrumentation)
&& bootstrapInjection.equals(((ClassReloadingStrategy) other).bootstrapInjection);
if (this == other) return true;
if (other == null || getClass() != other.getClass()) return false;
ClassReloadingStrategy that = (ClassReloadingStrategy) other;
return instrumentation.equals(that.instrumentation)
&& engine == that.engine
&& bootstrapInjection.equals(that.bootstrapInjection)
&& preregisteredTypes.equals(that.preregisteredTypes);
}

@Override
public int hashCode() {
return 31 * 31 * instrumentation.hashCode() + 31 * engine.hashCode() + bootstrapInjection.hashCode();
int result = instrumentation.hashCode();
result = 31 * result + engine.hashCode();
result = 31 * result + bootstrapInjection.hashCode();
result = 31 * result + preregisteredTypes.hashCode();
return result;
}

@Override
Expand All @@ -221,6 +248,7 @@ public String toString() {
"instrumentation=" + instrumentation +
", engine=" + engine +
", bootstrapInjection=" + bootstrapInjection +
", preregisteredTypes=" + preregisteredTypes +
'}';
}

Expand All @@ -239,6 +267,14 @@ protected void apply(Instrumentation instrumentation,
Map<Class<?>, ClassDefinition> classDefinitions) throws UnmodifiableClassException, ClassNotFoundException {
instrumentation.redefineClasses(classDefinitions.values().toArray(new ClassDefinition[classDefinitions.size()]));
}

@Override
protected Engine validate(Instrumentation instrumentation) {
if (!instrumentation.isRedefineClassesSupported()) {
throw new IllegalArgumentException("Does not support redefinition: " + instrumentation);
}
return this;
}
},

/**
Expand All @@ -250,7 +286,7 @@ protected void apply(Instrumentation instrumentation,
@Override
protected void apply(Instrumentation instrumentation, Map<Class<?>, ClassDefinition> classDefinitions) throws UnmodifiableClassException {
ClassRedefinitionTransformer classRedefinitionTransformer = new ClassRedefinitionTransformer(classDefinitions);
synchronized (instrumentation) {
synchronized (this) {
instrumentation.addTransformer(classRedefinitionTransformer, REDEFINE_CLASSES);
try {
instrumentation.retransformClasses(classDefinitions.keySet().toArray(new Class<?>[classDefinitions.size()]));
Expand All @@ -260,6 +296,14 @@ protected void apply(Instrumentation instrumentation, Map<Class<?>, ClassDefinit
}
classRedefinitionTransformer.assertTransformation();
}

@Override
protected Engine validate(Instrumentation instrumentation) {
if (!instrumentation.isRetransformClassesSupported()) {
throw new IllegalArgumentException("Does not support retransformation: " + instrumentation);
}
return this;
}
};

/**
Expand Down Expand Up @@ -294,6 +338,14 @@ protected void apply(Instrumentation instrumentation, Map<Class<?>, ClassDefinit
protected abstract void apply(Instrumentation instrumentation,
Map<Class<?>, ClassDefinition> classDefinitions) throws UnmodifiableClassException, ClassNotFoundException;

/**
* Validates that this engine supports a given transformation type.
*
* @param instrumentation The instrumentation instance being used.
* @return This engine.
*/
protected abstract Engine validate(Instrumentation instrumentation);

/**
* Returns {@code true} if this engine represents {@link net.bytebuddy.dynamic.loading.ClassReloadingStrategy.Engine#REDEFINITION}.
*
Expand All @@ -313,6 +365,11 @@ public String toString() {
*/
protected static class ClassRedefinitionTransformer implements ClassFileTransformer {

/**
* Indicates that a class is not redefined.
*/
private static final byte[] NO_REDEFINITION = null;

/**
* A mapping of classes to be redefined to their redefined class definitions.
*/
Expand All @@ -328,16 +385,15 @@ protected ClassRedefinitionTransformer(Map<Class<?>, ClassDefinition> redefinedC
}

@Override
public byte[] transform(ClassLoader loader,
String className,
public byte[] transform(ClassLoader classLoader,
String internalTypeName,
Class<?> classBeingRedefined,
ProtectionDomain protectionDomain,
byte[] classfileBuffer) throws IllegalClassFormatException {
ClassDefinition redefinedClass = redefinedClasses.remove(classBeingRedefined);
if (redefinedClass == null) {
throw new IllegalArgumentException("Encountered class without redefinition information");
}
return redefinedClass.getDefinitionClassFile();
return redefinedClass == null
? NO_REDEFINITION
: redefinedClass.getDefinitionClassFile();
}

/**
Expand Down
48 changes: 42 additions & 6 deletions byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassReloadingStrategyTest.java
View file
Expand Up @@ -2,20 +2,24 @@

import net.bytebuddy.ByteBuddy;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.implementation.FixedValue;
import net.bytebuddy.test.utility.AgentAttachmentRule;
import net.bytebuddy.test.utility.ObjectPropertyAssertion;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.MethodRule;
import org.mockito.ArgumentCaptor;

import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.util.Collections;

import static junit.framework.TestCase.assertEquals;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static org.hamcrest.CoreMatchers.*;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.*;

public class ClassReloadingStrategyTest {

Expand Down Expand Up @@ -103,20 +107,52 @@ public void testRetransformationReloadingStrategy() throws Exception {
}

@Test
public void testRetransformationFunctional() throws Exception {
public void testPreregisteredType() throws Exception {
Instrumentation instrumentation = mock(Instrumentation.class);
ClassLoader classLoader = mock(ClassLoader.class);
when(instrumentation.isRedefineClassesSupported()).thenReturn(true);
when(instrumentation.getInitiatedClasses(classLoader)).thenReturn(new Class<?>[0]);
ClassReloadingStrategy classReloadingStrategy = ClassReloadingStrategy.of(instrumentation).preregistered(Object.class);
ArgumentCaptor<ClassDefinition> classDefinition = ArgumentCaptor.forClass(ClassDefinition.class);
classReloadingStrategy.load(classLoader, Collections.singletonMap(TypeDescription.OBJECT, new byte[]{1, 2, 3}));
verify(instrumentation).redefineClasses(classDefinition.capture());
assertEquals(Object.class, classDefinition.getValue().getDefinitionClass());
assertThat(classDefinition.getValue().getDefinitionClassFile(), is(new byte[]{1, 2, 3}));
}

@Test
public void testRetransformationDiscovery() throws Exception {
Instrumentation instrumentation = mock(Instrumentation.class);
when(instrumentation.isRetransformClassesSupported()).thenReturn(true);
assertThat(new ClassReloadingStrategy(instrumentation), notNullValue(ClassReloadingStrategy.class));
assertThat(ClassReloadingStrategy.of(instrumentation), notNullValue(ClassReloadingStrategy.class));
}

@Test(expected = IllegalArgumentException.class)
public void testNonCompatible() throws Exception {
new ClassReloadingStrategy(mock(Instrumentation.class));
ClassReloadingStrategy.of(mock(Instrumentation.class));
}

@Test(expected = IllegalArgumentException.class)
public void testNoRedefinition() throws Exception {
new ClassReloadingStrategy(mock(Instrumentation.class), ClassReloadingStrategy.Engine.REDEFINITION);
}

@Test(expected = IllegalArgumentException.class)
public void testNoRetransformation() throws Exception {
new ClassReloadingStrategy(mock(Instrumentation.class), ClassReloadingStrategy.Engine.RETRANSFORMATION);
}

@Test(expected = IllegalStateException.class)
public void testResetNotSupported() throws Exception {
new ClassReloadingStrategy(mock(Instrumentation.class), ClassReloadingStrategy.Engine.RETRANSFORMATION).reset();
Instrumentation instrumentation = mock(Instrumentation.class);
when(instrumentation.isRetransformClassesSupported()).thenReturn(true);
new ClassReloadingStrategy(instrumentation, ClassReloadingStrategy.Engine.RETRANSFORMATION).reset();
}

@Test
public void testEngineSelfReport() throws Exception {
assertThat(ClassReloadingStrategy.Engine.REDEFINITION.isRedefinition(), is(true));
assertThat(ClassReloadingStrategy.Engine.RETRANSFORMATION.isRedefinition(), is(false));
}

@Test
Expand Down

0 comments on commit b564fc1

Please sign in to comment.