/**
 * Copyright 2015 Palantir Technologies
 *
 * Licensed under the BSD-3 License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://opensource.org/licenses/BSD-3-Clause
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.palantir.common.concurrent;

import java.util.Map;
import java.util.concurrent.ConcurrentMap;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.MapMaker;

/**
 * This class is like {@link InheritableThreadLocal} but it works with tasks spawned from the static
 * methods in {@link PTExecutors}. We don't usually use plain old threads for task management, so
 * {code InheritableThreadLocal} is not that useful. Use this class as a replacement.
 */
public class ExecutorInheritableThreadLocal<T> {

    private final static ThreadLocal<ConcurrentMap<ExecutorInheritableThreadLocal<?>, Object>> mapForThisThread
            = new ThreadLocal<ConcurrentMap<ExecutorInheritableThreadLocal<?>,Object>>() {
        @Override
        protected ConcurrentMap<ExecutorInheritableThreadLocal<?>, Object> initialValue() {
            return makeNewMap();
        }
    };

    private enum NullWrapper {
        INSTANCE
    }

    public void set(T value) {
        mapForThisThread.get().put(this, wrapNull(value));
    }

    public void remove() {
        Map<ExecutorInheritableThreadLocal<?>, Object> map = mapForThisThread.get();
        map.remove(this);
        if (map.isEmpty()) {
            mapForThisThread.remove();
        }
    }

    public T get() {
        if (mapForThisThread.get().containsKey(this)) {
            Object ret = mapForThisThread.get().get(this);
            return unwrapNull(ret);
        } else {
            T ret = initialValue();
            set(ret);
            return ret;
        }
    }

    /**
     * Computes the child's initial value for this thread-local variable as a function of the
     * parent's value. This method is called from within the parent thread during the submit call to
     * the Executor. The value returned will be that of the child's when it is run by the executor.
     * <p>
     * This method merely returns its input argument, and should be overridden if a different
     * behavior is desired.
     *
     * @param parentValue the parent thread's value
     * @return the child thread's initial value
     */
    protected T childValue(T parentValue) {
        return parentValue;
    }

    /**
     * Computes the child's initial value for this thread-local variable as a function of the result
     * of {@link #childValue(Object)} (which is called on the parent thread). This method is run on
     * the child thread.
     * <p>
     * This method merely returns its input argument, and should be overridden if a different
     * behavior is desired.
     *
     * @param childValue the return value from {@link #childValue(Object)}
     * @return the initial value that will be used for this thread
     */
    protected T installOnChildThread(T childValue) {
        return childValue;
    }

    /**
     * Called by the executor service when the submitted task is complete. This method can be used
     * to perform cleanup activities on the child thread.
     * <p>
     * By default this method is a no-op, and should be overridden if different behavior is desired.
     */
    protected void uninstallOnChildThread() {
        /* Do nothing. */
    }

    /**
     * Returns the current thread's initial value for this thread-local variable. This method will
     * be invoked at most once per accessing thread for each thread-local, the first time the thread
     * accesses the variable with the {@link #get} method. The <tt>initialValue</tt> method will not
     * be invoked in a thread if the thread invokes the {@link #set} method prior to the
     * <tt>get</tt> method.
     * <p>
     * This implementation simply returns <tt>null</tt>; if the programmer desires thread-local
     * variables to be initialized to some value other than <tt>null</tt>, <tt>ThreadLocal</tt> must
     * be subclassed, and this method overridden. Typically, an anonymous inner class will be used.
     * Typical implementations of <tt>initialValue</tt> will invoke an appropriate constructor and
     * return the newly constructed object.
     *
     * @return the initial value for this thread-local
     */
    protected T initialValue() {
        return null;
    }

    static ImmutableMap<ExecutorInheritableThreadLocal<?>, Object> getMapForNewThread() {
        ConcurrentMap<ExecutorInheritableThreadLocal<?>, Object> currentMap = mapForThisThread.get();
        if (currentMap.isEmpty()) {
            mapForThisThread.remove();
            return ImmutableMap.of();
        }
        ImmutableMap.Builder<ExecutorInheritableThreadLocal<?>, Object> ret = ImmutableMap.builder();
        for (ExecutorInheritableThreadLocal<?> e : currentMap.keySet()) {
            @SuppressWarnings("unchecked")
            ExecutorInheritableThreadLocal<Object> eitl = (ExecutorInheritableThreadLocal<Object>) e;
            ret.put(eitl, wrapNull(eitl.callChildValue(eitl.get())));
        }

        return ret.build();
    }

    /**
     * @return the old map installed on that thread
     */
    static ConcurrentMap<ExecutorInheritableThreadLocal<?>, Object> installMapOnThread(ImmutableMap<ExecutorInheritableThreadLocal<?>, Object> map) {
        ConcurrentMap<ExecutorInheritableThreadLocal<?>, Object> oldMap = mapForThisThread.get();
        if (map.isEmpty()) {
            mapForThisThread.remove();
        } else {
            ConcurrentMap<ExecutorInheritableThreadLocal<?>, Object> newMap = makeNewMap();
            newMap.putAll(map);

            // Install the map in case callInstallOnChildThread makes use
            // of existing thread locals (UserSessionClientInfo does this).
            mapForThisThread.set(newMap);

            for (ExecutorInheritableThreadLocal<?> e : map.keySet()) {
                @SuppressWarnings("unchecked")
                ExecutorInheritableThreadLocal<Object> eitl = (ExecutorInheritableThreadLocal<Object>) e;
                eitl.set(eitl.callInstallOnChildThread(eitl.get()));
            }
        }
        return oldMap;
    }

    private static ConcurrentMap<ExecutorInheritableThreadLocal<?>, Object> makeNewMap() {
        return new MapMaker().weakKeys().concurrencyLevel(1).makeMap();
    }

    private static Object wrapNull(Object value) {
        return value == null ? NullWrapper.INSTANCE : value;
    }

    @SuppressWarnings("unchecked")
    private T unwrapNull(Object ret) {
        if (ret == NullWrapper.INSTANCE) {
            return null;
        } else {
            return (T) ret;
        }
    }

    static void uninstallMapOnThread(ConcurrentMap<ExecutorInheritableThreadLocal<?>, Object> oldMap) {
        try {
            for (ExecutorInheritableThreadLocal<?> eitl : mapForThisThread.get().keySet()) {
                eitl.uninstallOnChildThread();
            }
        } finally {
            if (oldMap.isEmpty()) {
                mapForThisThread.remove();
            } else {
                mapForThisThread.set(oldMap);
            }
        }
    }

    @Override
    public final boolean equals(Object obj) {
        return super.equals(obj);
    }

    @Override
    public final int hashCode() {
        return super.hashCode();
    }

    @SuppressWarnings("unchecked")
    private T callChildValue(Object o) {
        return childValue((T) o);
    }

    @SuppressWarnings("unchecked")
    private T callInstallOnChildThread(Object o) {
        return installOnChildThread((T) o);
    }
}

@@ -3664,24 +3664,321 @@
 is desired.%0A
+     * %3Cp%3E%0A     * This will be run from a finally block, so it should not throw.%0A     * %3Cp%3E%0A     * NOTE: This code isn't guaranteed to finish by the time future.get() returns in the calling thread.%0A     * The completed Future is marked %3Ccode%3EisDone%3C/code%3E and then this is run in a finally block.%0A
      */%0A    

add comment