Migrated attachment for Guice issue 100, comment 8

GuiceThreadScope.txt

Simple thread scope implementation for Guice, Apache 2.0 licensed. Enjoy!

Robbie Vanbrabant (robbie.vanbrabant @ google mail)


/**
* @author Robbie Vanbrabant
* @see CustomScopes#THREAD
*/
public class ThreadCache {
    // use lazy init to avoid memory overhead when not using the scope?
    private static final ThreadLocal<Cache> THREAD_LOCAL = 
        new ThreadLocal<Cache>() {
            @Override protected Cache initialValue() {
                return new Cache();
            }
    };
    
    public ThreadCache(){}
    
    public Cache getCache() {
        return THREAD_LOCAL.get();
    }

    /**
    * Execute this if you plan to reuse the same thread,
    * e.g. in a servlet environment threads might get
    * reused. Preferably, call this method in a finally
    * block to make sure that it executes, so that you
    * avoid possible memory leaks.
    */
    public void reset() {
        THREAD_LOCAL.remove();
    }
    
    /**
    * Cache class for type capture and minimizing
    * ThreadLocal lookups.
    */
    public static class Cache {
        private Map<Key<?>, Object> map = new HashMap<Key<?>, Object>();
        public Cache() {}
        // suppress warnings because the add method
        // captures the type
        @SuppressWarnings("unchecked")
        public <T> T get(Key<T> key) {
            return (T) map.get(key);
        }
        public <T> void add(Key<T> key, T value) {
            map.put(key, value);
        }
    }
}




/**
* @author Robbie Vanbrabant
* @see CustomScopes#THREAD
*/
public class ThreadScope implements Scope {
    private ThreadCache thread;
    ThreadScope(ThreadCache thread) {
        this.thread = thread;
    }
    
    /**
    * @see com.google.inject.Scope#scope(com.google.inject.Key, com.google.inject.Provider)
    */
    public <T> Provider<T> scope(final Key<T> key, final Provider<T> creator) {
        return new Provider<T>() {
            public T get() {
                Cache cache = thread.getCache();
                T value = cache.get(key);
                if (value == null) {
                    value = creator.get();
                    cache.add(key, value);
                }
                return value;
            }
        };
    }
}




/**
* Guice scope additons.
* @author Robbie Vanbrabant
*/
public class CustomScopes {
    /**
    * Thread scope, backed by a {@link java.lang.ThreadLocal}.
    * Example usage:
    * <pre>
    * Injector i = Guice.createInjector(new Module() {
    *     public void configure(Binder binder) {
    *         binder.bindScope(ThreadScoped.class, CustomScopes.THREAD);
    *         binder.bind(ThreadCache.class).in(Scopes.SINGLETON);
    *         binder.bind(SomeClass.class).in(CustomScopes.THREAD);
    *     }
    * });
    * </pre>
    * In thread pooling scenario's, never forget to reset the scope
    * at the end of a request:
    * <pre>
    * i.getInstance(ThreadCache.class).reset();
    * </pre>
    * Note that if you create new threads within this scope, they will
    * start with a clean slate.
    */
    public static final Scope THREAD = new ThreadScope(new ThreadCache());
}




/**
* Indicates that an object needs to be
* {@link CustomScopes#THREAD} scoped.
* @author Robbie Vanbrabant
* @see CustomScopes#THREAD
*/
@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
@ScopeAnnotation
public @interface ThreadScoped {}




public class ThreadScopeTest extends TestCase {
    private class SomeClass {@Inject public SomeClass() {}}
    private static Injector i = Guice.createInjector(new Module() {
        public void configure(Binder binder) {
            binder.bindScope(ThreadScoped.class, CustomScopes.THREAD);
            binder.bind(ThreadCache.class).in(Scopes.SINGLETON);
            binder.bind(SomeClass.class).in(CustomScopes.THREAD);
        }
    });

    public void testReset() {
        SomeClass someClass = i.getInstance(SomeClass.class);
        assertTrue(someClass == i.getInstance(SomeClass.class));
        ThreadCache c = i.getInstance(ThreadCache.class);
        c.reset();
        assertFalse(someClass == i.getInstance(SomeClass.class));
    }

    public void testLocality() {
        SomeClass someClass = i.getInstance(SomeClass.class);
        final SomeClass[] innerSomeClass = new SomeClass[1];
        final CountDownLatch done = new CountDownLatch(1);

        new Thread(new Runnable() {
            public void run() {
                innerSomeClass[0] = i.getInstance(SomeClass.class);
                done.countDown();
            }
        }).start();

        try {
            done.await();
        } catch (InterruptedException e) {
            fail("unexpected thread interruption");
        }

        assertFalse(someClass == innerSomeClass[0]);
    }
    
    // probably makes no sense to test this,
    // but it makes me sleep better at night
    public void testConcurrency() {
        final CountDownLatch done = new CountDownLatch(1);
        Executor executor = Executors.newFixedThreadPool(50);
        final Injector fi = i;
        for (int i = 0; i < 200; i++) {
            final int index = i;
            executor.execute(new Runnable() {
                public void run() {
                    System.out.println("test "+index);
                    assertTrue(fi.getInstance(SomeClass.class) == fi.getInstance(SomeClass.class));
                    fi.getInstance(ThreadCache.class).reset();
                    if (index == 199)
                        done.countDown();
                }
            });
            
        }
        try {
            done.await();
        } catch (InterruptedException e) {
            fail("unexpected thread interruption");
        }
    }
}