package com.atlassian.sal.core.executor;
   
   import com.atlassian.sal.api.executor.ThreadLocalContextManager;
   import com.google.common.base.Function;
   import com.google.common.collect.Collections2;
   
   import java.util.Collection;
   import java.util.List;
   import java.util.concurrent.Callable;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Future;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  
  import static com.google.common.base.Preconditions.checkNotNull;
  
  /**
   * Executor service that wraps executing callables and runnables in a wrapper that transfers the thread local state of
   * the caller to the thread of the executing task.
   *
   * @since 2.0
   */
  public class ThreadLocalDelegateExecutorService<C> implements ExecutorService {
      protected final ThreadLocalContextManager<C> manager;
      private final ExecutorService delegate;
  
      public ThreadLocalDelegateExecutorService(ThreadLocalContextManager<C> manager, ExecutorService delegate) {
          this.manager = checkNotNull(manager);
          this.delegate = checkNotNull(delegate);
      }
  
      public void shutdown() {
          delegate.shutdown();
      }
  
      public List<Runnable> shutdownNow() {
          return delegate.shutdownNow();
      }
  
      public boolean isShutdown() {
          return delegate.isShutdown();
      }
  
      public boolean isTerminated() {
          return delegate.isTerminated();
      }
  
      public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
          return delegate.awaitTermination(timeout, unit);
      }
  
      public <T> Future<T> submit(Callable<T> callable) {
          return delegate.submit(threadLocalDelegateCallable(callable));
      }
  
      public <T> Future<T> submit(Runnable runnable, T t) {
          return delegate.submit(threadLocalDelegateRunnable(runnable), t);
      }
  
      public Future<?> submit(Runnable runnable) {
          return delegate.submit(threadLocalDelegateRunnable(runnable));
      }
  
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> callables) throws InterruptedException {
          return delegate.invokeAll(threadLocalDelegateCallableCollection(callables));
      }
  
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> callables, long timeout, TimeUnit unit) throws InterruptedException {
          return delegate.invokeAll(threadLocalDelegateCallableCollection(callables), timeout, unit);
      }
  
      public <T> T invokeAny(Collection<? extends Callable<T>> callables) throws InterruptedException, ExecutionException {
          return delegate.invokeAny(threadLocalDelegateCallableCollection(callables));
      }
  
      public <T> T invokeAny(Collection<? extends Callable<T>> callables, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
          return delegate.invokeAny(threadLocalDelegateCallableCollection(callables), timeout, unit);
      }
  
      public void execute(Runnable runnable) {
          delegate.execute(threadLocalDelegateRunnable(runnable));
      }
  
      private ThreadLocalDelegateRunnable threadLocalDelegateRunnable(Runnable runnable) {
          return new ThreadLocalDelegateRunnable<C>(manager, runnable);
      }
  
      private <T> ThreadLocalDelegateCallable<C, T> threadLocalDelegateCallable(Callable<T> callable) {
          return new ThreadLocalDelegateCallable<C, T>(manager, callable);
      }
  
      private <T> Collection<ThreadLocalDelegateCallable<C, T>> threadLocalDelegateCallableCollection(Collection<? extends Callable<T>> callables) {
          return Collections2.transform(callables, new Function<Callable<T>, ThreadLocalDelegateCallable<C, T>>() {
              public ThreadLocalDelegateCallable<C, T> apply(Callable<T> callable) {
                  return threadLocalDelegateCallable(callable);
              }
          });
      }
 }