add beforeExecute/afterExecute callback

alibaba · Jun 16, 2015 · b963d6c · b963d6c
1 parent 29e3b14
commit b963d6c
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/src/main/java/com/alibaba/mtc/MtContextThreadLocal.java b/src/main/java/com/alibaba/mtc/MtContextThreadLocal.java
@@ -32,6 +32,12 @@ protected T copy(T parentValue) {
         return parentValue;
     }
 
+    protected void beforeExecute() {
+    }
+
+    protected void afterExecute() {
+    }
+
     @Override
     public final T get() {
         T value = super.get();
@@ -105,7 +111,7 @@ static Map<MtContextThreadLocal<?>, Object> backupAndSet(Map<MtContextThreadLoca
                 threadLocal.superRemove();
             }
         }
-        setMtContexts(copied);
+        setMtContexts(copied, true);
         return backup;
     }
 
@@ -120,14 +126,19 @@ 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();
             threadLocal.set(entry.getValue());
+            if (isStore) {
+                threadLocal.beforeExecute();
+            } else {
+                threadLocal.afterExecute();
+            }
         }
     }
 }