/*
 * Written by Dawid Kurzyniec and released to the public domain, as explained
 * at http://creativecommons.org/licenses/publicdomain
 */

package edu.emory.mathcs.util.concurrent;

import java.security.*;
import java.util.*;
import edu.emory.mathcs.backport.java.util.concurrent.*;

/**
 * Represents an immutable snapshot of a thread state. The state consists of
 * {@link DelegatableThreadLocal delegatable} thread locals,
 * context class loader, and a priority. After the snapshot has been taken,
 * it is later possible to execute tasks within that previously stored context
 * using one of {@link #perform perform} methods. The thread executing the task
 * need not be the same as the one for which the snapshot was originally taken;
 * that is, the state can be recovered (for the duration of the task) in any thread
 * invoking "perform".
 * <p>
 * This class is particularly useful in developing thread pools and
 * asynchronous invocation frameworks, where it is common that tasks are
 * executed in generic asynchronous threads on behalf of invokers initiating
 * them from other threads. This class allows such worker threads to
 * temporarily inherit initiator thread's properties, such as delegatable
 * thread locals or a context class loader.
 * <p>
 * The snapshot represented by this class is immutable; that is, it is not
 * altered during the execution of "perform" methods. However, if desired,
 * it is possible to implement state persistence via chained propagation:
 * that is, a replacement snapshot should be re-taken by the task at
 * the end of the execution.
 * <p>
 * Note that non-delegatable thread locals (e.g. normal ThreadLocals and
 * InheritableThreadLocals) cannot be detached from their threads and thus
 * are not propagated through the thread context.
 * For instance, non-delegatable thread locals set during a
 * "perform"
 * call will remain set in the current thread after the call completes.
 */
public class ThreadContext {

    final ClassLoader ccl;
    final int priority;
    final Map delegThreadLocals;
    int hash = 0;

    private ThreadContext(ClassLoader ccl, int priority,
                          Map delegThreadLocals) {
        this.ccl = ccl;
        this.priority = priority;
        this.delegThreadLocals = delegThreadLocals;
    }

    /**
     * Takes the snapshot of the current thread's state. The state
     * consists of context class loader, delegatable thread locals, and a
     * priority. The snapshot is immutable.
     */
    public static ThreadContext getContext() {
        final Thread current = Thread.currentThread();
        ClassLoader ccl = (ClassLoader)AccessController.doPrivileged(
            new PrivilegedAction() {
                public Object run() {
                    return current.getContextClassLoader();
                }
        });
        int priority = current.getPriority();
        Map delegThreadLocals = DelegatableThreadLocal.takeSnapshot();
        return new ThreadContext(ccl, priority, delegThreadLocals);
    }

    public static ThreadContext create(ClassLoader ccl, int priority) {
        return new ThreadContext(ccl, priority, Collections.EMPTY_MAP);
    }

    public static ThreadContext create(ClassLoader ccl) {
        return create(ccl, Thread.NORM_PRIORITY);
    }

    /**
     * Temporarily recovers current thread's state from this snapshot and then
     * executes the specified task. In result, the task will observe
     * access control context, delegatable thread locals, and the context
     * class loader, as they existed when the snapshot was taken
     * (perhaps in a different thread). Upon completion, the original thread
     * state is restored.
     *
     * @param task the task to execute
     */
    public void perform(final Runnable task) {
        perform(new PrivilegedAction() {
            public Object run() {
                task.run();
                return null;
            }
        });
    }

    /**
     * Temporarily recovers current thread's state from this snapshot and then
     * executes the specified call. In result, the call will observe
     * access control context, delegatable thread locals, and the context
     * class loader, as they existed when the snapshot was taken
     * (perhaps in a different thread). Upon completion, the original thread
     * state is restored.
     *
     * @param call the call to execute
     * @return the result of the call
     * @throws Exception thrown from the call
     */
    public Object perform(final Callable call) throws Exception {
        try {
            return perform(new PrivilegedExceptionAction() {
                public Object run() throws Exception {
                    return call.call();
                }
            });
        }
        catch (PrivilegedActionException pae) {
            throw pae.getException();
        }
    }

    /**
     * Temporarily recovers current thread's state from this snapshot and then
     * performs the specified action. In result, the action will observe
     * access control context, delegatable thread locals, and the context
     * class loader, as they existed when the snapshot was taken
     * (perhaps in a different thread). Upon completion, the original thread
     * state is restored.
     *
     * @param action the action to perform
     * @return the result of the action
     */
    public Object perform(final PrivilegedAction action) {
        final Thread current = Thread.currentThread();
        final AccessControlContext acc = AccessController.getContext();
        // make sure that we can modify this worker thread even if we're
        // called from un-privileged code
        return AccessController.doPrivileged(new PrivilegedAction() { public Object run() {
            ClassLoader savedCcl = current.getContextClassLoader();
            int savedPriority = current.getPriority();
            Map savedDtl;
            try {
                if (ccl != savedCcl) current.setContextClassLoader(ccl);
                // do not boost priority if the request comes from
                // high-priority thread: due to high-cost of setPriority
                // method, it would usually result in a slow-down. However,
                // we do prohibit low-priority threads to schedule
                // higher-priority tasks
                if (priority < savedPriority) current.setPriority(priority);
                savedDtl = DelegatableThreadLocal.delegate(delegThreadLocals);
                try {
                    // now, invoke the real thing in the caller context
                    return AccessController.doPrivileged(action, acc);
                }
                finally {
                    DelegatableThreadLocal.restore(savedDtl);
                }
            }
            finally {
                if (ccl != savedCcl) current.setContextClassLoader(savedCcl);
                if (priority < savedPriority) current.setPriority(savedPriority);
            }
        }});
    }

    /**
     * Temporarily recovers current thread's state from this snapshot and then
     * performs the specified action. In result, the action will observe
     * access control context, delegatable thread locals, and the context
     * class loader, as they existed when the snapshot was taken
     * (perhaps in a different thread). Upon completion, the original thread
     * state is restored.
     *
     * @param action the action to perform
     * @return the result of the action
     * @throws PrivilegedActionException if action threw an exception
     */
    public Object perform(final PrivilegedExceptionAction action)
        throws PrivilegedActionException
    {
        // verbatim copy of perform(PrivilegedAction), but with exception thrown

        final Thread current = Thread.currentThread();
        final AccessControlContext acc = AccessController.getContext();
        // make sure that we can modify this worker thread even if we're
        // called from un-privileged code
        return AccessController.doPrivileged(new PrivilegedExceptionAction() {
          public Object run() throws Exception {
            ClassLoader savedCcl = current.getContextClassLoader();
            int savedPriority = current.getPriority();
            Map savedDtl;
            try {
                if (ccl != savedCcl) current.setContextClassLoader(ccl);
                // do not boost priority if the request comes from
                // high-priority thread: due to high-cost of setPriority
                // method, it would usually result in a slow-down. However,
                // we do prohibit low-priority threads to schedule
                // higher-priority tasks
                if (priority < savedPriority) current.setPriority(priority);
                savedDtl = DelegatableThreadLocal.delegate(delegThreadLocals);
                try {
                    // now, invoke the real thing in the caller context
                    return AccessController.doPrivileged(action, acc);
                }
                catch (PrivilegedActionException pae) {
                    throw pae.getException();
                }
                finally {
                    DelegatableThreadLocal.restore(savedDtl);
                }
            }
            finally {
                if (ccl != savedCcl) current.setContextClassLoader(savedCcl);
                if (priority < savedPriority) current.setPriority(savedPriority);
            }
        }});
    }

    public PrivilegedAction wrap(final PrivilegedAction action) {
        return new PrivilegedAction() {
            public Object run() {
                return perform(action);
            }
        };
    }

    public PrivilegedExceptionAction wrap(final PrivilegedExceptionAction action) {
        return new PrivilegedExceptionAction() {
            public Object run() throws PrivilegedActionException {
                return perform(action);
            }
        };
    }

    public int hashCode() {
        if (hash != 0) return hash;
        hash = priority + ((ccl != null ? ccl.hashCode() : 1) ^ delegThreadLocals.hashCode());
        return hash;
    }

    public boolean equals(Object other) {
        if (other == this) return true;
        if (! (other instanceof ThreadContext)) return false;
        ThreadContext that = (ThreadContext) other;
        return this.priority == that.priority &&
            (this.ccl == null ? that.ccl == null : this.ccl.equals(that.ccl)) &&
            this.delegThreadLocals.equals(that.delegThreadLocals);
    }
}
