ContextStore

src/main/java/mina/context/ContextStore.java

@@ -0,0 +1,11 @@
+package mina.context;
+
+public interface ContextStore {
+    MinaContext createOrGetContext();
+
+    MinaContext getContext();
+
+    void removeContext();
+
+    void close();
+}

src/main/java/mina/context/GlobalContextStore.java

@@ -0,0 +1,59 @@
+package mina.context;
+
+import java.lang.ref.WeakReference;
+import java.util.concurrent.atomic.AtomicReference;
+
+public class GlobalContextStore implements ContextStore {
+    private final AtomicReference<MinaContext> globalContext = new AtomicReference<>();
+    private final AtomicReference<WeakReference<Thread>> ownerThread = new AtomicReference<>();
+
+    @Override
+    public synchronized MinaContext createOrGetContext() {
+        assertParallelAccessToGlobalContext();
+        return globalContext.updateAndGet(context -> context == null ? new MinaContext() : context);
+    }
+
+    @Override
+    public MinaContext getContext() {
+        return globalContext.get();
+    }
+
+    @Override
+    public synchronized void removeContext() {
+        assertParallelAccessToGlobalContext();
+        globalContext.set(null);
+        ownerThread.set(null);
+    }
+
+    @Override
+    public void close() {
+        removeContext();
+    }
+
+    private synchronized void assertParallelAccessToGlobalContext() {
+        if (globalContext.get() == null) {
+            return;
+        }
+
+        Thread currentThread = Thread.currentThread();
+
+        Thread ownerThread = this.ownerThread.updateAndGet((weak) -> {
+            Thread thread = weak != null ? weak.get() : null;
+            if (thread == null || !thread.isAlive()) {
+                return new WeakReference<>(currentThread);
+            } else {
+                return weak;
+            }
+        }).get();
+
+        if (currentThread == ownerThread) {
+            return;
+        }
+
+        throw new IllegalStateException(
+                "Mina is configured to use GLOBAL context in a single thread tests but multi access to the context detected. " +
+                        "Current thread is [" + currentThread.getName() + "] but the owner thread is [" + (ownerThread != null ? ownerThread.getName() : null) + "]"
+        );
+    }
+
+}

src/main/java/mina/context/MinaContextHolder.java

@@ -1,97 +1,57 @@
 package mina.context;
 
-import java.lang.ref.WeakReference;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 public final class MinaContextHolder {
     public static final String GLOBAL_CONTEXT_PROPERTY_KEY = "mina.context.global";
 
-    private static final ThreadLocal<MinaContext> THREAD_LOCAL_CONTEXT = new InheritableThreadLocal<>();
-
-    private static final AtomicBoolean useGlobalContext = new AtomicBoolean(false);
-    private static final AtomicReference<MinaContext> GLOBAL_CONTEXT = new AtomicReference<>();
-    private static final AtomicReference<WeakReference<Thread>> GLOBAL_CONTEXT_THREAD_OWNER = new AtomicReference<>();
+    private static final AtomicReference<ContextStore> CONTEXT_STORE = new AtomicReference<>();
 
     static {
-        useGlobalContext.set(Boolean.parseBoolean(System.getProperty(GLOBAL_CONTEXT_PROPERTY_KEY)));
+        boolean useGlobalContext = Boolean.parseBoolean(System.getProperty(GLOBAL_CONTEXT_PROPERTY_KEY));
+        if (useGlobalContext) {
+            CONTEXT_STORE.set(new GlobalContextStore());
+        } else {
+            CONTEXT_STORE.set(new ThreadLocalContextStore());
+        }
     }
 
     private MinaContextHolder() {
     }
 
     public static MinaContext createOrGetContext() {
-        assertParallelAccessToGlobalContext();
-        if (useGlobalContext.get()) {
-            return GLOBAL_CONTEXT.updateAndGet(context -> context == null ? new MinaContext() : context);
-        } else {
-            return createOrGetThreadLocalContext();
-        }
+        return CONTEXT_STORE.get().createOrGetContext();
     }
 
     public static MinaContext getContext() {
-        if (useGlobalContext.get()) {
-            return GLOBAL_CONTEXT.get();
-        } else {
-            return THREAD_LOCAL_CONTEXT.get();
-        }
-    }
-
-    private static MinaContext createOrGetThreadLocalContext() {
-        MinaContext minaContext = THREAD_LOCAL_CONTEXT.get();
-        if (minaContext == null) {
-            minaContext = new MinaContext();
-            THREAD_LOCAL_CONTEXT.set(minaContext);
-        }
-        return minaContext;
+        return CONTEXT_STORE.get().getContext();
     }
 
     public static void removeContext() {
-        assertParallelAccessToGlobalContext();
-        if (useGlobalContext.get()) {
-            GLOBAL_CONTEXT.set(null);
-            GLOBAL_CONTEXT_THREAD_OWNER.set(null);
-        } else {
-            THREAD_LOCAL_CONTEXT.remove();
-        }
+        CONTEXT_STORE.get().removeContext();
     }
 
     public static void useGlobalContext() {
-        useGlobalContext.set(true);
-    }
-
-    public static void useThreadLocalContext() {
-        assertParallelAccessToGlobalContext();
-        useGlobalContext.set(false);
-    }
-
-    private static synchronized void assertParallelAccessToGlobalContext() {
-        if (!useGlobalContext.get()) {
-            return;
-        }
+        CONTEXT_STORE.updateAndGet((store) -> {
+            if (store instanceof GlobalContextStore) {
+                return store;
+            }
 
-        if (GLOBAL_CONTEXT.get() == null) {
-            return;
-        }
+            store.close();
 
-        Thread currentThread = Thread.currentThread();
+            return new GlobalContextStore();
+        });
+    }
 
-        Thread ownerThread = GLOBAL_CONTEXT_THREAD_OWNER.updateAndGet((weak) -> {
-            Thread thread = weak != null ? weak.get() : null;
-            if (thread == null || !thread.isAlive()) {
-                return new WeakReference<>(currentThread);
-            } else {
-                return weak;
+    public static void useThreadLocalContext() {
+        CONTEXT_STORE.updateAndGet((store) -> {
+            if (store instanceof ThreadLocalContextStore) {
+                return store;
             }
-        }).get();
 
-        if (currentThread == ownerThread) {
-            return;
-        }
+            store.close();
 
-        throw new IllegalStateException(
-                "Mina is configured to use GLOBAL context in a single thread tests but multi access to the context detected. " +
-                        "Current thread is [" + currentThread.getName() + "] but the owner thread is [" + (ownerThread != null ? ownerThread.getName() : null) + "]"
-        );
+            return new ThreadLocalContextStore();
+        });
     }
 }

src/main/java/mina/context/ThreadLocalContextStore.java

@@ -0,0 +1,32 @@
+package mina.context;
+
+public class ThreadLocalContextStore implements ContextStore {
+    private final ThreadLocal<MinaContext> threadLocalContext = new InheritableThreadLocal<>();
+
+    @Override
+    public MinaContext createOrGetContext() {
+        MinaContext minaContext = threadLocalContext.get();
+        if (minaContext == null) {
+            minaContext = new MinaContext();
+            threadLocalContext.set(minaContext);
+        }
+        return minaContext;
+
+    }
+
+    @Override
+    public MinaContext getContext() {
+        return threadLocalContext.get();
+    }
+
+    @Override
+    public void removeContext() {
+        threadLocalContext.remove();
+
+    }
+
+    @Override
+    public void close() {
+        // Nothing
+    }
+}

src/test/java/mina/MinaContextHolderTest.java

@@ -0,0 +1,15 @@
+package mina;
+
+import mina.context.MinaContextHolder;
+import org.junit.jupiter.api.Test;
+
+public class MinaContextHolderTest {
+    @Test
+    public void testDontChangeContextStore() {
+        MinaContextHolder.useGlobalContext();
+        MinaContextHolder.useGlobalContext();
+
+        MinaContextHolder.useThreadLocalContext();
+        MinaContextHolder.useThreadLocalContext();
+    }
+}