package com.atlassian.sal.core.executor;
   
   import com.atlassian.sal.api.executor.ThreadLocalDelegateExecutorFactory;
   
   import javax.annotation.Nonnull;
   import javax.annotation.ParametersAreNonnullByDefault;
   import java.util.concurrent.Callable;
   import java.util.concurrent.ScheduledExecutorService;
   import java.util.concurrent.ScheduledFuture;
  import java.util.concurrent.TimeUnit;
  
  import static com.google.common.base.Preconditions.checkNotNull;
  
  /**
   * Scheduled 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
   */
  @ParametersAreNonnullByDefault
  public class ThreadLocalDelegateScheduledExecutorService extends ThreadLocalDelegateExecutorService implements ScheduledExecutorService {
      private final ScheduledExecutorService delegate;
      private final ThreadLocalDelegateExecutorFactory deletegateExecutorFactory;
  
      public ThreadLocalDelegateScheduledExecutorService(ScheduledExecutorService delegate,
                                                         ThreadLocalDelegateExecutorFactory deletegateExecutorFactory) {
          super(delegate, deletegateExecutorFactory);
          this.delegate = checkNotNull(delegate);
          this.deletegateExecutorFactory = checkNotNull(deletegateExecutorFactory);
      }
  
      @Override
      @Nonnull
      public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
          return delegate.schedule(deletegateExecutorFactory.createRunnable(command), delay, unit);
      }
  
      @Override
      @Nonnull
      public <V> ScheduledFuture<V> schedule(Callable<V> callable, long delay, TimeUnit unit) {
          return delegate.schedule(deletegateExecutorFactory.createCallable(callable), delay, unit);
      }
  
      @Override
      @Nonnull
      public ScheduledFuture<?> scheduleAtFixedRate(Runnable command, long initialDelay, long period, TimeUnit unit) {
          return delegate.scheduleAtFixedRate(deletegateExecutorFactory.createRunnable(command), initialDelay, period, unit);
      }
  
      @Override
      @Nonnull
      public ScheduledFuture<?> scheduleWithFixedDelay(Runnable command, long initialDelay, long delay, TimeUnit unit) {
          return delegate.scheduleWithFixedDelay(deletegateExecutorFactory.createRunnable(command), initialDelay, delay, unit);
      }
  }