add beforeExecute/afterExecute callback

alibaba · Jun 17, 2015 · 3731a9c · 3731a9c
1 parent 29e3b14
commit 3731a9c
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 4 deletions.
diff --git a/src/main/java/com/alibaba/mtc/MtContextRunnable.java b/src/main/java/com/alibaba/mtc/MtContextRunnable.java
@@ -137,7 +137,7 @@ public static List<MtContextRunnable> gets(Collection<? extends Runnable> tasks,
      *
      * @param tasks                    task to be wrapped
      * @param releaseMtContextAfterRun release MtContext after run, avoid memory leak even if {@link MtContextRunnable} is referred.
-     * @param idempotent                is idempotent or not. {@code true} will cover up bugs! <b>DO NOT</b> set, only when you know why.
+     * @param idempotent               is idempotent or not. {@code true} will cover up bugs! <b>DO NOT</b> set, only when you know why.
      * @return wrapped tasks
      */
     public static List<MtContextRunnable> gets(Collection<? extends Runnable> tasks, boolean releaseMtContextAfterRun, boolean idempotent) {

diff --git a/src/main/java/com/alibaba/mtc/MtContextThreadLocal.java b/src/main/java/com/alibaba/mtc/MtContextThreadLocal.java
@@ -32,6 +32,32 @@ protected T copy(T parentValue) {
         return parentValue;
     }
 
+    /**
+     * Callback method before task object({@link MtContextRunnable}/{@link MtContextCallable}) execute.
+     * <p/>
+     * Default behavior is do nothing, and should be overridden
+     * if a different behavior is desired.
+     * <p/>
+     * Do not throw any exception, just ignored.
+     *
+     * @since 1.2.0
+     */
+    protected void beforeExecute() {
+    }
+
+    /**
+     * Callback method after task object({@link MtContextRunnable}/{@link MtContextCallable}) execute.
+     * <p/>
+     * Default behavior is do nothing, and should be overridden
+     * if a different behavior is desired.
+     * <p/>
+     * Do not throw any exception, just ignored.
+     *
+     * @since 1.2.0
+     */
+    protected void afterExecute() {
+    }
+
     @Override
     public final T get() {
         T value = super.get();
@@ -105,7 +131,7 @@ static Map<MtContextThreadLocal<?>, Object> backupAndSet(Map<MtContextThreadLoca
                 threadLocal.superRemove();
             }
         }
-        setMtContexts(copied);
+        setMtContexts(copied, true);
         return backup;
     }
 
@@ -120,14 +146,29 @@ static void restore(Map<MtContextThreadLocal<?>, Object> backup) {
                 threadLocal.superRemove();
             }
         }
-        setMtContexts(backup);
+        setMtContexts(backup, false);
     }
 
-    static void setMtContexts(Map<MtContextThreadLocal<?>, Object> set) {
+    static void setMtContexts(Map<MtContextThreadLocal<?>, Object> set, boolean isStore) {
         for (Map.Entry<MtContextThreadLocal<?>, Object> entry : set.entrySet()) {
             @SuppressWarnings("unchecked")
             MtContextThreadLocal<Object> threadLocal = (MtContextThreadLocal<Object>) entry.getKey();
+
+            if (!isStore) {
+                try {
+                    threadLocal.afterExecute();
+                } catch (Throwable t) {
+                    t.printStackTrace();
+                }
+            }
             threadLocal.set(entry.getValue());
+            if (isStore) {
+                try {
+                    threadLocal.beforeExecute();
+                } catch (Throwable t) {
+                    t.printStackTrace();
+                }
+            }
         }
     }
 }