add option to release MtContext after run #26

src/main/java/com/alibaba/mtc/MtContextCallable.java

@@ -19,12 +19,15 @@ import java.util.concurrent.Callable;
  * @since 0.9.0
  */
 public final class MtContextCallable<V> implements Callable<V> {
-    private final Map<MtContextThreadLocal<?>, Object> copied;
+    private volatile Map<MtContextThreadLocal<?>, Object> copied;
     private final Callable<V> callable;
+    private final boolean releaseMtContextAfterCall;
 
-    private MtContextCallable(Callable<V> callable) {
-        copied = MtContextThreadLocal.copy();
+
+    private MtContextCallable(Callable<V> callable, boolean releaseMtContextAfterCall) {
         this.callable = callable;
+        this.releaseMtContextAfterCall = releaseMtContextAfterCall;
+        copied = MtContextThreadLocal.copy();
     }
 
     /**
@@ -32,13 +35,18 @@ public final class MtContextCallable<V> implements Callable<V> {
      */
     @Override
     public V call() throws Exception {
+        if (null == copied) {
+            throw new IllegalStateException("MtContext is released!");
+        }
         Map<MtContextThreadLocal<?>, Object> backup = MtContextThreadLocal.backupAndSet(copied);
         try {
             return callable.call();
         } finally {
             MtContextThreadLocal.restore(backup);
+            if (releaseMtContextAfterCall) {
+                copied = null;
+            }
         }
-        // FIXME add attribute so as to release copied after call
     }
 
     public Callable<V> getCallable() {
@@ -54,6 +62,19 @@ public final class MtContextCallable<V> implements Callable<V> {
      * @return Wrapped {@link Callable}
      */
     public static <T> MtContextCallable<T> get(Callable<T> callable) {
+        return get(callable, false);
+    }
+
+    /**
+     * Factory method, wrapper input {@link Callable} to {@link MtContextCallable}.
+     * <p/>
+     * This method is idempotent.
+     *
+     * @param callable                  input {@link Callable}
+     * @param releaseMtContextAfterCall release MtContext after run, avoid memory leak even if {@link MtContextRunnable} is referred.
+     * @return Wrapped {@link Callable}
+     */
+    public static <T> MtContextCallable<T> get(Callable<T> callable, boolean releaseMtContextAfterCall) {
         if (null == callable) {
             return null;
         }
@@ -61,19 +82,33 @@ public final class MtContextCallable<V> implements Callable<V> {
         if (callable instanceof MtContextCallable) { // avoid redundant decoration, and ensure idempotency
             return (MtContextCallable<T>) callable;
         }
-        return new MtContextCallable<T>(callable);
+        return new MtContextCallable<T>(callable, releaseMtContextAfterCall);
     }
 
     /**
      * wrapper input {@link Callable} Collection to {@link MtContextCallable} Collection.
+     *
+     * @param tasks task to be wrapped
+     * @return Wrapped {@link Callable}
      */
     public static <T> List<MtContextCallable<T>> gets(Collection<? extends Callable<T>> tasks) {
+        return gets(tasks, false);
+    }
+
+    /**
+     * wrapper input {@link Callable} Collection to {@link MtContextCallable} Collection.
+     *
+     * @param tasks                     task to be wrapped
+     * @param releaseMtContextAfterCall release MtContext after run, avoid memory leak even if {@link MtContextRunnable} is referred.
+     * @return Wrapped {@link Callable}
+     */
+    public static <T> List<MtContextCallable<T>> gets(Collection<? extends Callable<T>> tasks, boolean releaseMtContextAfterCall) {
         if (null == tasks) {
             return null;
         }
         List<MtContextCallable<T>> copy = new ArrayList<MtContextCallable<T>>();
         for (Callable<T> task : tasks) {
-            copy.add(MtContextCallable.get(task));
+            copy.add(MtContextCallable.get(task, releaseMtContextAfterCall));
         }
         return copy;
     }

src/main/java/com/alibaba/mtc/MtContextRunnable.java

@@ -20,12 +20,14 @@ import java.util.Map;
  * @since 0.9.0
  */
 public final class MtContextRunnable implements Runnable {
-    private final Map<MtContextThreadLocal<?>, Object> copied;
+    private volatile Map<MtContextThreadLocal<?>, Object> copied;
     private final Runnable runnable;
+    private final boolean releaseMtContextAfterRun;
 
-    private MtContextRunnable(Runnable runnable) {
-        copied = MtContextThreadLocal.copy();
+    private MtContextRunnable(Runnable runnable, boolean releaseMtContextAfterRun) {
         this.runnable = runnable;
+        this.releaseMtContextAfterRun = releaseMtContextAfterRun;
+        copied = MtContextThreadLocal.copy();
     }
 
     /**
@@ -33,13 +35,18 @@ public final class MtContextRunnable implements Runnable {
      */
     @Override
     public void run() {
+        if(null == copied) {
+            throw new IllegalStateException("MtContext is released!");
+        }
         Map<MtContextThreadLocal<?>, Object> backup = MtContextThreadLocal.backupAndSet(copied);
         try {
             runnable.run();
         } finally {
             MtContextThreadLocal.restore(backup);
+            if (releaseMtContextAfterRun) {
+                copied = null;
+            }
         }
-        // FIXME add attribute so as to release copied after run
     }
 
     public Runnable getRunnable() {
@@ -55,6 +62,19 @@ public final class MtContextRunnable implements Runnable {
      * @return Wrapped {@link Runnable}
      */
     public static MtContextRunnable get(Runnable runnable) {
+        return get(runnable, false);
+    }
+
+    /**
+     * Factory method, wrapper input {@link Runnable} to {@link MtContextRunnable}.
+     * <p/>
+     * This method is idempotent.
+     *
+     * @param runnable                 input {@link Runnable}
+     * @param releaseMtContextAfterRun release MtContext after run, avoid memory leak even if {@link MtContextRunnable} is referred.
+     * @return Wrapped {@link Runnable}
+     */
+    public static MtContextRunnable get(Runnable runnable, boolean releaseMtContextAfterRun) {
         if (null == runnable) {
             return null;
         }
@@ -62,19 +82,33 @@ public final class MtContextRunnable implements Runnable {
         if (runnable instanceof MtContextRunnable) { // avoid redundant decoration, and ensure idempotency
             return (MtContextRunnable) runnable;
         }
-        return new MtContextRunnable(runnable);
+        return new MtContextRunnable(runnable, releaseMtContextAfterRun);
     }
 
     /**
      * wrapper input {@link Runnable} Collection to {@link MtContextRunnable} Collection.
+     *
+     * @param tasks task to be wrapped
+     * @return wrapped tasks
      */
     public static List<MtContextRunnable> gets(Collection<? extends Runnable> tasks) {
+        return gets(tasks, false);
+    }
+
+    /**
+     * wrapper input {@link Runnable} Collection to {@link MtContextRunnable} Collection.
+     *
+     * @param tasks                    task to be wrapped
+     * @param releaseMtContextAfterRun release MtContext after run, avoid memory leak even if {@link MtContextRunnable} is referred.
+     * @return wrapped tasks
+     */
+    public static List<MtContextRunnable> gets(Collection<? extends Runnable> tasks, boolean releaseMtContextAfterRun) {
         if (null == tasks) {
             return null;
         }
         List<MtContextRunnable> copy = new ArrayList<MtContextRunnable>();
         for (Runnable task : tasks) {
-            copy.add(MtContextRunnable.get(task));
+            copy.add(MtContextRunnable.get(task, releaseMtContextAfterRun));
         }
         return copy;
     }