package com.alibaba.mtc;

import java.lang.ref.WeakReference;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

/**
 * {@link MtContextThreadLocal} can transmit context from the thread of submitting task to the thread of executing task.
 * <p/>
 * Note: this class extends {@link java.lang.InheritableThreadLocal},
 * so {@link com.alibaba.mtc.MtContextThreadLocal} first is a {@link java.lang.InheritableThreadLocal}.
 *
 * @author ding.lid
 * @see MtContextRunnable
 * @see MtContextCallable
 * @since 0.10.0
 */
public class MtContextThreadLocal<T> extends InheritableThreadLocal<T> {
    /**
     * Computes the context value for this multi-thread context variable
     * as a function of the source thread's value at the time the task
     * Object is created.  This method is called from {@link com.alibaba.mtc.MtContextRunnable} or
     * {@link com.alibaba.mtc.MtContextCallable} when it create, before the task is started.
     * <p/>
     * This method merely returns reference of its source thread value, and should be overridden
     * if a different behavior is desired.
     *
     * @since 1.0.0
     */
    protected T copyValue(T parentValue) {
        return parentValue;
    }

    /**
     * Override this method to have an initial value other than <tt>null</tt>.
     */
    @Override
    protected T initialValue() {
        return super.initialValue();
    }

    @Override
    protected T childValue(T parentValue) {
        return super.childValue(parentValue);
    }

    @Override
    public final T get() {
        T value = super.get();
        if (null != value) {
            addMtContextThreadLocal();
        }
        return value;
    }

    @Override
    public final void set(T value) {
        super.set(value);
        if (null == value) { // may set null to remove value
            removeMtContextThreadLocal();
        } else {
            addMtContextThreadLocal();
        }
    }

    @Override
    public final void remove() {
        removeMtContextThreadLocal();
        super.remove();
    }

    T copyMtContextValue() {
        return copyValue(get());
    }

    static class Reference<T> extends WeakReference<MtContextThreadLocal<T>> {
        public Reference(MtContextThreadLocal<T> referent) {
            super(referent);
        }
    }

    static ThreadLocal<Set<Reference<?>>> holder =
            new ThreadLocal<Set<Reference<?>>>() {
                @Override
                protected Set<Reference<?>> initialValue() {
                    return new HashSet<Reference<?>>();
                }
            };

    void addMtContextThreadLocal() {
        for (Iterator<Reference<?>> iterator = holder.get().iterator(); iterator.hasNext(); ) {
            Reference<?> reference = iterator.next();
            MtContextThreadLocal<?> threadLocal = reference.get();
            if (null == threadLocal) {
                iterator.remove();
                continue;
            }
            if (threadLocal == this) {
                return;
            }
        }
        holder.get().add(new Reference<T>(this));
    }

    void removeMtContextThreadLocal() {
        for (Iterator<Reference<?>> iterator = holder.get().iterator(); iterator.hasNext(); ) {
            Reference<?> reference = iterator.next();
            MtContextThreadLocal<?> threadLocal = reference.get();
            if (null == threadLocal) {
                iterator.remove();
                continue;
            }
            if (threadLocal == this) {
                iterator.remove();
            }
        }
    }

    static Map<MtContextThreadLocal<?>, Object> copy() {
        Map<MtContextThreadLocal<?>, Object> copy = new HashMap<MtContextThreadLocal<?>, Object>();
        for (Iterator<Reference<?>> iterator = holder.get().iterator(); iterator.hasNext(); ) {
            Reference<?> reference = iterator.next();
            MtContextThreadLocal<?> threadLocal = reference.get();
            if (threadLocal == null) {
                iterator.remove();
            } else {
                copy.put(threadLocal, threadLocal.copyMtContextValue());
            }
        }
        return copy;
    }

    static Map<MtContextThreadLocal<?>, Object> backupAndSet(Map<MtContextThreadLocal<?>, Object> set) {
        // backup MtContext
        Map<MtContextThreadLocal<?>, Object> backup = new HashMap<MtContextThreadLocal<?>, Object>();
        for (Map.Entry<MtContextThreadLocal<?>, Object> entry : set.entrySet()) {
            @SuppressWarnings("unchecked")
            MtContextThreadLocal<Object> threadLocal = (MtContextThreadLocal<Object>) entry.getKey();
            backup.put(threadLocal, threadLocal.get());
            threadLocal.set(entry.getValue());
        }
        return backup;
    }

    static void restore(Map<MtContextThreadLocal<?>, Object> backup) {
        // restore MtContext
        for (Map.Entry<MtContextThreadLocal<?>, Object> entry : backup.entrySet()) {
            @SuppressWarnings("unchecked")
            MtContextThreadLocal<Object> threadLocal = (MtContextThreadLocal<Object>) entry.getKey();
            threadLocal.set(entry.getValue());
        }
    }
}

@@ -85,34 +85,8 @@
 ap;%0A
-import java.util.HashSet;%0A
 impo
@@ -147,19 +147,27 @@
 va.util.
-Set
+WeakHashMap
 ;%0A%0A/**%0A 
@@ -2175,222 +2175,50 @@
 tic 
-class Reference%3CT%3E extends WeakReference%3CMtContextThreadLocal%3CT%3E%3E %7B%0A        public Reference(MtContextThreadLocal%3CT%3E referent) %7B%0A            super(referent);%0A        %7D%0A    %7D%0A%0A    static ThreadLocal%3CSet%3CReference%3C?%3E
+ThreadLocal%3CMap%3CMtContextThreadLocal%3C?%3E, ?
 %3E%3E h
@@ -2253,32 +2253,46 @@
 adLocal%3C
-Set%3CReference%3C?%3E
+Map%3CMtContextThreadLocal%3C?%3E, ?
 %3E%3E() %7B%0A 
@@ -2342,32 +2342,46 @@
 otected 
-Set%3CReference%3C?%3E
+Map%3CMtContextThreadLocal%3C?%3E, ?
 %3E initia
@@ -2426,28 +2426,51 @@
 new 
-HashSet%3CReference%3C?%3E
+WeakHashMap%3CMtContextThreadLocal%3C?%3E, Object
 %3E();
@@ -2553,943 +2553,179 @@
     
-for (Iterator%3CReference%3C?%3E%3E iterator = holder.get().iterator(); iterator.hasNext(); ) %7B%0A            Reference%3C?%3E reference = iterator.next();%0A            MtContextThreadLocal%3C?%3E threadLocal = reference.get();%0A            if (null == threadLocal) %7B%0A                iterator.remove();%0A                continue;%0A            %7D%0A            if (threadLocal == this) %7B%0A                return;%0A            %7D%0A        %7D%0A        holder.get().add(new Reference%3CT%3E(this));%0A    %7D%0A%0A    void removeMtContextThreadLocal() %7B%0A        for (Iterator%3CReference%3C?%3E%3E iterator = holder.get().iterator(); iterator.hasNext(); ) %7B%0A            Reference%3C?%3E reference = iterator.next();%0A            MtContextThreadLocal%3C?%3E threadLocal = reference.get();%0A            if (null == threadLocal) %7B%0A                iterator.remove();%0A                continue;%0A            %7D%0A            if (threadLocal == this) %7B%0A                iterator.remove();%0A            %7D%0A        %7D
+if(!holder.get().containsKey(this)) %7B%0A            holder.get().put(this, null);%0A        %7D%0A    %7D%0A%0A    void removeMtContextThreadLocal() %7B%0A        holder.get().remove(this);
 %0A   
@@ -2907,25 +2907,36 @@
 terator%3C
-Reference
+MtContextThreadLocal
 %3C?%3E%3E ite
@@ -2956,16 +2956,25 @@
 r.get().
+keySet().
 iterator
@@ -3017,216 +3017,63 @@
     
-Reference%3C?%3E reference = iterator.next();%0A            MtContextThreadLocal%3C?%3E threadLocal = reference.get();%0A            if (threadLocal == null) %7B%0A                iterator.remove();%0A            %7D else %7B%0A    
+MtContextThreadLocal%3C?%3E threadLocal = iterator.next();%0A
     
@@ -3137,30 +3137,16 @@
 lue());%0A
-            %7D%0A
         

change MtContextThreadLocal.holder to WeakHashMap #28