spring CachingByTypeBeanFactory to cache lookups by type

CachingByTypeBeanFactory.java

package com.example.spring;

import org.apache.log4j.Logger;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;

import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Caches certain calls to get bean names for type. Profiling showed that significant cpu time was spent on
 * reflection because this was not cached. It does not change, so should be safe to cache.
 */
public class CachingByTypeBeanFactory extends DefaultListableBeanFactory {
    private static Logger log = Logger.getLogger(CachingByTypeBeanFactory.class);
    ConcurrentHashMap<TypeKey, String[]> cachedBeanNamesForType = new ConcurrentHashMap<TypeKey, String[]>();

    @Override
    public String[] getBeanNamesForType(Class type) {
        return getBeanNamesForType(type, true, true);
    }

    @Override
    public String[] getBeanNamesForType(Class type, boolean includeNonSingletons, boolean allowEagerInit) {
        TypeKey typeKey = new TypeKey(type, includeNonSingletons, allowEagerInit);
        if (cachedBeanNamesForType.containsKey(typeKey)) {
            log.debug("will retrieve from cache: " + typeKey);
            return cachedBeanNamesForType.get(typeKey);
        }
        String[] value = super.getBeanNamesForType(type, includeNonSingletons, allowEagerInit);
        if (log.isDebugEnabled()) {
            log.debug("will add to cache: " + typeKey + " " + Arrays.asList(value));
        }
        cachedBeanNamesForType.putIfAbsent(typeKey, value);
        return value;
    }

    // This is the input parameters, which we memoize.
    // We conservatively cache based on the possible parameters passed in. Assuming that state does not change within the
    // super.getBeanamesForType() call between subsequent requests.
    static class TypeKey {
        Class type;
        boolean includeNonSingletons;
        boolean allowEagerInit;

        TypeKey(Class type, boolean includeNonSingletons, boolean allowEagerInit) {
            this.type = type;
            this.includeNonSingletons = includeNonSingletons;
            this.allowEagerInit = allowEagerInit;
        }

        @Override
        public String toString() {
            return "TypeKey{" + type + " " + includeNonSingletons + " " + allowEagerInit + "}";
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;

            TypeKey typeKey = (TypeKey) o;

            if (allowEagerInit != typeKey.allowEagerInit) return false;
            if (includeNonSingletons != typeKey.includeNonSingletons) return false;
            if (type != null ? !type.equals(typeKey.type) : typeKey.type != null) return false;

            return true;
        }

        @Override
        public int hashCode() {
            int result = type != null ? type.hashCode() : 0;
            result = 31 * result + (includeNonSingletons ? 1 : 0);
            result = 31 * result + (allowEagerInit ? 1 : 0);
            return result;
        }
    }

}

CachingByTypeBeanFactoryTest.java

package com.example.spring;

import com.example.common.testing.MoreAsserts;
import org.junit.Before;
import org.junit.Test;

import java.util.concurrent.ConcurrentHashMap;

import static org.junit.Assert.assertEquals;

public class CachingByTypeBeanFactoryTest { // not the best test
    CachingByTypeBeanFactory bf = new CachingByTypeBeanFactory();

    int putsCalled;
    ConcurrentHashMap map = new ConcurrentHashMap() {
        @Override
        public Object putIfAbsent(Object key, Object value) {
            putsCalled++;
            return super.putIfAbsent(key, value);
        }
    };

    @Before
    public void setup() {
        bf.cachedBeanNamesForType = map;
    }

    @Test
    public void shouldCacheTypeSig() {
        bf.getBeanNamesForType(Object.class);
        bf.getBeanNamesForType(Object.class);
        assertEquals(1, putsCalled);
    }

    @Test
    public void shouldCache3ArgsSig() {
        bf.getBeanNamesForType(Integer.class, true, false);
        bf.getBeanNamesForType(Integer.class, true, false);
        bf.getBeanNamesForType(Object.class, true, false);
        bf.getBeanNamesForType(Object.class, true, true);
        bf.getBeanNamesForType(Object.class, true, true);
        assertEquals(3, putsCalled);
    }

    @Test
    public void typeKeyHashEquals() {
        CachingByTypeBeanFactory.TypeKey key1 = new CachingByTypeBeanFactory.TypeKey(Object.class, true, true);
        MoreAsserts.checkEqualsAndHashCodeMethods(key1, new CachingByTypeBeanFactory.TypeKey(Object.class, true, true), true);

        MoreAsserts.checkEqualsAndHashCodeMethods(key1, new CachingByTypeBeanFactory.TypeKey(Object.class, true, false), false);
        MoreAsserts.checkEqualsAndHashCodeMethods(key1, new CachingByTypeBeanFactory.TypeKey(Object.class, false, true), false);
        MoreAsserts.checkEqualsAndHashCodeMethods(key1, new CachingByTypeBeanFactory.TypeKey(Integer.class, true, true), false);
    }
}

CachingWebApplicationContext.java

package com.example.spring;

import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.web.context.support.XmlWebApplicationContext;

/**
 * Takes advantage of a spring hook to create a caching Bean Factory. Profiling identified the
 * high amount of reflection in spring 2.5.6 as a very large proportion of CPU time per request.
 * This class metadata should not change, so we cache it.
 */
public class CachingWebApplicationContext extends XmlWebApplicationContext {

    @Override
    protected DefaultListableBeanFactory createBeanFactory() {
        return new CachingByTypeBeanFactory();
    }
}

CachingWebApplicationContextTest.java

package com.example.spring;

import org.junit.Test;

import static org.junit.Assert.assertEquals;

public class CachingWebApplicationContextTest {

    @Test
    public void shouldCreateCachingByTypeBeanFactory() {
        assertEquals(CachingByTypeBeanFactory.class, new CachingWebApplicationContext().createBeanFactory().getClass());
    }
}

web.xml

<!-- assuming you use the standard org.springframework.web.context.ContextLoaderListener listener, you'll need to add this config param -->
    <context-param>
	<param-name>contextClass</param-name>
	<param-value>com.example.spring.CachingWebApplicationContext</param-value>
    </context-param>