content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.executors;
  
  import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
  import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.List;
  import java.util.concurrent.AbstractExecutorService;
  import java.util.concurrent.ArrayBlockingQueue;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.Callable;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.ForkJoinPool;
  import java.util.concurrent.ForkJoinTask;
  import java.util.concurrent.Future;
  import java.util.concurrent.RunnableFuture;
  import java.util.concurrent.ScheduledThreadPoolExecutor;
  import java.util.concurrent.ThreadPoolExecutor;
  import java.util.concurrent.TimeUnit;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.extension.RegisterExtension;
  
  abstract class ExecutorInstrumentationTest<T extends ExecutorService>
      extends AbstractExecutorServiceTest<T, JavaAsyncChild> {
  
    @RegisterExtension
    static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
  
    ExecutorInstrumentationTest(T executor) {
      super(executor, testing);
    }
  
    @Override
    protected JavaAsyncChild newTask(boolean doTraceableWork, boolean blockThread) {
      return new JavaAsyncChild(doTraceableWork, blockThread);
    }
  
    static class ThreadPoolExecutorTest extends ExecutorInstrumentationTest<ThreadPoolExecutor> {
      ThreadPoolExecutorTest() {
        super(new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<>(20)));
      }
    }
  
    static class WorkStealingPoolTest extends ExecutorInstrumentationTest<ExecutorService> {
      public WorkStealingPoolTest() {
        super(Executors.newWorkStealingPool(2));
      }
    }
  
    static class ScheduledThreadPoolExecutorTest
        extends ExecutorInstrumentationTest<ScheduledThreadPoolExecutor> {
      ScheduledThreadPoolExecutorTest() {
        super(new ScheduledThreadPoolExecutor(1));
      }
  
      @Test
      void scheduleRunnable() {
        executeTwoTasks(task -> executor().schedule((Runnable) task, 10, TimeUnit.MICROSECONDS));
      }
  
      @Test
      void scheduleCallable() {
        executeTwoTasks(task -> executor().schedule((Callable<?>) task, 10, TimeUnit.MICROSECONDS));
      }
  
      @Test
      void scheduleLambdaRunnable() {
        executeTwoTasks(task -> executor().schedule(() -> task.run(), 10, TimeUnit.MICROSECONDS));
      }
  
      @Test
      void scheduleLambdaCallable() {
        executeTwoTasks(
            task ->
                executor()
                    .schedule(
                        () -> {
                          task.run();
                          return null;
                        },
                        10,
                        TimeUnit.MICROSECONDS));
      }
  
      @Test
      void scheduleRunnableAndCancel() {
        executeAndCancelTasks(
            task -> executor().schedule((Runnable) task, 10, TimeUnit.MICROSECONDS));
      }
  
      @Test
      void scheduleCallableAndCancel() {
        executeAndCancelTasks(
            task -> executor().schedule((Callable<?>) task, 10, TimeUnit.MICROSECONDS));
      }
    }
  
    static class ForkJoinPoolTest extends ExecutorInstrumentationTest<ForkJoinPool> {
      ForkJoinPoolTest() {
        super(new ForkJoinPool(20));
      }
  
      @Test
      void invokeForkJoinTask() {
        executeTwoTasks(task -> executor().invoke((ForkJoinTask<?>) task));
      }
  
      @Test
      void submitForkJoinTask() {
        executeTwoTasks(task -> executor().submit((ForkJoinTask<?>) task));
      }
    }
  
    // CustomThreadPoolExecutor would normally be disabled except enabled by system property.
    static class CustomThreadPoolExecutorTest
        extends ExecutorInstrumentationTest<CustomThreadPoolExecutor> {
      CustomThreadPoolExecutorTest() {
        super(new CustomThreadPoolExecutor());
      }
    }
  
    @SuppressWarnings("RedundantOverride")
    private static class CustomThreadPoolExecutor extends AbstractExecutorService {
  
      private volatile boolean running = true;
  
      private final BlockingQueue<Runnable> workQueue = new ArrayBlockingQueue<>(20);
  
      private final Thread workerThread =
          new Thread(
              () -> {
                try {
                  while (running) {
                    Runnable runnable = workQueue.take();
                    runnable.run();
                  }
                } catch (InterruptedException e) {
                  Thread.currentThread().interrupt();
                  throw new AssertionError(e);
                } catch (Throwable t) {
                  throw new AssertionError(t);
                }
              },
              "ExecutorTestThread");
  
      private CustomThreadPoolExecutor() {
        workerThread.start();
      }
  
      @Override
      public void shutdown() {
        running = false;
        workerThread.interrupt();
      }
  
      @Override
      public List<Runnable> shutdownNow() {
        running = false;
        workerThread.interrupt();
        return Collections.emptyList();
      }
  
      @Override
      public boolean isShutdown() {
        return !running;
      }
  
      @Override
      public boolean isTerminated() {
        return workerThread.isAlive();
      }
  
      @Override
      public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        workerThread.join(unit.toMillis(timeout));
        return true;
      }
  
      @Override
      public <T> Future<T> submit(Callable<T> task) {
        RunnableFuture<T> future = newTaskFor(task);
        execute(future);
        return future;
      }
  
      @Override
      public <T> Future<T> submit(Runnable task, T result) {
        RunnableFuture<T> future = newTaskFor(task, result);
        execute(future);
        return future;
      }
  
      @Override
      public Future<?> submit(Runnable task) {
        RunnableFuture<?> future = newTaskFor(task, null);
        execute(future);
        return future;
      }
  
      @Override
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) {
        return Collections.singletonList(submit(tasks.stream().findFirst().get()));
      }
  
      @Override
      public <T> List<Future<T>> invokeAll(
          Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) {
        return Collections.singletonList(submit(tasks.stream().findFirst().get()));
      }
  
      @Override
      public <T> T invokeAny(Collection<? extends Callable<T>> tasks) {
        submit(tasks.stream().findFirst().get());
        return null;
      }
  
      @Override
      public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) {
        submit(tasks.stream().findFirst().get());
        return null;
      }
  
      @Override
      public void execute(Runnable command) {
        workQueue.add(command);
      }
    }
  }
  