package com.atlassian.sal.core.executor;
   
   import com.atlassian.sal.api.executor.ThreadLocalDelegateExecutorFactory;
   import com.google.common.base.Function;
   import com.google.common.collect.Collections2;
   
   import javax.annotation.Nonnull;
   import javax.annotation.Nullable;
   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 implements ExecutorService {
      private final ExecutorService delegate;
      private final ThreadLocalDelegateExecutorFactory delegateExecutorFactory;
  
      public ThreadLocalDelegateExecutorService(ExecutorService delegate, ThreadLocalDelegateExecutorFactory delegateExecutorFactory) {
          this.delegate = checkNotNull(delegate);
          this.delegateExecutorFactory = checkNotNull(delegateExecutorFactory);
      }
  
      @Override
      public void shutdown() {
          delegate.shutdown();
      }
  
      @Override
      @Nonnull
      public List<Runnable> shutdownNow() {
          return delegate.shutdownNow();
      }
  
      @Override
      public boolean isShutdown() {
          return delegate.isShutdown();
      }
  
      @Override
      public boolean isTerminated() {
          return delegate.isTerminated();
      }
  
      @Override
      public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
          return delegate.awaitTermination(timeout, unit);
      }
  
      @Override
      @Nonnull
      public <T> Future<T> submit(Callable<T> callable) {
          return delegate.submit(threadLocalDelegateCallable(callable));
      }
  
      @Override
      @Nonnull
      public <T> Future<T> submit(Runnable runnable, @Nullable T result) {
          return delegate.submit(threadLocalDelegateRunnable(runnable), result);
      }
  
      @Override
      @Nonnull
      public Future<?> submit(Runnable runnable) {
          return delegate.submit(threadLocalDelegateRunnable(runnable));
      }
  
      @Override
      @Nonnull
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> callables) throws InterruptedException {
          return delegate.invokeAll(threadLocalDelegateCallableCollection(callables));
      }
  
      @Override
      @Nonnull
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> callables, long timeout, TimeUnit unit) throws InterruptedException {
          return delegate.invokeAll(threadLocalDelegateCallableCollection(callables), timeout, unit);
      }
  
      @Override
      @Nonnull
      public <T> T invokeAny(Collection<? extends Callable<T>> callables) throws InterruptedException, ExecutionException {
          return delegate.invokeAny(threadLocalDelegateCallableCollection(callables));
      }
  
      @Override
      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 Runnable threadLocalDelegateRunnable(Runnable runnable) {
         return delegateExecutorFactory.createRunnable(runnable);
     }
 
     private <T> Callable<T> threadLocalDelegateCallable(Callable<T> callable) {
         return delegateExecutorFactory.createCallable(callable);
     }
 
     private <T> Collection<? extends Callable<T>> threadLocalDelegateCallableCollection(Collection<? extends Callable<T>> callables) {
         return Collections2.transform(callables, (Function<Callable<T>, Callable<T>>) this::threadLocalDelegateCallable);
     }
 }