Saturday, October 20, 2007

ClassLoader (ab)use

A weird piece of code it was and I had the misfortune of having to get around this very piece of code - a hand me down, from an old version of the product.

It was some kind of a store/map where the keys were Classes instead of Strings as keys, which you would've seen in most "regular" Maps. So, it automatically limited the range of keys one could use. New Classes had to be created and compiled and then used as keys, if required. I had to use this package because there a lot of business logic tied to this store and it couldn't be modified. There were a lot of versions that were already in production.

I needed to use the very same store but my keys were going to be dynamic. The only choices I had was to either compile new classes beforehand or generate new Classes on the fly in the JVM using BCEL or ASM or any of those scary Bytecode libraries. That was when I realized that there was a simpler solution. One that didn't really need an over-the-top solution. I chose to use (rather abuse) ClassLoaders.

I realized that I didn't really need the Key-Classes anywhere other than to store data into the map. So, I created an empty Class and then every time I needed a new Key-Class I would just load my fixed Class from a new ClassLoader. Since the same Class can be loaded into the JVM by multiple ClassLoaders, the Store/map would have no idea that they were all the same Classes but from different ClassLoaders. So, I created a small custom ClassLoader that would do this everytime a new Key was required.

I must mention that I was allowed to sub-Class this store.


This is what the "weird" old Store looks like:


public class AbsurdLegacyStore {
protected final Map<Class, Set<Object>> perTypeStore;

public AbsurdLegacyStore() {
this.perTypeStore = new HashMap<Class, Set<Object>>();
}

public Map<Class, Set<Object>> getPerTypeStore() {
return perTypeStore;
}

public void addToStoreAndDoWork(Employee e) {
Set<Object> set = perTypeStore.get(Employee.class);
if (set == null) {
set = new HashSet<Object>();
perTypeStore.put(Employee.class, set);
}

set.add(e);
}

public void addToStoreAndDoWork(Customer c) {
Set<Object> set = perTypeStore.get(Customer.class);
if (set == null) {
set = new HashSet<Object>();
perTypeStore.put(Customer.class, set);
}

set.add(c);
}

public void addToStoreAndDoWork(Supplier s) {
Set<Object> set = perTypeStore.get(Supplier.class);
if (set == null) {
set = new HashSet<Object>();
perTypeStore.put(Supplier.class, set);
}

set.add(s);
}
}


All the fixed key Classes:


public class Customer {
protected final String name;

public Customer(String name) {
this.name = name;
}

public String getName() {
return name;
}

@Override
public String toString() {
return name;
}
}



public class Employee {
protected final String name;

public Employee(String name) {
this.name = name;
}

public String getName() {
return name;
}

@Override
public String toString() {
return name;
}
}



public class Supplier {
protected final String name;

public Supplier(String name) {
this.name = name;
}

public String getName() {
return name;
}

@Override
public String toString() {
return name;
}
}



public class ChannelPartner {
protected final String name;

public ChannelPartner(String name) {
this.name = name;
}

@Override
public String toString() {
return name;
}
}


The "hacky" sub-Class of the old Store with its own custom ClassLoader:


public class HackyTypeBasedStore extends AbsurdLegacyStore {
protected Map<String, Class> dynamicSalesRegionTypes;

public HackyTypeBasedStoere() {
this.dynamicSalesRegionTypes = new HashMap<String, Class>();
}

public void registerNewSalesRegion(String salesRegionName) throws ClassNotFoundException {
if (dynamicSalesRegionTypes.containsKey(salesRegionName)) {
return;
}

ClassLoader customLoader = new CustomClassLoader();
Class dynamicType = customLoader.loadClass(SalesRegion.class.getName());
dynamicSalesRegionTypes.put(salesRegionName, dynamicType);
}

public void addToStoreAndDoWork(String salesRegionName, ChannelPartner partner) {
Class regionType = dynamicSalesRegionTypes.get(salesRegionName);

Set<Object> set = perTypeStore.get(regionType);
if (set == null) {
set = new HashSet<Object>();
perTypeStore.put(regionType, set);
}

set.add(partner);
}

public static void main(String[] args) throws ClassNotFoundException {
HackyTypeBasedStore store = new HackyTypeBasedStore();

String salesRegion1 = "Europe";
String salesRegion2 = "Americas";
String salesRegion3 = "Asia";

store.registerNewSalesRegion(salesRegion1);
store.registerNewSalesRegion(salesRegion2);
store.registerNewSalesRegion(salesRegion3);

ChannelPartner partnerABCD = new ChannelPartner("ABCD");
store.addToStoreAndDoWork(salesRegion1, partnerABCD);
store.addToStoreAndDoWork(salesRegion3, partnerABCD);

ChannelPartner partnerKLMN = new ChannelPartner("KLMN");
store.addToStoreAndDoWork(salesRegion2, partnerKLMN);
store.addToStoreAndDoWork(salesRegion3, partnerKLMN);

ChannelPartner partnerWXYZ = new ChannelPartner("WXYZ");
store.addToStoreAndDoWork(salesRegion2, partnerWXYZ);

// --------

Customer customer = new Customer("Customer-A");
store.addToStoreAndDoWork(customer);

Employee employee = new Employee("Employee-1");
store.addToStoreAndDoWork(employee);

Supplier supplier5 = new Supplier("Supplier-5");
store.addToStoreAndDoWork(supplier5);

Supplier supplier6 = new Supplier("Supplier-6");
store.addToStoreAndDoWork(supplier6);

// ---------

Map<Class, Set<Object>> wickedMap = store.getPerTypeStore();
for (Class key : wickedMap.keySet()) {
Set<Object> set = wickedMap.get(key);

System.out.println("Key: " + key.getName() + ", ClassLoader: " + key.getClassLoader());
System.out.println("Data: " + set);
System.out.println();
}
}

// --------

public static class CustomClassLoader extends URLClassLoader {
protected final String nameAsPath;

public CustomClassLoader() {
super(new URL[0]);

String name = SalesRegion.class.getName().replace('.', '/');
this.nameAsPath = "/" + name + ".class";
}

@Override
public Class<?> loadClass(String name) throws ClassNotFoundException {
if (name.equals(SalesRegion.class.getName()) == false) {
// Delegate to System ClassLoader.
return getParent().loadClass(name);
}

Class clazz = null;

try {
InputStream inputStream = SalesRegion.class.getResourceAsStream(nameAsPath);
if (inputStream == null) {
throw new NullPointerException("Could not find Class file");
}

ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

byte[] buffer = new byte[512];
int c = 0;
while ((c = inputStream.read(buffer)) != -1) {
outputStream.write(buffer, 0, c);
}

byte[] classAsBytes = outputStream.toByteArray();

clazz = defineClass(name, classAsBytes, 0, classAsBytes.length);
}
catch (Exception e) {
throw new ClassNotFoundException(e.getMessage(), e);
}

return clazz;
}
}
}


The dummy Key that would be loaded by multiple ClassLoaders:


public interface SalesRegion {
}


Sample output:


Key: temp.demo.SalesRegion, ClassLoader: temp.demo.HackyTypeBasedStore$CustomClassLoader@c2ea3f
Data: [ABCD]

Key: temp.demo.Customer, ClassLoader: sun.misc.Launcher$AppClassLoader@fabe9
Data: [Customer-A]

Key: temp.demo.Supplier, ClassLoader: sun.misc.Launcher$AppClassLoader@fabe9
Data: [Supplier-6, Supplier-5]

Key: temp.demo.SalesRegion, ClassLoader: temp.demo.HackyTypeBasedStore$CustomClassLoader@1034bb5
Data: [KLMN, WXYZ]

Key: temp.demo.SalesRegion, ClassLoader: temp.demo.HackyTypeBasedStore$CustomClassLoader@b162d5
Data: [KLMN, ABCD]

Key: temp.demo.Employee, ClassLoader: sun.misc.Launcher$AppClassLoader@fabe9
Data: [Employee-1]

0 comments: