Temporarily escaping ThreadLocal's
Many web frameworks make use of ThreadLocal's to expose request-sensitive information while avoiding the need to pass such details through the call hierarchy. Finding usages of ThreadLocal in the project I currently have open, it's in Guava, EhCache, Tiles, Jackson, Hibernate, Spring and Liferay (to name a few of the 35 jars with matches).
Occasionally we have the need to temporarily reset ThreadLocals to execute some code and then restore those values. One concrete use case is to temporarily change the security context for the current thread to execute some privileged access before reverting back.
The naive approach to this implementation looks something like:
import java.util.Locale; public class BypassThreadLocal_1 { private static final ThreadLocal<Locale> LOCALE_THREAD_LOCAL = new ThreadLocal<Locale>(); private static final ThreadLocal<String> USERNAME_THREAD_LOCAL = new ThreadLocal<String>(); private static final ThreadLocal<Long> USER_ID_THREAD_LOCAL = new ThreadLocal<Long>(); public void doSomething() { Locale locale = LOCALE_THREAD_LOCAL.get(); String username = USERNAME_THREAD_LOCAL.get(); Long userId = USER_ID_THREAD_LOCAL.get(); try { LOCALE_THREAD_LOCAL.set(Locale.getDefault()); USERNAME_THREAD_LOCAL.remove(); USER_ID_THREAD_LOCAL.remove(); System.out.println("Locale: " + LOCALE_THREAD_LOCAL.get()); System.out.println("Username: " + USERNAME_THREAD_LOCAL.get()); System.out.println("UserID: " + USER_ID_THREAD_LOCAL.get()); } finally { LOCALE_THREAD_LOCAL.set(locale); USERNAME_THREAD_LOCAL.set(username); USER_ID_THREAD_LOCAL.set(userId); } } }
The problems here are pretty obvious:
the code is not reusable, if I need to do this again elsewhere then ctrl-c and ctrl-v is going to get a work over
the get current value, remove/replace, execute, restore pattern is duplicated per ThreadLocal
There's also the third feature/issue of needing to know exactly what to change. Generally I'd rather be in control of what's being modified but sometimes it's just quicker to start anew (we'll get to that shortly).
So, in order to address the first two issues here, let's do some refactoring, starting with the code we wish we had (written using the Java 8 lambda syntax, because it's easier to - substitute anonymous or top level classes at your leisure):
// BypassThreadLocal_2 import java.util.Locale; import java.util.concurrent.Callable; public class BypassThreadLocal_2 { private static final ThreadLocal<Locale> LOCALE_THREAD_LOCAL = new ThreadLocal<Locale>(); private static final ThreadLocal<String> USERNAME_THREAD_LOCAL = new ThreadLocal<String>(); private static final ThreadLocal<Long> USER_ID_THREAD_LOCAL = new ThreadLocal<Long>(); public void doSomething() throws Exception { new BypassBuilder() .set(LOCALE_THREAD_LOCAL, Locale.getDefault()) .remove(USERNAME_THREAD_LOCAL) .remove(USER_ID_THREAD_LOCAL) .call(() -> { System.out.println("Locale: " + LOCALE_THREAD_LOCAL.get()); System.out.println("Username: " + USERNAME_THREAD_LOCAL.get()); System.out.println("UserID: " + USER_ID_THREAD_LOCAL.get()); return null; }); } }
So, obviously we need a BypassBuilder we can build (sic) upon. Here's the skeleton:
// BypassBuilder.java import java.util.concurrent.Callable; public class BypassBuilder { public <T> T call(final Callable<T> callable) throws Exception { // Do something with callable } public <T> BypassBuilder remove(final ThreadLocal<T> threadLocal) { // Do something with the thread local return this; } public <T> BypassBuilder set(final ThreadLocal<T> threadLocal, final T value) { // Do something with the thread local and value return this; } }
Looking at this, we have two things to manage: Mutation of the ThreadLocal's and calling of the supplied Callable. We'll throw together those SAM interfaces (we'll implement them shortly) as well as provide the framework which will manage the ThreadLocal get/mutate/reset process:
// BypassBuilder.java import java.util.concurrent.Callable; public class BypassBuilder { static interface CallableCaller { <T> T call(Callable<T> callable) throws Exception; } static interface ThreadLocalMutator<V> { void mutate(ThreadLocal<V> threadLocal); } static class MutatingCaller<T> implements CallableCaller { private final CallableCaller delegate; private final ThreadLocalMutator<T> mutator; private final ThreadLocal<T> threadLocal; public MutatingCaller(final CallableCaller delegate, final ThreadLocal<T> threadLocal, final ThreadLocalMutator<T> mutator) { this.delegate = delegate; this.mutator = mutator; this.threadLocal = threadLocal; } @Override public <V> V call(final Callable<V> callable) throws Exception { T original = threadLocal.get(); try { mutator.mutate(threadLocal); return delegate.call(callable); } finally { threadLocal.set(original); } } } // The rest of the class as above }
So now we have a couple of SAM interfaces that we can implement via lambdas to mutate ThreadLocals and pass around nested CallableCallers, and a nested class which manages (sets and resets)the ThreadLocal state through mutation. So, we can finish this sucker off. The class in full:
// BypassBuilder.java import java.util.concurrent.Callable; public class BypassBuilder { static interface CallableCaller { <T> T call(Callable<T> callable) throws Exception; } static interface ThreadLocalMutator<V> { void mutate(ThreadLocal<V> threadLocal); } static class MutatingCaller<T> implements CallableCaller { private final CallableCaller delegate; private final ThreadLocalMutator<T> mutator; private final ThreadLocal<T> threadLocal; public MutatingCaller(final CallableCaller delegate, final ThreadLocal<T> threadLocal, final ThreadLocalMutator<T> mutator) { this.delegate = delegate; this.mutator = mutator; this.threadLocal = threadLocal; } @Override public <V> V call(final Callable<V> callable) throws Exception { T original = threadLocal.get(); try { mutator.mutate(threadLocal); return delegate.call(callable); } finally { threadLocal.set(original); } } } private CallableCaller caller = (callable) -> { return callable.call(); }; public <T> call(final Callable<T> callable) throws Exception { return caller.call(callable); } public <T> ThreadLocalBypassBuilder clear(final ThreadLocal<T> threadLocal) { caller = new MutatingCaller<T>( caller, threadLocal, (threadLocal) -> { threadLocal.remove(); } ); return this; } public <T> ThreadLocalBypassBuilder set(final ThreadLocal<T> threadLocal, final T value) { caller = new MutatingCaller<T>( caller, threadLocal, (threadLocal) -> { threadLocal.set(value); } ); return this; } }
Let's briefly describe what's going on here:
The builder tracks a single instance of CallableCaller which will eventually be called to do "stuff"
Initially this field is set to just return the value of calling the passed parameter
When remove or set is called, the field is updated with a new value which:
references the previous field value (think turtles all the way down)
has access to the ThreadLocal passed to remove or set accordingly
knows what to do to the thread local based on the calling method
This is all well and good, but there's on small issue... We assume we know each and every ThreadLocal to clear (or that we have enough patience to clear them all). I was shown an unexpectedly elegant and subtle solution to this by a colleague (Chun). Let's add the target code to our business logic first:
// BypassThreadLocal_2 import java.util.Locale; import java.util.concurrent.Callable; public class BypassThreadLocal_2 { private static final ThreadLocal<Locale> LOCALE_THREAD_LOCAL = new ThreadLocal<Locale>(); private static final ThreadLocal<String> USERNAME_THREAD_LOCAL = new ThreadLocal<String>(); private static final ThreadLocal<Long> USER_ID_THREAD_LOCAL = new ThreadLocal<Long>(); public void doSomething() throws Exception { new BypassBuilder() .isolate(() -> { System.out.println("Locale: " + LOCALE_THREAD_LOCAL.get()); System.out.println("Username: " + USERNAME_THREAD_LOCAL.get()); System.out.println("UserID: " + USER_ID_THREAD_LOCAL.get()); return null; }); } }
And now for the clever implementation:
// BypassBuilder.java import java.util.concurrent.Callable; public class BypassBuilder { // Implementation as above public <T> T isolate(final Callable<T> callable) throws Exception { return Executors.newSingleThreadExecutor() .submit(callable) .get(); } }
So simple and clean, isn't it? A single thread executor spawns a new thread to execute the task and returns a future, which we then get() (and block while we wait) the result of. Because we've spawned a new thread, all of the existing ThreadLocal's are associated elsewhere and not visible!