6712: Make agent retransform classes when loaded dynamically

Reviewed-by: hirt

agent/src/main/java/org/openjdk/jmc/agent/Agent.java

@@ -34,10 +34,13 @@
 
 import java.io.File;
 import java.io.FileInputStream;
-import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.lang.instrument.Instrumentation;
+import java.lang.instrument.UnmodifiableClassException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -56,6 +59,7 @@ public class Agent {
 	 */
 	public final static String VERSION = "0.0.2"; //$NON-NLS-1$
 	private final static String DEFAULT_CONFIG = "jfrprobes.xml"; //$NON-NLS-1$
+	private static boolean loadedDynamically = false;
 
 	/**
 	 * This method is run when the agent is started from the command line.
@@ -82,6 +86,7 @@ public static void premain(String agentArguments, Instrumentation instrumentatio
 	public static void agentmain(String agentArguments, Instrumentation instrumentation) {
 		printVersion();
 		getLogger().fine("Starting from agentmain"); //$NON-NLS-1$
+		loadedDynamically = true;
 		initializeAgent(agentArguments, instrumentation);
 	}
 
@@ -100,6 +105,9 @@ public static void initializeAgent(InputStream configuration, Instrumentation in
 		TransformRegistry registry = DefaultTransformRegistry.from(configuration);
 		instrumentation.addTransformer(new Transformer(registry), true);
 		AgentManagementFactory.createAndRegisterAgentControllerMBean(instrumentation, registry);
+		if (loadedDynamically) {
+			retransformClasses(registry.getClassNames(), instrumentation);
+		}
 	}
 
 	/**
@@ -130,6 +138,31 @@ private static void initializeAgent(String agentArguments, Instrumentation instr
 		}
 	}
 
+	/**
+	 * Retransforms the required classes when the agent is loaded dynamically.
+	 *
+	 * @param clazzes
+	 *            list of names of classes to retransform
+	 * @param instrumentation
+	 *            the {@link Instrumentation} instance.
+	 */
+	private static void retransformClasses(Set<String> clazzes, Instrumentation instrumentation) {
+		List<Class<?>> classesToRetransform = new ArrayList<Class<?>>();
+		for (String clazz : clazzes) {
+			try {
+				Class<?> classToRetransform = Class.forName(clazz.replace('/', '.'));
+				classesToRetransform.add(classToRetransform);
+			} catch (ClassNotFoundException cnfe) {
+				getLogger().log(Level.SEVERE, "Unable to find class: " + clazz, cnfe);
+			}
+		}
+		try {
+			instrumentation.retransformClasses(classesToRetransform.toArray(new Class<?>[0]));
+		} catch (UnmodifiableClassException e) {
+			getLogger().log(Level.SEVERE, "Unable to retransform classes", e);
+		}
+	}
+
 	private static void printVersion() {
 		getLogger().info(String.format("JMC BCI agent v%s", VERSION)); //$NON-NLS-1$
 	}

agent/src/main/java/org/openjdk/jmc/agent/TransformRegistry.java

@@ -33,6 +33,7 @@
 package org.openjdk.jmc.agent;
 
 import java.util.List;
+import java.util.Set;
 
 public interface TransformRegistry {
 	/**
@@ -44,7 +45,7 @@ public interface TransformRegistry {
 	 */
 	boolean hasPendingTransforms(String className);
 
-	/**O
+	/**
 	 * Returns the list of {@link TransformDescriptor}s for the named class.
 	 *
 	 * @param className
@@ -53,6 +54,13 @@ public interface TransformRegistry {
 	 */
 	List<TransformDescriptor> getTransformData(String className);
 
+	/**
+	 * Returns the names of all classes stored in the registry.
+	 *
+	 * @return the unmodifiable set of class names.
+	 */
+	Set<String> getClassNames();
+
 	/**
 	 * Modifies class information in the registry according to the xml description.
 	 *

agent/src/main/java/org/openjdk/jmc/agent/impl/DefaultTransformRegistry.java

@@ -35,11 +35,13 @@
 import java.io.InputStream;
 import java.io.StringReader;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Set;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -456,6 +458,10 @@ public List<String> clearAllTransformData() {
 		return classNames;
 	}
 
+	public Set<String> getClassNames() {
+		return Collections.unmodifiableSet(transformData.keySet());
+	}
+
 	public void setRevertInstrumentation(boolean shouldRevert) {
 		this.revertInstrumentation = shouldRevert;
 	}