Tuesday, December 27, 2011

Using a custom classloader

So, recently I was using a code generation tool, and needed to invoke the code generation multiple times on the same input files but pass in different configuration settings each time. But I was disappointed when the second (and subsequent) invocations would not generate any output files. After downloading the source code for the tool and spending some time debugging, I discovered that there was a private static Set<File> files field in a class which would keep track of the files it had already handled.
Now to get around this there were two solutions:

  1. After each invocation of code generation, clear the set using reflection
  2. Use a different classloader to load the code generator each time
Now for this particular scenario, it was easier to simply use reflection. However, there can be cases when you want to concurrently use multiple instances of some class but avoid cross-talk, e.g. due to static fields for instance. A custom classloader can be used in such cases. You will have to resort to reflection for making the actual method invocations though.

Example:
package com.usta;

public class Service {
    private static final Service INSTANCE = new Service();

    public static Service getInstance() { return INSTANCE; }

    //private constructor
    private Service() {}

    public void doIt(String arg) {
        System.out.println("Hey " + arg + " " + hashCode());
    }
}

package com.usta;

import java.io.DataInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FilenameFilter;
import java.io.IOException;
import java.io.InputStream;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

public class CustomClassLoader extends ClassLoader {
    private final File                      classpathDir;
    private ConcurrentMap<String, Class<?>> classes = new ConcurrentHashMap<String, Class<?>>();

    public CustomClassLoader(String directory) {
        super();
        this.classpathDir = new File(directory);
        if (!classpathDir.isDirectory())
            throw new RuntimeException("Not a folder " + directory);
    }

    @Override
    public Class<?> loadClass(String name) throws ClassNotFoundException {
        return name.startsWith("com.usta") ? findClass(name) : findSystemClass(name);
    }

    // TODO support inner classes
    private String binaryNameToFileName(String binaryName) {
        return binaryName.replaceAll("\\.", "/") + ".class";
    }

    @Override
    protected Class<?> findClass(String clazz) throws ClassNotFoundException {
        if (classes.containsKey(clazz)) {
            return classes.get(clazz);
        }
        synchronized (clazz.intern()) {
            // we might have already defined the class on a separate thread
            if (classes.containsKey(clazz)) {
                return classes.get(clazz);
            }
            InputStream is = null;
            int size = 0;
            for (File f : classpathDir.listFiles(new FilenameFilter() {
                public boolean accept(File arg0, String arg1) {
                    return arg1.endsWith(".jar");
                }
            })) {
                try {
                    JarFile jarFile = new JarFile(f);
                    JarEntry clazzEntry = jarFile.getJarEntry(binaryNameToFileName(clazz));
                    if (clazzEntry != null) {
                        is = jarFile.getInputStream(clazzEntry);
                        size = (int) clazzEntry.getSize();
                        break;
                    }
                } catch (IOException e) {
                    throw new ClassNotFoundException(e.getMessage());
                }
            }
            if (is == null) {
                String file = classpathDir.getPath() + System.getProperty("file.separator")
                        + binaryNameToFileName(clazz);
                File classFile = new File(file);
                if (!classFile.exists() || !classFile.isFile())
                    return super.findClass(clazz);
                size = (int) classFile.length();

                try {
                    is = new FileInputStream(classFile);
                } catch (FileNotFoundException e) {
                    throw new ClassNotFoundException(e.getMessage());
                }
            }

            DataInputStream dis = new DataInputStream(is);
            byte buff[] = new byte[size];
            try {
                dis.readFully(buff);
                dis.close();
            } catch (IOException e) {
                throw new ClassNotFoundException(e.getMessage());
            }

            classes.putIfAbsent(clazz, super.defineClass(clazz, buff, 0, size));
            return classes.get(clazz);
        }
    }
}

Now the use of this custom class loader to load multiple instances of Service would be:


Service service = Service.getInstance();
service.doIt("Buddy");
ClassLoader loader = new CustomClassLoader("/users/ustamansangat/workspace/test/target");
Class<?> serviceClass = loader.loadClass("com.usta.Service");
Object service2 = serviceClass.getMethod("getInstance").invoke(null);
serviceClass.getMethod("doIt", String.class).invoke(service2, "Dude");