optimize classpool usage

erkieh · Oct 9, 2014 · 627a1b6 · 627a1b6
1 parent 36aea6d
commit 627a1b6
Show file tree

Hide file tree

Showing 7 changed files with 1,993 additions and 39 deletions.
diff --git a/src/main/java/io/github/proxyhotswap/AbstractProxyTransformer.java b/src/main/java/io/github/proxyhotswap/AbstractProxyTransformer.java
@@ -22,7 +22,7 @@
  */
 public abstract class AbstractProxyTransformer implements ClassFileTransformer {
 	protected static final String INIT_FIELD_PREFIX = "initCalled";
-	protected static final ClassPool classPool = ClassPool.getDefault();
+	protected static final ClassPool classPool = TransformationUtils.getClassPool();
 
 	protected Instrumentation inst;
 	protected Map<Class<?>, TransformationState> transformationStates;
@@ -87,23 +87,20 @@ protected byte[] generateNewProxyClass(ClassLoader loader, String className, Cla
 
 		byte[] newByteCode = getNewByteCode(loader, className, classBeingRedefined);
 
-		CtClass cc = classPool.makeClass(new ByteArrayInputStream(newByteCode));
-		try {
-			String random = generateRandomString();
-			String initFieldName = INIT_FIELD_PREFIX + random;
-			addStaticInitStateField(cc, initFieldName);
-
-			String method = getInitCall(cc, random);
-
-			addInitCallToMethods(cc, initFieldName, method);
-
-			System.out.println("writing");
-			cc.writeFile("C:\\Users\\Juhtla\\Desktop\\");
-			System.out.println("written " + cc.getName());
-			return cc.toBytecode();
-		} finally {
-			TransformationUtils.detachCtClass(cc);
-		}
+		CtClass cc = getCtClass(newByteCode, className);
+		String random = generateRandomString();
+		String initFieldName = INIT_FIELD_PREFIX + random;
+		addStaticInitStateField(cc, initFieldName);
+
+		String method = getInitCall(cc, random);
+
+		addInitCallToMethods(cc, initFieldName, method);
+
+		return cc.toBytecode();
+	}
+
+	protected CtClass getCtClass(byte[] newByteCode, String className) throws Exception {
+		return classPool.makeClass(new ByteArrayInputStream(newByteCode), false);
 	}
 
 	protected abstract String getInitCall(CtClass cc, String random) throws Exception;
@@ -159,5 +156,4 @@ protected TransformationState setClassAsWaiting(Class<?> classBeingRedefined) {
 	protected TransformationState removeClassState(Class<?> classBeingRedefined) {
 		return transformationStates.remove(classBeingRedefined);
 	}
-
 }
diff --git a/src/main/java/io/github/proxyhotswap/ClassfileBufferSigantureTransformer.java b/src/main/java/io/github/proxyhotswap/ClassfileBufferSigantureTransformer.java
@@ -21,20 +21,18 @@
 public class ClassfileBufferSigantureTransformer implements ClassFileTransformer {
 
 	private static Map<String, String> classSignatures = new ConcurrentHashMap<>();
-	protected static final ClassPool classPool = ClassPool.getDefault();
+	protected static final ClassPool classPool = TransformationUtils.getClassPool();
 
 	public byte[] transform(ClassLoader loader, String className, final Class<?> classBeingRedefined,
 			ProtectionDomain protectionDomain, final byte[] classfileBuffer) {
 		if (classBeingRedefined == null)
 			return null;
 		CtClass cc = null;
 		try {
-			cc = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
+			cc = classPool.makeClass(new ByteArrayInputStream(classfileBuffer), false);
 			classSignatures.put(classBeingRedefined.getName(), getSignature(cc));
 		} catch (IOException | RuntimeException e) {
 			TransformationUtils.logError(e);
-		} finally {
-			TransformationUtils.detachCtClass(cc);
 		}
 		return null;
 	}

diff --git a/src/main/java/io/github/proxyhotswap/TransformationUtils.java b/src/main/java/io/github/proxyhotswap/TransformationUtils.java
@@ -1,5 +1,6 @@
 package io.github.proxyhotswap;
 
+import io.github.proxyhotswap.javassist.ClassPool;
 import io.github.proxyhotswap.javassist.CtClass;
 
 /**
@@ -20,4 +21,12 @@ public static void detachCtClass(CtClass cc) {
 	public static void logError(Exception e) {
 		e.printStackTrace();
 	}
+
+	public static String getClassName(String name) {
+		return name.replaceAll("/", ".");
+	}
+
+	public static ClassPool getClassPool() {
+		return ClassPool.getDefault();
+	}
 }
diff --git a/src/main/java/io/github/proxyhotswap/cglib/EnhancerTransformer.java b/src/main/java/io/github/proxyhotswap/cglib/EnhancerTransformer.java
@@ -2,6 +2,7 @@
 
 import io.github.proxyhotswap.AbstractProxyTransformer;
 import io.github.proxyhotswap.TransformationState;
+import io.github.proxyhotswap.TransformationUtils;
 import io.github.proxyhotswap.javassist.CtClass;
 import io.github.proxyhotswap.javassist.CtMethod;
 
@@ -44,7 +45,8 @@ protected String getInitCall(CtClass cc, String random) throws Exception {
 	@Override
 	protected byte[] getNewByteCode(ClassLoader loader, String className, Class<?> classBeingRedefined)
 			throws Exception {
-		GeneratorParams param = GeneratorSpyTransformer.getGeneratorParams().get(className.replaceAll("/", "."));
+		GeneratorParams param = GeneratorSpyTransformer.getGeneratorParams().get(
+				TransformationUtils.getClassName(className));
 		if (param == null)
 			throw new RuntimeException("No Parameters found for redefinition!");
 
@@ -56,6 +58,12 @@ protected byte[] getNewByteCode(ClassLoader loader, String className, Class<?> c
 		return invoke;
 	}
 
+	@Override
+	protected CtClass getCtClass(byte[] newByteCode, String className) throws Exception {
+		// can use get because generator parameter spy has already loaded it to the clas pool
+		return classPool.get(TransformationUtils.getClassName(className));
+	}
+
 	private Method getGenerateMethod(Object generator) {
 		Method generateMethod = null;
 		Method[] methods = generator.getClass().getMethods();

diff --git a/src/main/java/io/github/proxyhotswap/cglib/GeneratorSpyTransformer.java b/src/main/java/io/github/proxyhotswap/cglib/GeneratorSpyTransformer.java
@@ -21,15 +21,15 @@
 public class GeneratorSpyTransformer implements ClassFileTransformer {
 
 	private static Map<String, GeneratorParams> generatorParams = new ConcurrentHashMap<>();
-	protected static final ClassPool classPool = ClassPool.getDefault();
+	protected static final ClassPool classPool = TransformationUtils.getClassPool();
 
 	public byte[] transform(ClassLoader loader, String className, final Class<?> classBeingRedefined,
 			ProtectionDomain protectionDomain, final byte[] classfileBuffer) {
 		if (classBeingRedefined != null)
 			return null;
 		CtClass cc;
 		try {
-			cc = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
+			cc = classPool.makeClass(new ByteArrayInputStream(classfileBuffer), false);
 			CtClass[] interfaces = cc.getInterfaces();
 			for (CtClass class1 : interfaces) {
 				// We use strings because some libraries repackage cglib to a different namespace to avoid conflicts.
@@ -38,7 +38,7 @@ public byte[] transform(ClassLoader loader, String className, final Class<?> cla
 					for (CtMethod method : declaredMethods) {
 						if (method.getName().equals("generate")
 								&& method.getReturnType().getSimpleName().equals("byte[]")) {
-							return addGenerationParameterCollector(classfileBuffer);
+							return addGenerationParameterCollector(cc);
 						}
 					}
 				}
@@ -52,25 +52,18 @@ public byte[] transform(ClassLoader loader, String className, final Class<?> cla
 	public static void register(Object generatorStrategy, Object classGenerator, byte[] bytes) {
 		CtClass cc = null;
 		try {
-			cc = classPool.makeClass(new ByteArrayInputStream(bytes));
+			cc = classPool.makeClass(new ByteArrayInputStream(bytes), false);
 			generatorParams.put(cc.getName(), new GeneratorParams(generatorStrategy, classGenerator));
 		} catch (IOException | RuntimeException e) {
 			TransformationUtils.logError(e);
-		} finally {
-			TransformationUtils.detachCtClass(cc);
 		}
 	}
 
-	private byte[] addGenerationParameterCollector(final byte[] classfileBuffer) throws IOException, NotFoundException,
+	private byte[] addGenerationParameterCollector(final CtClass cc) throws IOException, NotFoundException,
 			CannotCompileException {
-		CtClass cc2 = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
-		try {
-			CtMethod declaredMethod = cc2.getDeclaredMethod("generate");
-			declaredMethod.insertAfter(getClass().getName() + ".register($0, $1, $_);");
-			return cc2.toBytecode();
-		} finally {
-			TransformationUtils.detachCtClass(cc2);
-		}
+		CtMethod declaredMethod = cc.getDeclaredMethod("generate");
+		declaredMethod.insertAfter(getClass().getName() + ".register($0, $1, $_);");
+		return cc.toBytecode();
 	}
 
 	public static Map<String, GeneratorParams> getGeneratorParams() {