content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.executors;
  
  import io.opentelemetry.api.trace.SpanKind;
  import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  import java.util.concurrent.Callable;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Future;
  import java.util.concurrent.TimeUnit;
  import java.util.function.Function;
  import org.junit.jupiter.api.AfterAll;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.TestInstance;
  import org.junit.jupiter.api.function.ThrowingConsumer;
  
  @TestInstance(TestInstance.Lifecycle.PER_CLASS)
  public abstract class AbstractExecutorServiceTest<T extends ExecutorService, U extends TestTask> {
  
    private final T executor;
    private final InstrumentationExtension testing;
  
    protected AbstractExecutorServiceTest(T executor, InstrumentationExtension testing) {
      this.executor = executor;
      this.testing = testing;
    }
  
    protected abstract U newTask(boolean doTraceableWork, boolean blockThread);
  
    protected T executor() {
      return executor;
    }
  
    @AfterAll
    void shutdown() throws InterruptedException {
      executor.shutdown();
      executor.awaitTermination(10, TimeUnit.SECONDS);
    }
  
    @Test
    void executeRunnable() {
      executeTwoTasks(executor::execute);
    }
  
    @Test
    void submitRunnable() {
      executeTwoTasks(task -> executor.submit((Runnable) task));
    }
  
    @Test
    void submitCallable() {
      executeTwoTasks(task -> executor.submit((Callable<?>) task));
    }
  
    @Test
    void invokeAll() {
      executeTwoTasks(task -> executor.invokeAll(Collections.singleton(task)));
    }
  
    @Test
    void invokeAllWithTimeout() {
      executeTwoTasks(task -> executor.invokeAll(Collections.singleton(task), 10, TimeUnit.SECONDS));
    }
  
    @Test
    void invokeAny() {
      executeTwoTasks(task -> executor.invokeAny(Collections.singleton(task)));
    }
  
    @Test
    void invokeAnyWithTimeout() {
      executeTwoTasks(task -> executor.invokeAny(Collections.singleton(task), 10, TimeUnit.SECONDS));
    }
  
    @Test
    void executeLambdaRunnable() {
      executeTwoTasks(task -> executor.execute(() -> task.run()));
    }
  
    @Test
    void submitLambdaRunnable() {
      executeTwoTasks(task -> executor.submit(() -> task.run()));
    }
  
    @Test
    void submitLambdaCallable() {
      executeTwoTasks(
          task ->
              executor.submit(
                  () -> {
                    task.run();
                    return null;
                  }));
    }
  
    @Test
    void submitRunnableAndCancel() {
      executeAndCancelTasks(task -> executor.submit((Runnable) task));
    }
  
    @Test
    void submitCallableAndCancel() {
      executeAndCancelTasks(task -> executor.submit((Callable<?>) task));
    }
  
    protected final void executeTwoTasks(ThrowingConsumer<U> task) {
      testing.runWithSpan(
          "parent",
          () -> {
            // this child will have a span
            U child1 = newTask(true, false);
            // this child won't
            U child2 = newTask(false, false);
            try {
              task.accept(child1);
              task.accept(child2);
            } catch (Throwable t) {
              throw new AssertionError(t);
            }
            child1.waitForCompletion();
            child2.waitForCompletion();
          });
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName("asyncChild")
                          .hasKind(SpanKind.INTERNAL)
                          .hasParent(trace.getSpan(0))));
    }
  
    protected final void executeAndCancelTasks(Function<U, Future<?>> task) {
      List<U> children = new ArrayList<>();
      List<Future<?>> jobFutures = new ArrayList<>();
  
      testing.runWithSpan(
          "parent",
          () -> {
            for (int i = 0; i < 20; i++) {
              // Our current instrumentation instrumentation does not behave very well
              // if we try to reuse Callable/Runnable. Namely we would be getting 'orphaned'
              // child traces sometimes since state can contain only one parent span - and
              // we do not really have a good way for attributing work to correct parent span
              // if we reuse Callable/Runnable.
              // Solution for now is to never reuse a Callable/Runnable.
              U child = newTask(false, true);
              children.add(child);
              Future<?> f = task.apply(child);
              jobFutures.add(f);
            }
  
            jobFutures.forEach(f -> f.cancel(false));
            children.forEach(U::unblock);
          });
  
      // Just check there is a single trace, this test is primarily to make sure that scopes aren't
      // leaked on cancellation.
      testing.waitAndAssertTraces(trace -> {});
    }
  }
  