import com.google.common.util.concurrent.MoreExecutors
import datadog.trace.agent.test.AgentTestRunner
import datadog.trace.api.Trace
import datadog.trace.bootstrap.instrumentation.java.concurrent.RunnableWrapper
import datadog.trace.core.DDSpan
import org.apache.tomcat.util.threads.TaskQueue
import spock.lang.Shared

import java.lang.reflect.InvocationTargetException
import java.util.concurrent.AbstractExecutorService
import java.util.concurrent.ArrayBlockingQueue
import java.util.concurrent.Callable
import java.util.concurrent.CompletableFuture
import java.util.concurrent.ExecutionException
import java.util.concurrent.Executor
import java.util.concurrent.Executors
import java.util.concurrent.ForkJoinPool
import java.util.concurrent.ForkJoinTask
import java.util.concurrent.Future
import java.util.concurrent.LinkedBlockingQueue
import java.util.concurrent.PriorityBlockingQueue
import java.util.concurrent.RejectedExecutionException
import java.util.concurrent.ScheduledThreadPoolExecutor
import java.util.concurrent.ThreadPoolExecutor
import java.util.concurrent.TimeUnit
import java.util.concurrent.TimeoutException

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
import static org.junit.Assume.assumeTrue

class ExecutorInstrumentationTest extends AgentTestRunner {
  @Shared
  def executeRunnable = { e, c -> e.execute((Runnable) c) }
  @Shared
  def executePriorityTask = { e, c -> e.execute(new ComparableAsyncChild(0, (Runnable) c)) }
  @Shared
  def executeForkJoinTask = { e, c -> e.execute((ForkJoinTask) c) }
  @Shared
  def submitRunnable = { e, c -> e.submit((Runnable) c) }
  @Shared
  def submitCallable = { e, c -> e.submit((Callable) c) }
  @Shared
  def submitForkJoinTask = { e, c -> e.submit((ForkJoinTask) c) }
  @Shared
  def invokeAll = { e, c -> e.invokeAll([(Callable) c]) }
  @Shared
  def invokeAllTimeout = { e, c -> e.invokeAll([(Callable) c], 10, TimeUnit.SECONDS) }
  @Shared
  def invokeAny = { e, c -> e.invokeAny([(Callable) c]) }
  @Shared
  def invokeAnyTimeout = { e, c -> e.invokeAny([(Callable) c], 10, TimeUnit.SECONDS) }
  @Shared
  def invokeForkJoinTask = { e, c -> e.invoke((ForkJoinTask) c) }
  @Shared
  def scheduleRunnable = { e, c -> e.schedule((Runnable) c, 10, TimeUnit.MILLISECONDS) }
  @Shared
  def scheduleCallable = { e, c -> e.schedule((Callable) c, 10, TimeUnit.MILLISECONDS) }

  @Override
  void configurePreAgent() {
    super.configurePreAgent()

    injectSysConfig("dd.trace.executors", "ExecutorInstrumentationTest\$CustomThreadPoolExecutor")
  }

  def "#poolImpl '#name' propagates"() {
    setup:
    assumeTrue(poolImpl != null) // skip for Java 7 CompletableFuture, non-Linux Netty EPoll
    def pool = poolImpl
    def m = method

    new Runnable() {
      @Override
      @Trace(operationName = "parent")
      void run() {
        activeScope().setAsyncPropagation(true)
        // this child will have a span
        m(pool, new JavaAsyncChild())
        // this child won't
        m(pool, new JavaAsyncChild(false, false))
        blockUntilChildSpansFinished(1)
      }
    }.run()

    TEST_WRITER.waitForTraces(1)
    List<DDSpan> trace = TEST_WRITER.get(0)

    expect:
    TEST_WRITER.size() == 1
    trace.size() == 2
    trace.get(0).operationName == "parent"
    trace.get(1).operationName == "asyncChild"
    trace.get(1).parentId == trace.get(0).spanId

    cleanup:
    if (pool?.hasProperty("shutdown")) {
      pool?.shutdown()
    }

    // Unfortunately, there's no simple way to test the cross product of methods/pools.
    where:
    name                     | method              | poolImpl
    "execute Runnable"       | executeRunnable     | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1))
    "submit Runnable"        | submitRunnable      | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1))
    "submit Callable"        | submitCallable      | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1))
    "invokeAll"              | invokeAll           | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1))
    "invokeAll with timeout" | invokeAllTimeout    | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1))
    "invokeAny"              | invokeAny           | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1))
    "invokeAny with timeout" | invokeAnyTimeout    | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1))

    "execute Priority task"  | executePriorityTask | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new PriorityBlockingQueue<ComparableAsyncChild>(10))

    // Scheduled executor has additional methods and also may get disabled because it wraps tasks
    "execute Runnable"       | executeRunnable     | new ScheduledThreadPoolExecutor(1)
    "submit Runnable"        | submitRunnable      | new ScheduledThreadPoolExecutor(1)
    "submit Callable"        | submitCallable      | new ScheduledThreadPoolExecutor(1)
    "invokeAll"              | invokeAll           | new ScheduledThreadPoolExecutor(1)
    "invokeAll with timeout" | invokeAllTimeout    | new ScheduledThreadPoolExecutor(1)
    "invokeAny"              | invokeAny           | new ScheduledThreadPoolExecutor(1)
    "invokeAny with timeout" | invokeAnyTimeout    | new ScheduledThreadPoolExecutor(1)
    "schedule Runnable"      | scheduleRunnable    | new ScheduledThreadPoolExecutor(1)
    "schedule Callable"      | scheduleCallable    | new ScheduledThreadPoolExecutor(1)

    // ForkJoinPool has additional set of method overloads for ForkJoinTask to deal with
    "execute Runnable"       | executeRunnable     | new ForkJoinPool()
    "execute ForkJoinTask"   | executeForkJoinTask | new ForkJoinPool()
    "submit Runnable"        | submitRunnable      | new ForkJoinPool()
    "submit Callable"        | submitCallable      | new ForkJoinPool()
    "submit ForkJoinTask"    | submitForkJoinTask  | new ForkJoinPool()
    "invoke ForkJoinTask"    | invokeForkJoinTask  | new ForkJoinPool()
    "invokeAll"              | invokeAll           | new ForkJoinPool()
    "invokeAll with timeout" | invokeAllTimeout    | new ForkJoinPool()
    "invokeAny"              | invokeAny           | new ForkJoinPool()
    "invokeAny with timeout" | invokeAnyTimeout    | new ForkJoinPool()

    // CustomThreadPoolExecutor would normally be disabled except enabled above.
    "execute Runnable"       | executeRunnable     | new CustomThreadPoolExecutor()
    "submit Runnable"        | submitRunnable      | new CustomThreadPoolExecutor()
    "submit Callable"        | submitCallable      | new CustomThreadPoolExecutor()
    "invokeAll"              | invokeAll           | new CustomThreadPoolExecutor()
    "invokeAll with timeout" | invokeAllTimeout    | new CustomThreadPoolExecutor()
    "invokeAny"              | invokeAny           | new CustomThreadPoolExecutor()
    "invokeAny with timeout" | invokeAnyTimeout    | new CustomThreadPoolExecutor()


    "execute Runnable"       | executeRunnable     | new TypeAwareThreadPoolExecutor()
    "submit Runnable"        | submitRunnable      | new TypeAwareThreadPoolExecutor()
    "submit Callable"        | submitCallable      | new TypeAwareThreadPoolExecutor()
    "invokeAll"              | invokeAll           | new TypeAwareThreadPoolExecutor()
    "invokeAll with timeout" | invokeAllTimeout    | new TypeAwareThreadPoolExecutor()
    "invokeAny"              | invokeAny           | new TypeAwareThreadPoolExecutor()
    "invokeAny with timeout" | invokeAnyTimeout    | new TypeAwareThreadPoolExecutor()

    // Internal executor used by CompletableFuture
    "execute Runnable"       | executeRunnable     | java7SafeCompletableFutureThreadPerTaskExecutor()

    // java.util.concurrent.Executors$FinalizableDelegatedExecutorService
    "execute Runnable"       | executeRunnable     | Executors.newSingleThreadExecutor()
    "submit Runnable"        | submitRunnable      | Executors.newSingleThreadExecutor()
    "submit Callable"        | submitCallable      | Executors.newSingleThreadExecutor()
    "invokeAll"              | invokeAll           | Executors.newSingleThreadExecutor()
    "invokeAll with timeout" | invokeAllTimeout    | Executors.newSingleThreadExecutor()
    "invokeAny"              | invokeAny           | Executors.newSingleThreadExecutor()
    "invokeAny with timeout" | invokeAnyTimeout    | Executors.newSingleThreadExecutor()

    // java.util.concurrent.Executors$DelegatedExecutorService
    "execute Runnable"       | executeRunnable     | Executors.unconfigurableExecutorService(Executors.newSingleThreadExecutor())
    "submit Runnable"        | submitRunnable      | Executors.unconfigurableExecutorService(Executors.newSingleThreadExecutor())
    "submit Callable"        | submitCallable      | Executors.unconfigurableExecutorService(Executors.newSingleThreadExecutor())
    "invokeAll"              | invokeAll           | Executors.unconfigurableExecutorService(Executors.newSingleThreadExecutor())
    "invokeAll with timeout" | invokeAllTimeout    | Executors.unconfigurableExecutorService(Executors.newSingleThreadExecutor())
    "invokeAny"              | invokeAny           | Executors.unconfigurableExecutorService(Executors.newSingleThreadExecutor())
    "invokeAny with timeout" | invokeAnyTimeout    | Executors.unconfigurableExecutorService(Executors.newSingleThreadExecutor())


    "execute Runnable"       | executeRunnable     | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "submit Runnable"        | submitRunnable      | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "submit Callable"        | submitCallable      | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "invokeAll"              | invokeAll           | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "invokeAll with timeout" | invokeAllTimeout    | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "invokeAny"              | invokeAny           | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "invokeAny with timeout" | invokeAnyTimeout    | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "schedule Runnable"      | scheduleRunnable    | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "schedule Callable"      | scheduleCallable    | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())

    // tomcat
    "execute Runnable"       | executeRunnable     | new org.apache.tomcat.util.threads.ThreadPoolExecutor(1, 1, 5, TimeUnit.SECONDS, new TaskQueue())
    "submit Runnable"        | submitRunnable      | new org.apache.tomcat.util.threads.ThreadPoolExecutor(1, 1, 5, TimeUnit.SECONDS, new TaskQueue())
    "submit Callable"        | submitCallable      | new org.apache.tomcat.util.threads.ThreadPoolExecutor(1, 1, 5, TimeUnit.SECONDS, new TaskQueue())
    "invokeAll"              | invokeAll           | new org.apache.tomcat.util.threads.ThreadPoolExecutor(1, 1, 5, TimeUnit.SECONDS, new TaskQueue())
    "invokeAll with timeout" | invokeAllTimeout    | new org.apache.tomcat.util.threads.ThreadPoolExecutor(1, 1, 5, TimeUnit.SECONDS, new TaskQueue())
    "invokeAny"              | invokeAny           | new org.apache.tomcat.util.threads.ThreadPoolExecutor(1, 1, 5, TimeUnit.SECONDS, new TaskQueue())
    "invokeAny with timeout" | invokeAnyTimeout    | new org.apache.tomcat.util.threads.ThreadPoolExecutor(1, 1, 5, TimeUnit.SECONDS, new TaskQueue())

    // guava
    "execute Runnable"       | executeRunnable     | MoreExecutors.directExecutor()

    "execute Runnable"       | executeRunnable     | MoreExecutors.newDirectExecutorService()
    "submit Runnable"        | submitRunnable      | MoreExecutors.newDirectExecutorService()
    "submit Callable"        | submitCallable      | MoreExecutors.newDirectExecutorService()
    "invokeAll"              | invokeAll           | MoreExecutors.newDirectExecutorService()
    "invokeAll with timeout" | invokeAllTimeout    | MoreExecutors.newDirectExecutorService()
    "invokeAny"              | invokeAny           | MoreExecutors.newDirectExecutorService()
    "invokeAny with timeout" | invokeAnyTimeout    | MoreExecutors.newDirectExecutorService()

    "execute Runnable"       | executeRunnable     | MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor())
    "submit Runnable"        | submitRunnable      | MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor())
    "submit Callable"        | submitCallable      | MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor())
    "invokeAll"              | invokeAll           | MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor())
    "invokeAll with timeout" | invokeAllTimeout    | MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor())
    "invokeAny"              | invokeAny           | MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor())
    "invokeAny with timeout" | invokeAnyTimeout    | MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor())

    "execute Runnable"       | executeRunnable     | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
    "submit Runnable"        | submitRunnable      | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
    "submit Callable"        | submitCallable      | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
    "invokeAll"              | invokeAll           | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
    "invokeAll with timeout" | invokeAllTimeout    | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
    "invokeAny"              | invokeAny           | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
    "invokeAny with timeout" | invokeAnyTimeout    | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
    "schedule Runnable"      | scheduleRunnable    | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
    "schedule Callable"      | scheduleCallable    | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())

  }

  def "#poolImpl '#name' wraps"() {
    setup:
    def pool = poolImpl
    def m = method
    def w = wrap

    JavaAsyncChild child = new JavaAsyncChild(true, true)
    new Runnable() {
      @Override
      @Trace(operationName = "parent")
      void run() {
        activeScope().setAsyncPropagation(true)
        m(pool, w(child))
      }
    }.run()
    child.unblock()

    TEST_WRITER.waitForTraces(1)

    expect:
    TEST_WRITER.size() == 1
    TEST_WRITER.get(0).size() == 2
    TEST_WRITER.get(0).get(1).operationName == "parent"
    TEST_WRITER.get(0).get(0).operationName == "asyncChild"
    TEST_WRITER.get(0).get(0).parentId == TEST_WRITER.get(0).get(1).spanId

    cleanup:
    pool?.shutdown()

    where:
    name                | method           | wrap                        | poolImpl
    "execute Runnable"  | executeRunnable  | { new RunnableWrapper(it) } | new ScheduledThreadPoolExecutor(1)
    "submit Runnable"   | submitRunnable   | { new RunnableWrapper(it) } | new ScheduledThreadPoolExecutor(1)
    "schedule Runnable" | scheduleRunnable | { new RunnableWrapper(it) } | new ScheduledThreadPoolExecutor(1)
  }

  def "#poolImpl '#name' reports after canceled jobs"() {
    setup:
    assumeTrue(poolImpl != null) // skip for non-Linux Netty EPoll
    def pool = poolImpl
    def m = method
    List<JavaAsyncChild> children = new ArrayList<>()
    List<Future> jobFutures = new ArrayList<>()

    new Runnable() {
      @Override
      @Trace(operationName = "parent")
      void run() {
        activeScope().setAsyncPropagation(true)
        try {
          for (int i = 0; i < 20; ++i) {
            final JavaAsyncChild child = new JavaAsyncChild(false, true)
            children.add(child)
            try {
              Future f = m(pool, child)
              jobFutures.add(f)
            } catch (InvocationTargetException e) {
              throw e.getCause()
            }
          }
        } catch (RejectedExecutionException e) {
        }

        for (Future f : jobFutures) {
          f.cancel(false)
        }
        for (JavaAsyncChild child : children) {
          child.unblock()
        }
      }
    }.run()

    TEST_WRITER.waitForTraces(1)

    expect:
    // FIXME: we should improve this test to make sure continuations are actually closed
    TEST_WRITER.size() == 1

    where:
    name                  | method             | poolImpl
    "submit Runnable"     | submitRunnable     | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1))
    "submit Callable"     | submitCallable     | new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1))

    // Scheduled executor has additional methods and also may get disabled because it wraps tasks
    "submit Runnable"     | submitRunnable     | new ScheduledThreadPoolExecutor(1)
    "submit Callable"     | submitCallable     | new ScheduledThreadPoolExecutor(1)
    "schedule Runnable"   | scheduleRunnable   | new ScheduledThreadPoolExecutor(1)
    "schedule Callable"   | scheduleCallable   | new ScheduledThreadPoolExecutor(1)

    // ForkJoinPool has additional set of method overloads for ForkJoinTask to deal with
    "submit Runnable"     | submitRunnable     | new ForkJoinPool()
    "submit Callable"     | submitCallable     | new ForkJoinPool()
    "submit ForkJoinTask" | submitForkJoinTask | new ForkJoinPool()

    // java.util.concurrent.Executors$FinalizableDelegatedExecutorService
    "submit Runnable"     | submitRunnable     | Executors.newSingleThreadExecutor()
    "submit Callable"     | submitCallable     | Executors.newSingleThreadExecutor()

    // java.util.concurrent.Executors$DelegatedExecutorService
    "submit Runnable"     | submitRunnable     | Executors.unconfigurableExecutorService(Executors.newSingleThreadExecutor())
    "submit Callable"     | submitCallable     | Executors.unconfigurableExecutorService(Executors.newSingleThreadExecutor())


    "submit Runnable"     | submitRunnable     | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "submit Callable"     | submitCallable     | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "schedule Runnable"   | scheduleRunnable   | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())
    "schedule Callable"   | scheduleCallable   | Executors.unconfigurableScheduledExecutorService(Executors.newSingleThreadScheduledExecutor())

    // tomcat
    "submit Runnable"     | submitRunnable     | new org.apache.tomcat.util.threads.ThreadPoolExecutor(1, 1, 5, TimeUnit.SECONDS, new TaskQueue())
    "submit Callable"     | submitCallable     | new org.apache.tomcat.util.threads.ThreadPoolExecutor(1, 1, 5, TimeUnit.SECONDS, new TaskQueue())

    // guava
    // FIXME - these need better rejection handling to pass reliably
//    "submit Runnable"     | submitRunnable     | MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor())
//    "submit Callable"     | submitCallable     | MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor())

//    "submit Runnable"     | submitRunnable     | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
//    "submit Callable"     | submitCallable     | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
//    "schedule Runnable"   | scheduleRunnable   | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
//    "schedule Callable"   | scheduleCallable   | MoreExecutors.listeningDecorator(Executors.newSingleThreadScheduledExecutor())
  }

  private static Executor java7SafeCompletableFutureThreadPerTaskExecutor() {
    try {
      return new CompletableFuture.ThreadPerTaskExecutor()
    } catch (NoClassDefFoundError e) {
      return null
    }
  }

  static class CustomThreadPoolExecutor extends AbstractExecutorService {
    volatile running = true
    def workQueue = new LinkedBlockingQueue<Runnable>(10)

    def worker = new Runnable() {
      void run() {
        try {
          while (running) {
            def runnable = workQueue.take()
            runnable.run()
          }
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt()
        } catch (Exception e) {
          e.printStackTrace()
        }
      }
    }

    def workerThread = new Thread(worker, "ExecutorTestThread")

    private CustomThreadPoolExecutor() {
      workerThread.start()
    }

    @Override
    void shutdown() {
      running = false
      workerThread.interrupt()
    }

    @Override
    List<Runnable> shutdownNow() {
      running = false
      workerThread.interrupt()
      return []
    }

    @Override
    boolean isShutdown() {
      return !running
    }

    @Override
    boolean isTerminated() {
      return workerThread.isAlive()
    }

    @Override
    boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
      workerThread.join(unit.toMillis(timeout))
      return true
    }

    @Override
    def <T> Future<T> submit(Callable<T> task) {
      def future = newTaskFor(task)
      execute(future)
      return future
    }

    @Override
    def <T> Future<T> submit(Runnable task, T result) {
      def future = newTaskFor(task, result)
      execute(future)
      return future
    }

    @Override
    Future<?> submit(Runnable task) {
      def future = newTaskFor(task, null)
      execute(future)
      return future
    }

    @Override
    def <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) throws InterruptedException {
      return super.invokeAll(tasks)
    }

    @Override
    def <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException {
      return super.invokeAll(tasks)
    }

    @Override
    def <T> T invokeAny(Collection<? extends Callable<T>> tasks) throws InterruptedException, ExecutionException {
      return super.invokeAny(tasks)
    }

    @Override
    def <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
      return super.invokeAny(tasks)
    }

    @Override
    void execute(Runnable command) {
      workQueue.put(command)
    }
  }
}
