TransmittableThreadLocal.java

package com.alibaba.ttl3;

import com.alibaba.ttl3.executor.TtlExecutors;
import com.alibaba.ttl3.transmitter.Transmittee;
import com.alibaba.ttl3.transmitter.TransmitteeRegistry;
import edu.umd.cs.findbugs.annotations.NonNull;

import javax.annotation.ParametersAreNonnullByDefault;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.WeakHashMap;
import java.util.function.Supplier;
import java.util.function.UnaryOperator;

import static com.alibaba.ttl3.internal.util.Utils.newHashMap;

/**
 * {@link TransmittableThreadLocal}({@code TTL}) can transmit the value from the thread of submitting task
 * to the thread of executing task even using thread pooling components.
 * <p>
 * <b>Note</b>:<br>
 * {@link TransmittableThreadLocal} extends {@link InheritableThreadLocal},
 * so {@link TransmittableThreadLocal} first is a {@link InheritableThreadLocal}.<br>
 * If the <b>inheritable</b> ability from {@link InheritableThreadLocal} has <b>potential leaking problem</b>,
 * you can disable the <b>inheritable</b> ability:
 * <p>
 * ❶ For thread pooling components({@link java.util.concurrent.ThreadPoolExecutor},
 * {@link java.util.concurrent.ForkJoinPool}), Inheritable feature <b>should never</b> happen,
 * since threads in thread pooling components is pre-created and pooled, these threads is <b>neutral</b> to biz logic/data.
 * <br>
 * Disable inheritable for thread pooling components by wrapping thread factories using methods
 * {@link TtlExecutors#getDisableInheritableThreadFactory(java.util.concurrent.ThreadFactory) getDisableInheritableThreadFactory} /
 * {@link TtlExecutors#getDefaultDisableInheritableForkJoinWorkerThreadFactory() getDefaultDisableInheritableForkJoinWorkerThreadFactory}.
 * <br>
 * Or you can turn on "disable inheritable for thread pool" by {@code TTL Java Agent}
 * to wrap thread factories for thread pooling components automatically and transparently.
 * <p>
 * ❷ In other cases, disable inheritable by overriding method {@link #childValue(Object)}.
 * <br>
 * Whether the value should be inheritable or not can be controlled by the data owner,
 * disable it <b>carefully</b> when data owner have a clear idea.
 * <pre>{@code
 * TransmittableThreadLocal<String> transmittableThreadLocal = new TransmittableThreadLocal<>() {
 *     protected String childValue(String parentValue) {
 *         return initialValue();
 *     }
 * }}</pre>
 * <p>
 * More discussion about "disable the <b>inheritable</b> ability"
 * see <a href="https://github.com/alibaba/transmittable-thread-local/issues/100">
 * issue #100: disable Inheritable when it's not necessary and buggy</a>.
 *
 * @author Jerry Lee (oldratlee at gmail dot com)
 * @author Yang Fang (snoop dot fy at gmail dot com)
 * @see <a href="https://github.com/alibaba/transmittable-thread-local">user guide docs and code repo of TransmittableThreadLocal(TTL)</a>
 * @see TtlRunnable
 * @see TtlCallable
 * @see TtlExecutors
 * @see TtlExecutors#getTtlExecutor(java.util.concurrent.Executor)
 * @see TtlExecutors#getTtlExecutorService(java.util.concurrent.ExecutorService)
 * @see TtlExecutors#getTtlScheduledExecutorService(java.util.concurrent.ScheduledExecutorService)
 * @see TtlExecutors#getDefaultDisableInheritableThreadFactory()
 * @see TtlExecutors#getDisableInheritableThreadFactory(java.util.concurrent.ThreadFactory)
 * @see TtlExecutors#getDefaultDisableInheritableForkJoinWorkerThreadFactory()
 * @see TtlExecutors#getDisableInheritableForkJoinWorkerThreadFactory(java.util.concurrent.ForkJoinPool.ForkJoinWorkerThreadFactory)
 */
public class TransmittableThreadLocal<T> extends InheritableThreadLocal<T> {
    private final boolean disableIgnoreNullValueSemantics;

    /**
     * Default constructor. Create a {@link TransmittableThreadLocal} instance with "Ignore-Null-Value Semantics".
     * <p>
     * About "Ignore-Null-Value Semantics":
     *
     * <ol>
     *     <li>If value is {@code null}(check by {@link #get()} method), do NOT transmit this {@code ThreadLocal}.</li>
     *     <li>If set {@code null} value, also remove value(invoke {@link #remove()} method).</li>
     * </ol>
     * <p>
     * This is a pragmatic design decision:
     * <ol>
     * <li>use explicit value type rather than {@code null} value to express biz intent.</li>
     * <li>safer and more robust code(avoid {@code NPE} risk).</li>
     * </ol>
     * <p>
     * So it's strongly not recommended to use {@code null} value.
     * <p>
     * But the behavior of "Ignore-Null-Value Semantics" is NOT compatible with
     * {@link ThreadLocal} and {@link InheritableThreadLocal},
     * you can disable this behavior/semantics via using constructor {@link #TransmittableThreadLocal(boolean)}
     * and setting parameter {@code disableIgnoreNullValueSemantics} to {@code true}.
     * <p>
     * More discussion about "Ignore-Null-Value Semantics" see
     * <a href="https://github.com/alibaba/transmittable-thread-local/issues/157">Issue #157</a>.
     *
     * @see #TransmittableThreadLocal(boolean)
     */
    public TransmittableThreadLocal() {
        this(false);
    }

    /**
     * Constructor, create a {@link TransmittableThreadLocal} instance
     * with parameter {@code disableIgnoreNullValueSemantics} to control "Ignore-Null-Value Semantics".
     *
     * @param disableIgnoreNullValueSemantics disable "Ignore-Null-Value Semantics"
     * @see #TransmittableThreadLocal()
     */
    public TransmittableThreadLocal(boolean disableIgnoreNullValueSemantics) {
        this.disableIgnoreNullValueSemantics = disableIgnoreNullValueSemantics;
    }

    /**
     * Creates a transmittable thread local variable.
     * The initial value({@link #initialValue()}) of the variable is
     * determined by invoking the {@link #get()} method on the {@code Supplier}.
     *
     * @param <S>      the type of the thread local's value
     * @param supplier the supplier to be used to determine the initial value
     * @return a new transmittable thread local variable
     * @throws NullPointerException if the specified supplier is null
     * @see #withInitialAndGenerator(Supplier, UnaryOperator)
     */
    @NonNull
    @SuppressWarnings("ConstantConditions")
    public static <S> TransmittableThreadLocal<S> withInitial(@NonNull Supplier<? extends S> supplier) {
        if (supplier == null) throw new NullPointerException("supplier is null");

        return new SuppliedTransmittableThreadLocal<>(supplier, null, null);
    }

    /**
     * Creates a transmittable thread local variable.
     * The initial value({@link #initialValue()}) of the variable is
     * determined by invoking the {@link #get()} method on the {@code Supplier};
     * and the child value({@link #childValue(Object)}) and the transmittee value({@link #transmitteeValue(Object)}) of the variable is
     * determined by invoking the {@link UnaryOperator#apply(Object)} method on the {@code UnaryOperator}.
     *
     * @param <S>                                       the type of the thread local's value
     * @param supplier                                  the supplier to be used to determine the initial value
     * @param generatorForChildValueAndTransmitteeValue the value generator to be used to determine the child value and the transmittee value
     * @return a new transmittable thread local variable
     * @throws NullPointerException if the specified supplier or value generator is null
     * @see #withInitial(Supplier)
     */
    @NonNull
    @ParametersAreNonnullByDefault
    @SuppressWarnings("ConstantConditions")
    public static <S> TransmittableThreadLocal<S> withInitialAndGenerator(Supplier<? extends S> supplier, UnaryOperator<S> generatorForChildValueAndTransmitteeValue) {
        if (supplier == null) throw new NullPointerException("supplier is null");
        if (generatorForChildValueAndTransmitteeValue == null) throw new NullPointerException("value generator is null");

        return new SuppliedTransmittableThreadLocal<>(supplier, generatorForChildValueAndTransmitteeValue, generatorForChildValueAndTransmitteeValue);
    }

    /**
     * Creates a transmittable thread local variable.
     * The initial value({@link #initialValue()}) of the variable is
     * determined by invoking the {@link #get()} method on the {@code Supplier};
     * and the child value({@link #childValue(Object)})}) and the transmittee value({@link #transmitteeValue(Object)}) of the variable is
     * determined by invoking the {@link UnaryOperator#apply(Object)} method on the {@code UnaryOperator}.
     * <p>
     * <B><I>NOTE:</I></B><br>
     * Recommend use {@link #withInitialAndGenerator(Supplier, UnaryOperator)} instead of this method.
     * In most cases, the logic of determining the child value({@link #childValue(Object)})
     * and the transmittee value({@link #transmitteeValue(Object)}) should be the same.
     *
     * @param <S>                          the type of the thread local's value
     * @param supplier                     the supplier to be used to determine the initial value
     * @param generatorForChildValue       the value generator to be used to determine the child value
     * @param generatorForTransmitteeValue the value generator to be used to determine the transmittee value
     * @return a new transmittable thread local variable
     * @throws NullPointerException if the specified supplier or value generator is null
     * @see #withInitial(Supplier)
     * @see #withInitialAndGenerator(Supplier, UnaryOperator)
     */
    @NonNull
    @ParametersAreNonnullByDefault
    @SuppressWarnings("ConstantConditions")
    public static <S> TransmittableThreadLocal<S> withInitialAndGenerator(Supplier<? extends S> supplier, UnaryOperator<S> generatorForChildValue, UnaryOperator<S> generatorForTransmitteeValue) {
        if (supplier == null) throw new NullPointerException("supplier is null");
        if (generatorForChildValue == null) throw new NullPointerException("value generator for child value is null");
        if (generatorForTransmitteeValue == null) throw new NullPointerException("value generator for transmittee value is null");

        return new SuppliedTransmittableThreadLocal<>(supplier, generatorForChildValue, generatorForTransmitteeValue);
    }

    /**
     * An extension of ThreadLocal that obtains its initial value from the specified {@code Supplier}
     * and obtains its child value and transmittee value from the specified generator.
     */
    private static final class SuppliedTransmittableThreadLocal<T> extends TransmittableThreadLocal<T> {
        private final Supplier<? extends T> supplier;
        private final UnaryOperator<T> generatorForChildValue;
        private final UnaryOperator<T> generatorForTransmitteeValue;

        SuppliedTransmittableThreadLocal(Supplier<? extends T> supplier, UnaryOperator<T> generatorForChildValue, UnaryOperator<T> generatorForTransmitteeValue) {
            if (supplier == null) throw new NullPointerException("supplier is null");
            this.supplier = supplier;
            this.generatorForChildValue = generatorForChildValue;
            this.generatorForTransmitteeValue = generatorForTransmitteeValue;
        }

        @Override
        protected T initialValue() {
            return supplier.get();
        }

        @Override
        protected T childValue(T parentValue) {
            if (generatorForChildValue != null) return generatorForChildValue.apply(parentValue);
            else return super.childValue(parentValue);
        }

        @Override
        public T transmitteeValue(T parentValue) {
            if (generatorForTransmitteeValue != null) return generatorForTransmitteeValue.apply(parentValue);
            else return super.transmitteeValue(parentValue);
        }
    }

    /**
     * Computes the child's initial value for this transmittable thread-local
     * variable as a function of the parent's value at the time the child
     * thread is created. This method is called from within the parent
     * thread before the child is started.
     * <p>
     * <b>Note</b>:<br>
     * This method is overridden, and merely call {@link #transmitteeValue(Object)}.
     * In most cases, the logic of determining the child value({@link #childValue(Object)})
     * and the transmittee value({@link #transmitteeValue(Object)}) should be the same,
     * so it's NOT recommended to override this method in subclass.
     *
     * @param parentValue the parent thread's value
     * @return the child thread's initial value
     */
    @Override
    protected T childValue(T parentValue) {
        return transmitteeValue(parentValue);
    }

    /**
     * Computes the value for this transmittable thread-local variable
     * as a function of the source thread's value at the time the task
     * Object is created.
     * <p>
     * This method is called from {@link TtlRunnable} or
     * {@link TtlCallable} when it create, before the task is started.
     * <p>
     * <b>Note</b>:<br>
     * This method merely returns reference of its source thread value(the shadow copy),
     * and should be overridden if a different behavior is desired.
     * It's recommended to override this method in subclass.
     */
    protected T transmitteeValue(T parentValue) {
        return parentValue;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public final T get() {
        T value = super.get();
        if (disableIgnoreNullValueSemantics || value != null) addThisToHolder();
        return value;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public final void set(T value) {
        if (!disableIgnoreNullValueSemantics && value == null) {
            // may set null to remove value
            remove();
        } else {
            super.set(value);
            addThisToHolder();
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public final void remove() {
        removeThisFromHolder();
        super.remove();
    }

    private void superRemove() {
        super.remove();
    }

    private T getTransmitteeValue() {
        return transmitteeValue(get());
    }

    // Note about the holder:
    // 1. holder self is a InheritableThreadLocal(a *ThreadLocal*).
    // 2. The type of value in the holder is WeakHashMap<TransmittableThreadLocal<Object>, ?>.
    //    2.1 but the WeakHashMap is used as a *Set*:
    //        the value of WeakHashMap is *always* null, and never used.
    //    2.2 WeakHashMap support *null* value.
    private static final InheritableThreadLocal<WeakHashMap<TransmittableThreadLocal<Object>, ?>> holder =
            new InheritableThreadLocal<WeakHashMap<TransmittableThreadLocal<Object>, ?>>() {
                @Override
                protected WeakHashMap<TransmittableThreadLocal<Object>, ?> initialValue() {
                    return new WeakHashMap<>();
                }

                @Override
                protected WeakHashMap<TransmittableThreadLocal<Object>, ?> childValue(WeakHashMap<TransmittableThreadLocal<Object>, ?> parentValue) {
                    return new WeakHashMap<>(parentValue);
                }
            };

    @SuppressWarnings("unchecked")
    private void addThisToHolder() {
        if (!holder.get().containsKey(this)) {
            holder.get().put((TransmittableThreadLocal<Object>) this, null); // WeakHashMap supports null value.
        }
    }

    private void removeThisFromHolder() {
        holder.get().remove(this);
    }


    private static class TtlTransmittee implements Transmittee<HashMap<TransmittableThreadLocal<Object>, Object>, HashMap<TransmittableThreadLocal<Object>, Object>> {
        @NonNull
        @Override
        public HashMap<TransmittableThreadLocal<Object>, Object> capture() {
            final HashMap<TransmittableThreadLocal<Object>, Object> ttl2Value = newHashMap(holder.get().size());
            for (TransmittableThreadLocal<Object> threadLocal : holder.get().keySet()) {
                ttl2Value.put(threadLocal, threadLocal.getTransmitteeValue());
            }
            return ttl2Value;
        }

        @NonNull
        @Override
        public HashMap<TransmittableThreadLocal<Object>, Object> replay(@NonNull HashMap<TransmittableThreadLocal<Object>, Object> captured) {
            final HashMap<TransmittableThreadLocal<Object>, Object> backup = newHashMap(holder.get().size());

            for (final Iterator<TransmittableThreadLocal<Object>> iterator = holder.get().keySet().iterator(); iterator.hasNext(); ) {
                TransmittableThreadLocal<Object> threadLocal = iterator.next();

                // backup
                backup.put(threadLocal, threadLocal.get());

                // clear the TTL values that is not in captured
                // avoid the extra TTL values after replay when run task
                if (!captured.containsKey(threadLocal)) {
                    iterator.remove();
                    threadLocal.superRemove();
                }
            }

            // set TTL values to captured
            setTtlValuesTo(captured);

            return backup;
        }

        @NonNull
        @Override
        public HashMap<TransmittableThreadLocal<Object>, Object> clear() {
            return replay(newHashMap(0));
        }

        @Override
        public void restore(@NonNull HashMap<TransmittableThreadLocal<Object>, Object> backup) {
            for (final Iterator<TransmittableThreadLocal<Object>> iterator = holder.get().keySet().iterator(); iterator.hasNext(); ) {
                TransmittableThreadLocal<Object> threadLocal = iterator.next();

                // clear the TTL values that is not in backup
                // avoid the extra TTL values after restore
                if (!backup.containsKey(threadLocal)) {
                    iterator.remove();
                    threadLocal.superRemove();
                }
            }

            // restore TTL values
            setTtlValuesTo(backup);
        }

        private static void setTtlValuesTo(@NonNull HashMap<TransmittableThreadLocal<Object>, Object> ttlValues) {
            for (Map.Entry<TransmittableThreadLocal<Object>, Object> entry : ttlValues.entrySet()) {
                TransmittableThreadLocal<Object> threadLocal = entry.getKey();
                threadLocal.set(entry.getValue());
            }
        }
    }

    private static final TtlTransmittee ttlTransmittee = new TtlTransmittee();

    static {
        TransmitteeRegistry.registerTransmittee(ttlTransmittee);
    }
}