// Copyright 2018 The Bazel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package com.google.devtools.build.lib.dynamic;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.truth.Truth.assertThat;
import static junit.framework.TestCase.fail;
import static org.junit.Assert.assertThrows;

import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import com.google.common.flogger.GoogleLogger;
import com.google.common.util.concurrent.Uninterruptibles;
import com.google.devtools.build.lib.actions.ActionContext;
import com.google.devtools.build.lib.actions.ActionExecutionContext;
import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.ActionKeyContext;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.actions.ArtifactRoot;
import com.google.devtools.build.lib.actions.BaseSpawn;
import com.google.devtools.build.lib.actions.DynamicStrategyRegistry;
import com.google.devtools.build.lib.actions.EmptyRunfilesSupplier;
import com.google.devtools.build.lib.actions.ExecException;
import com.google.devtools.build.lib.actions.Executor;
import com.google.devtools.build.lib.actions.ResourceSet;
import com.google.devtools.build.lib.actions.SandboxedSpawnStrategy;
import com.google.devtools.build.lib.actions.Spawn;
import com.google.devtools.build.lib.actions.SpawnResult;
import com.google.devtools.build.lib.actions.SpawnStrategy;
import com.google.devtools.build.lib.actions.UserExecException;
import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
import com.google.devtools.build.lib.actions.util.ActionsTestUtil.NullAction;
import com.google.devtools.build.lib.exec.BlazeExecutor;
import com.google.devtools.build.lib.exec.ExecutionOptions;
import com.google.devtools.build.lib.exec.ModuleActionContextRegistry;
import com.google.devtools.build.lib.exec.SpawnStrategyRegistry;
import com.google.devtools.build.lib.testutil.TestThread;
import com.google.devtools.build.lib.testutil.TestUtils;
import com.google.devtools.build.lib.util.AbruptExitException;
import com.google.devtools.build.lib.util.io.FileOutErr;
import com.google.devtools.build.lib.vfs.FileSystemUtils;
import com.google.devtools.build.lib.vfs.Path;
import com.google.devtools.build.lib.vfs.Root;
import com.google.devtools.build.lib.vfs.util.FileSystems;
import com.google.devtools.common.options.OptionsParser;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import javax.annotation.Nullable;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
import org.junit.runners.Parameterized.Parameter;
import org.junit.runners.Parameterized.Parameters;

/** Tests for {@link DynamicSpawnStrategy}. */
@RunWith(Parameterized.class)
public class DynamicSpawnStrategyTest {
  private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

  private Path testRoot;
  private ExecutorService executorServiceForCleanup;
  private FileOutErr outErr;
  private final ActionKeyContext actionKeyContext = new ActionKeyContext();

  @Parameters(name = "{index}: legacy={0}")
  public static Collection<Object[]> data() {
    return Arrays.asList(
        new Object[][] {
          {true}, {false},
        });
  }

  @Parameter public boolean legacyBehavior;

  /** Syntactic sugar to decrease and await for a latch in a single line. */
  private static void countDownAndWait(CountDownLatch countDownLatch) throws InterruptedException {
    countDownLatch.countDown();
    countDownLatch.await();
  }

  /** Hook to implement per-test custom logic in the {@link MockSpawnStrategy}. */
  @FunctionalInterface
  interface DoExec {
    DoExec NOTHING = (self, spawn, actionExecutionContext) -> {};

    void run(MockSpawnStrategy self, Spawn spawn, ActionExecutionContext actionExecutionContext)
        throws ExecException, InterruptedException;
  }

  /**
   * Minimal implementation of a strategy for testing purposes.
   *
   * <p>All the logic in here must be applicable to all tests. If any test needs to special-case
   * some aspect of this logic, then it must extend this subclass as necessary.
   */
  private class MockSpawnStrategy implements SandboxedSpawnStrategy {
    /** Identifier of this class for error reporting purposes. */
    private final String name;

    /** Lazily set to the spawn passed to {@link #exec} as soon as that hook is invoked. */
    @Nullable private volatile Spawn executedSpawn;

    /** Tracks whether {@link #exec} completed successfully or not. */
    private CountDownLatch succeeded = new CountDownLatch(1);

    /** Hook to implement per-test custom logic. */
    private final DoExec doExecBeforeStop;

    private final DoExec doExecAfterStop;

    MockSpawnStrategy(String name) {
      this(name, DoExec.NOTHING, DoExec.NOTHING);
    }

    MockSpawnStrategy(String name, DoExec doExecBeforeStop, DoExec doExecAfterStop) {
      this.name = name;
      this.doExecBeforeStop = doExecBeforeStop;
      this.doExecAfterStop = doExecAfterStop;
    }

    /** Helper to record an execution failure from within {@link #doExecBeforeStop}. */
    void failExecution(ActionExecutionContext actionExecutionContext) throws ExecException {
      try {
        FileSystemUtils.appendIsoLatin1(
            actionExecutionContext.getFileOutErr().getOutputPath(), "action failed with " + name);
      } catch (IOException e) {
        throw new IllegalStateException(e);
      }
      throw new UserExecException(name + " failed to execute the Spawn");
    }

    @Override
    public ImmutableList<SpawnResult> exec(
        Spawn spawn,
        ActionExecutionContext actionExecutionContext,
        @Nullable SandboxedSpawnStrategy.StopConcurrentSpawns stopConcurrentSpawns)
        throws ExecException, InterruptedException {
      executedSpawn = spawn;

      doExecBeforeStop.run(this, spawn, actionExecutionContext);
      if (stopConcurrentSpawns != null) {
        stopConcurrentSpawns.stop();
        doExecAfterStop.run(this, spawn, actionExecutionContext);
      }

      for (ActionInput output : spawn.getOutputFiles()) {
        try {
          FileSystemUtils.writeIsoLatin1(testRoot.getRelative(output.getExecPath()), name);
        } catch (IOException e) {
          throw new IllegalStateException(e);
        }
      }

      try {
        FileSystemUtils.appendIsoLatin1(
            actionExecutionContext.getFileOutErr().getOutputPath(),
            "output files written with " + name);
      } catch (IOException e) {
        throw new IllegalStateException(e);
      }

      succeeded.countDown();

      return ImmutableList.of();
    }

    @Override
    public ImmutableList<SpawnResult> exec(
        Spawn spawn, ActionExecutionContext actionExecutionContext) {
      throw new IllegalStateException("Not expected to be called");
    }

    @Override
    public boolean canExec(Spawn spawn, ActionContext.ActionContextRegistry actionContextRegistry) {
      return true;
    }

    @Nullable
    public Spawn getExecutedSpawn() {
      return executedSpawn;
    }

    /** Returns true if {@link #exec} was called and completed successfully; does not block. */
    boolean succeeded() {
      return succeeded.getCount() == 0;
    }
  }

  @Before
  public void setUp() throws Exception {
    testRoot = FileSystems.getNativeFileSystem().getPath(TestUtils.tmpDir());
    testRoot.deleteTreesBelow();
    outErr = new FileOutErr(testRoot.getRelative("stdout"), testRoot.getRelative("stderr"));
  }

  /**
   * Creates a new dynamic spawn strategy with different strategies for local and remote execution
   * and a default multi-threaded executor service.
   *
   * @param localStrategy the strategy for local execution
   * @param remoteStrategy the strategy for remote execution
   * @return the constructed dynamic strategy
   * @throws AbruptExitException if creating the strategy with the given parameters fails
   */
  private StrategyAndContext createSpawnStrategy(
      MockSpawnStrategy localStrategy, MockSpawnStrategy remoteStrategy)
      throws AbruptExitException {
    return createSpawnStrategyWithExecutor(
        localStrategy, remoteStrategy, Executors.newCachedThreadPool());
  }

  /**
   * Creates a new dynamic spawn strategy with different strategies for local, remote, and sandboxed
   * execution.
   *
   * <p>TODO(jmmv): This overload should not be necessary now that we do not special-case the
   * handling of sandboxed strategies any longer. Remove once the sandbox-specific flags are gone.
   *
   * @param localStrategy the default strategy for local execution
   * @param remoteStrategy the default strategy for remote execution
   * @param sandboxedStrategy the strategy to use when the mnemonic matches {@code testMnemonic}.
   * @return the constructed dynamic strategy
   * @throws AbruptExitException if creating the strategy with the given parameters fails
   */
  private StrategyAndContext createSpawnStrategy(
      MockSpawnStrategy localStrategy,
      MockSpawnStrategy remoteStrategy,
      @Nullable MockSpawnStrategy sandboxedStrategy)
      throws AbruptExitException {
    return createSpawnStrategyWithExecutor(
        localStrategy, remoteStrategy, sandboxedStrategy, Executors.newCachedThreadPool());
  }

  /**
   * Creates a new dynamic spawn strategy with different strategies for local and remote execution.
   *
   * @param localStrategy the strategy for local execution
   * @param remoteStrategy the strategy for remote execution
   * @param executorService the executor to pass to the dynamic strategy
   * @return the constructed dynamic strategy
   * @throws AbruptExitException if creating the strategy with the given parameters fails
   */
  private StrategyAndContext createSpawnStrategyWithExecutor(
      MockSpawnStrategy localStrategy,
      MockSpawnStrategy remoteStrategy,
      ExecutorService executorService)
      throws AbruptExitException {
    return createSpawnStrategyWithExecutor(localStrategy, remoteStrategy, null, executorService);
  }

  /**
   * Creates a new dynamic spawn strategy with different strategies for local, remote, and sandboxed
   * execution.
   *
   * <p>TODO(jmmv): This overload should not be necessary now that we do not special-case the
   * handling of sandboxed strategies any longer. Remove once the sandbox-specific flags are gone.
   *
   * @param localStrategy the default strategy for local execution
   * @param remoteStrategy the default strategy for remote execution
   * @param sandboxedStrategy the strategy to use when the mnemonic matches {@code testMnemonic}.
   * @param executorService the executor to pass to the dynamic strategy
   * @return the constructed dynamic strategy
   * @throws AbruptExitException if creating the strategy with the given parameters fails
   */
  private StrategyAndContext createSpawnStrategyWithExecutor(
      MockSpawnStrategy localStrategy,
      MockSpawnStrategy remoteStrategy,
      @Nullable MockSpawnStrategy sandboxedStrategy,
      ExecutorService executorService)
      throws AbruptExitException {
    ImmutableList.Builder<Map.Entry<String, List<String>>> dynamicLocalStrategies =
        ImmutableList.<Map.Entry<String, List<String>>>builder()
            .add(Maps.immutableEntry("", ImmutableList.of("mock-local")));
    ImmutableList.Builder<Map.Entry<String, List<String>>> dynamicRemoteStrategies =
        ImmutableList.<Map.Entry<String, List<String>>>builder()
            .add(Maps.immutableEntry("", ImmutableList.of("mock-remote")));

    if (sandboxedStrategy != null) {
      dynamicLocalStrategies.add(
          Maps.immutableEntry("testMnemonic", ImmutableList.of("mock-sandboxed")));
      dynamicRemoteStrategies.add(
          Maps.immutableEntry("testMnemonic", ImmutableList.of("mock-sandboxed")));
    }

    DynamicExecutionOptions options = new DynamicExecutionOptions();
    options.dynamicLocalStrategy = dynamicLocalStrategies.build();
    options.dynamicRemoteStrategy = dynamicRemoteStrategies.build();
    options.dynamicWorkerStrategy = "mock-local";
    options.internalSpawnScheduler = true;
    options.localExecutionDelay = 0;
    options.legacySpawnScheduler = legacyBehavior;

    checkState(executorServiceForCleanup == null);
    executorServiceForCleanup = executorService;

    ModuleActionContextRegistry.Builder moduleActionContextRegistryBuilder =
        ModuleActionContextRegistry.builder();
    SpawnStrategyRegistry.Builder spawnStrategyRegistryBuilder = SpawnStrategyRegistry.builder();

    spawnStrategyRegistryBuilder.registerStrategy(localStrategy, "mock-local");
    spawnStrategyRegistryBuilder.registerStrategy(remoteStrategy, "mock-remote");

    if (sandboxedStrategy != null) {
      spawnStrategyRegistryBuilder.registerStrategy(sandboxedStrategy, "mock-sandboxed");
    }

    DynamicExecutionModule dynamicExecutionModule = new DynamicExecutionModule(executorService);
    dynamicExecutionModule.registerSpawnStrategies(spawnStrategyRegistryBuilder, options);

    SpawnStrategyRegistry spawnStrategyRegistry = spawnStrategyRegistryBuilder.build();

    moduleActionContextRegistryBuilder.register(SpawnStrategyRegistry.class, spawnStrategyRegistry);
    moduleActionContextRegistryBuilder.register(
        DynamicStrategyRegistry.class, spawnStrategyRegistry);
    ModuleActionContextRegistry moduleActionContextRegistry =
        moduleActionContextRegistryBuilder.build();

    Executor executor =
        new BlazeExecutor(
            null,
            testRoot,
            null,
            null,
            OptionsParser.builder()
                .optionsClasses(ImmutableList.of(ExecutionOptions.class))
                .build(),
            moduleActionContextRegistry,
            spawnStrategyRegistry);

    ActionExecutionContext actionExecutionContext =
        ActionsTestUtil.createContext(
            /*executor=*/ executor,
            /*eventHandler=*/ null,
            actionKeyContext,
            outErr,
            testRoot,
            /*metadataHandler=*/ null,
            /*actionGraph=*/ null);

    List<? extends SpawnStrategy> dynamicStrategies =
        spawnStrategyRegistry.getStrategies(
            newCustomSpawn("RunDynamic", ImmutableMap.of()), event -> {});

    Optional<? extends SpawnStrategy> optionalContext =
        dynamicStrategies.stream()
            .filter(
                c -> c instanceof DynamicSpawnStrategy || c instanceof LegacyDynamicSpawnStrategy)
            .findAny();
    checkState(optionalContext.isPresent(), "Expected module to register a dynamic strategy");

    return new AutoValue_DynamicSpawnStrategyTest_StrategyAndContext(
        (SpawnStrategy) optionalContext.get(), actionExecutionContext);
  }

  private static class NullActionWithMnemonic extends NullAction {
    private final String mnemonic;

    private NullActionWithMnemonic(String mnemonic, List<Artifact> inputs, Artifact... outputs) {
      super(inputs, outputs);
      this.mnemonic = mnemonic;
    }

    @Override
    public String getMnemonic() {
      return mnemonic;
    }
  }

  @After
  public void tearDown() throws Exception {
    if (executorServiceForCleanup != null) {
      executorServiceForCleanup.shutdownNow();
    }
  }

  /** Constructs a new spawn with a custom mnemonic and execution info. */
  private Spawn newCustomSpawn(String mnemonic, ImmutableMap<String, String> executionInfo) {
    Artifact inputArtifact =
        ActionsTestUtil.createArtifact(
            ArtifactRoot.asSourceRoot(Root.fromPath(testRoot)), "input.txt");
    Artifact outputArtifact =
        ActionsTestUtil.createArtifact(
            ArtifactRoot.asSourceRoot(Root.fromPath(testRoot)), "output.txt");

    ActionExecutionMetadata action =
        new NullActionWithMnemonic(mnemonic, ImmutableList.of(inputArtifact), outputArtifact);
    return new BaseSpawn(
        ImmutableList.of(),
        ImmutableMap.of(),
        executionInfo,
        EmptyRunfilesSupplier.INSTANCE,
        action,
        ResourceSet.create(1, 0, 0));
  }

  /** Constructs a new spawn that can be run locally and remotely with arbitrary settings. */
  private Spawn newDynamicSpawn() {
    return newCustomSpawn("Null", ImmutableMap.of());
  }

  @Test
  public void nonRemotableSpawnRunsLocally() throws Exception {
    MockSpawnStrategy localStrategy = new MockSpawnStrategy("MockLocalSpawnStrategy");
    MockSpawnStrategy remoteStrategy = new MockSpawnStrategy("MockRemoteSpawnStrategy");
    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);

    Spawn spawn = newCustomSpawn("Null", ImmutableMap.of("local", "1"));
    strategyAndContext.exec(spawn);

    assertThat(localStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(localStrategy.succeeded()).isTrue();
    assertThat(remoteStrategy.getExecutedSpawn()).isNull();
    assertThat(remoteStrategy.succeeded()).isFalse();

    assertThat(outErr.outAsLatin1()).contains("output files written with MockLocalSpawnStrategy");
    assertThat(outErr.outAsLatin1()).doesNotContain("MockRemoteSpawnStrategy");
  }

  @Test
  public void localSpawnUsesStrategyByMnemonicWithWorkerFlagDisabled() throws Exception {
    MockSpawnStrategy localStrategy = new MockSpawnStrategy("MockLocalSpawnStrategy");
    MockSpawnStrategy remoteStrategy = new MockSpawnStrategy("MockRemoteSpawnStrategy");
    MockSpawnStrategy sandboxedStrategy = new MockSpawnStrategy("MockSandboxedSpawnStrategy");
    StrategyAndContext strategyAndContext =
        createSpawnStrategy(localStrategy, remoteStrategy, sandboxedStrategy);

    Spawn spawn = newCustomSpawn("testMnemonic", ImmutableMap.of("local", "1"));
    strategyAndContext.exec(spawn);

    assertThat(localStrategy.getExecutedSpawn()).isNull();
    assertThat(localStrategy.succeeded()).isFalse();
    assertThat(remoteStrategy.getExecutedSpawn()).isNull();
    assertThat(remoteStrategy.succeeded()).isFalse();
    assertThat(sandboxedStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(sandboxedStrategy.succeeded()).isTrue();

    assertThat(outErr.outAsLatin1())
        .contains("output files written with MockSandboxedSpawnStrategy");
    assertThat(outErr.outAsLatin1()).doesNotContain("MockLocalSpawnStrategy");
    assertThat(outErr.outAsLatin1()).doesNotContain("MockRemoteSpawnStrategy");
  }

  @Test
  public void remoteSpawnUsesStrategyByMnemonic() throws Exception {
    MockSpawnStrategy localStrategy = new MockSpawnStrategy("MockLocalSpawnStrategy");
    MockSpawnStrategy remoteStrategy = new MockSpawnStrategy("MockRemoteSpawnStrategy");
    MockSpawnStrategy sandboxedStrategy = new MockSpawnStrategy("MockSandboxedSpawnStrategy");
    StrategyAndContext strategyAndContext =
        createSpawnStrategy(localStrategy, remoteStrategy, sandboxedStrategy);

    Spawn spawn = newCustomSpawn("testMnemonic", ImmutableMap.of("remote", "1"));
    strategyAndContext.exec(spawn);

    assertThat(localStrategy.getExecutedSpawn()).isNull();
    assertThat(localStrategy.succeeded()).isFalse();
    assertThat(remoteStrategy.getExecutedSpawn()).isNull();
    assertThat(remoteStrategy.succeeded()).isFalse();
    assertThat(sandboxedStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(sandboxedStrategy.succeeded()).isTrue();

    assertThat(outErr.outAsLatin1())
        .contains("output files written with MockSandboxedSpawnStrategy");
    assertThat(outErr.outAsLatin1()).doesNotContain("MockLocalSpawnStrategy");
    assertThat(outErr.outAsLatin1()).doesNotContain("MockRemoteSpawnStrategy");
  }

  @Test
  public void actionSucceedsIfLocalExecutionSucceedsEvenIfRemoteFailsLater() throws Exception {
    CountDownLatch countDownLatch = new CountDownLatch(2);

    MockSpawnStrategy localStrategy =
        new MockSpawnStrategy(
            "MockLocalSpawnStrategy",
            (self, spawn, actionExecutionContext) -> countDownAndWait(countDownLatch),
            DoExec.NOTHING);

    MockSpawnStrategy remoteStrategy =
        new MockSpawnStrategy(
            "MockRemoteSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              Thread.sleep(2000);
              self.failExecution(actionExecutionContext);
            },
            DoExec.NOTHING);

    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);

    Spawn spawn = newDynamicSpawn();
    strategyAndContext.exec(spawn);

    assertThat(localStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(localStrategy.succeeded()).isTrue();
    assertThat(remoteStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(remoteStrategy.succeeded()).isFalse();

    assertThat(outErr.outAsLatin1()).contains("output files written with MockLocalSpawnStrategy");
    assertThat(outErr.outAsLatin1()).doesNotContain("MockRemoteSpawnStrategy");
  }

  @Test
  public void actionSucceedsIfRemoteExecutionSucceedsEvenIfLocalFailsLater() throws Exception {
    CountDownLatch countDownLatch = new CountDownLatch(2);

    MockSpawnStrategy localStrategy =
        new MockSpawnStrategy(
            "MockLocalSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              Thread.sleep(2000);
              self.failExecution(actionExecutionContext);
            },
            DoExec.NOTHING);

    MockSpawnStrategy remoteStrategy =
        new MockSpawnStrategy(
            "MockRemoteSpawnStrategy",
            (self, spawn, actionExecutionContext) -> countDownAndWait(countDownLatch),
            DoExec.NOTHING);

    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);

    Spawn spawn = newDynamicSpawn();
    strategyAndContext.exec(spawn);

    assertThat(localStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(localStrategy.succeeded()).isFalse();
    assertThat(remoteStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(remoteStrategy.succeeded()).isTrue();

    assertThat(outErr.outAsLatin1()).contains("output files written with MockRemoteSpawnStrategy");
    assertThat(outErr.outAsLatin1()).doesNotContain("MockLocalSpawnStrategy");
  }

  @Test
  public void actionFailsIfLocalFailsImmediatelyEvenIfRemoteSucceedsLater() throws Exception {
    CountDownLatch countDownLatch = new CountDownLatch(2);

    MockSpawnStrategy localStrategy =
        new MockSpawnStrategy(
            "MockLocalSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              self.failExecution(actionExecutionContext);
            },
            DoExec.NOTHING);

    MockSpawnStrategy remoteStrategy =
        new MockSpawnStrategy(
            "MockRemoteSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              Thread.sleep(2000);
            },
            DoExec.NOTHING);

    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);

    Spawn spawn = newDynamicSpawn();
    ExecException e = assertThrows(ExecException.class, () -> strategyAndContext.exec(spawn));
    assertThat(e).hasMessageThat().matches("MockLocalSpawnStrategy failed to execute the Spawn");

    assertThat(localStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(localStrategy.succeeded()).isFalse();
    assertThat(remoteStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(remoteStrategy.succeeded()).isFalse();

    assertThat(outErr.outAsLatin1()).contains("action failed with MockLocalSpawnStrategy");
    assertThat(outErr.outAsLatin1()).doesNotContain("MockRemoteSpawnStrategy");
  }

  @Test
  public void actionFailsIfRemoteFailsImmediatelyEvenIfLocalSucceedsLater() throws Exception {
    CountDownLatch countDownLatch = new CountDownLatch(2);

    MockSpawnStrategy localStrategy =
        new MockSpawnStrategy(
            "MockLocalSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              Thread.sleep(2000);
            },
            DoExec.NOTHING);

    MockSpawnStrategy remoteStrategy =
        new MockSpawnStrategy(
            "MockRemoteSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              self.failExecution(actionExecutionContext);
            },
            DoExec.NOTHING);

    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);

    Spawn spawn = newDynamicSpawn();
    ExecException e = assertThrows(ExecException.class, () -> strategyAndContext.exec(spawn));
    assertThat(e).hasMessageThat().matches("MockRemoteSpawnStrategy failed to execute the Spawn");

    assertThat(localStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(localStrategy.succeeded()).isFalse();
    assertThat(remoteStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(remoteStrategy.succeeded()).isFalse();

    assertThat(outErr.outAsLatin1()).contains("action failed with MockRemoteSpawnStrategy");
    assertThat(outErr.outAsLatin1()).doesNotContain("MockLocalSpawnStrategy");
  }

  @Test
  public void actionFailsIfLocalAndRemoteFail() throws Exception {
    CountDownLatch countDownLatch = new CountDownLatch(2);

    MockSpawnStrategy localStrategy =
        new MockSpawnStrategy(
            "MockLocalSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              self.failExecution(actionExecutionContext);
            },
            DoExec.NOTHING);

    MockSpawnStrategy remoteStrategy =
        new MockSpawnStrategy(
            "MockRemoteSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              self.failExecution(actionExecutionContext);
            },
            DoExec.NOTHING);

    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);

    Spawn spawn = newDynamicSpawn();
    ExecException e = assertThrows(ExecException.class, () -> strategyAndContext.exec(spawn));
    assertThat(e)
        .hasMessageThat()
        .matches("Mock(Local|Remote)SpawnStrategy failed to execute the Spawn");

    assertThat(localStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(localStrategy.succeeded()).isFalse();
    assertThat(remoteStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(remoteStrategy.succeeded()).isFalse();
  }

  @Test
  public void stopConcurrentSpawnsWaitForCompletion() throws Exception {
    if (legacyBehavior) {
      // The legacy spawn scheduler does not implement cross-cancellations of the two parallel
      // branches so this test makes no sense in that case.
      logger.atInfo().log("Skipping test");
      return;
    }

    CountDownLatch countDownLatch = new CountDownLatch(2);

    AtomicBoolean slowCleanupFinished = new AtomicBoolean(false);
    MockSpawnStrategy localStrategy =
        new MockSpawnStrategy(
            "MockLocalSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              try {
                countDownAndWait(countDownLatch);
                // Block indefinitely waiting for the remote branch to interrupt us.
                Thread.sleep(TestUtils.WAIT_TIMEOUT_MILLISECONDS);
                fail("Should have been interrupted");
              } catch (InterruptedException e) {
                // Wait for "long enough" hoping that the remoteStrategy will have enough time to
                // check the value of slowCleanupFinished before we finish this sleep, in case we
                // have a bug.
                Uninterruptibles.sleepUninterruptibly(5, TimeUnit.SECONDS);
                slowCleanupFinished.set(true);
              }
            },
            DoExec.NOTHING);

    MockSpawnStrategy remoteStrategy =
        new MockSpawnStrategy(
            "MockRemoteSpawnStrategy",
            (self, spawn, actionExecutionContext) -> countDownAndWait(countDownLatch),
            (self, spawn, actionExecutionContext) -> {
              // This runs after we have asked the local spawn to complete and, in theory, awaited
              // for InterruptedException to propagate. Make sure that's the case here by checking
              // that we did indeed wait for the slow process.
              if (!slowCleanupFinished.get()) {
                fail("Did not await for the other branch to do its cleanup");
              }
            });

    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);

    Spawn spawn = newDynamicSpawn();
    strategyAndContext.exec(spawn);

    assertThat(localStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(localStrategy.succeeded()).isFalse();
    assertThat(remoteStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(remoteStrategy.succeeded()).isTrue();
  }

  @Test
  public void noDeadlockWithSingleThreadedExecutor() throws Exception {
    MockSpawnStrategy localStrategy = new MockSpawnStrategy("MockLocalSpawnStrategy");
    MockSpawnStrategy remoteStrategy = new MockSpawnStrategy("MockRemoteSpawnStrategy");
    StrategyAndContext strategyAndContext =
        createSpawnStrategyWithExecutor(
            localStrategy, remoteStrategy, Executors.newSingleThreadExecutor());

    Spawn spawn = newDynamicSpawn();
    strategyAndContext.exec(spawn);

    assertThat(localStrategy.getExecutedSpawn()).isEqualTo(spawn);
    assertThat(localStrategy.succeeded()).isTrue();

    /*
     * The single-threaded executorService#invokeAny does not comply to the contract where
     * the callables are *always* called sequentially. In this case, both spawns will start
     * executing, but the local one will always succeed as it's the first to be called. The remote
     * one will then be cancelled, or is null if the local one completes before the remote one
     * starts.
     *
     * See the documentation of {@link BoundedExectorService#invokeAny(Collection)}, specifically:
     * "The following is less efficient (it goes on submitting tasks even if there is some task
     * already finished), but quite straight-forward.".
     */
    assertThat(remoteStrategy.getExecutedSpawn()).isAnyOf(spawn, null);
    assertThat(remoteStrategy.succeeded()).isFalse();
  }

  @Test
  public void interruptDuringExecutionDoesActuallyInterruptTheExecution() throws Exception {
    CountDownLatch countDownLatch = new CountDownLatch(2);

    MockSpawnStrategy localStrategy =
        new MockSpawnStrategy(
            "MockLocalSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              Thread.sleep(60000);
            },
            DoExec.NOTHING);

    MockSpawnStrategy remoteStrategy =
        new MockSpawnStrategy(
            "MockRemoteSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              countDownAndWait(countDownLatch);
              Thread.sleep(60000);
            },
            DoExec.NOTHING);

    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);

    TestThread testThread =
        new TestThread(
            () -> {
              try {
                Spawn spawn = newDynamicSpawn();
                strategyAndContext.exec(spawn);
              } catch (InterruptedException e) {
                // This is expected.
              }
            });
    testThread.start();
    countDownLatch.await(5, TimeUnit.SECONDS);
    testThread.interrupt();
    testThread.joinAndAssertState(5000);

    assertThat(outErr.getOutputPath().exists()).isFalse();
    assertThat(outErr.getErrorPath().exists()).isFalse();
  }

  /** Hook to validate the result of the strategy's execution. */
  @FunctionalInterface
  interface CheckExecResult {
    void check(@Nullable Exception e) throws Exception;
  }

  /**
   * Runs a test to check that both spawns finished under various conditions before the strategy's
   * {@code exec} method returns control.
   *
   * @param executionFails causes one of the branches in the execution to terminate with an
   *     execution exception
   * @param interruptThread causes the strategy's execution to be interrupted while it is waiting
   *     for its branches to complete
   * @param checkExecResult a lambda to validate the result of the execution. Receives null if the
   *     execution completed successfully, or else the raised exception.
   */
  private void assertThatStrategyWaitsForBothSpawnsToFinish(
      boolean executionFails, boolean interruptThread, CheckExecResult checkExecResult)
      throws Exception {
    if (!legacyBehavior) {
      // TODO(jmmv): I've spent *days* trying to make these tests work reliably with the new dynamic
      // spawn scheduler implementation but I keep encountering tricky race conditions everywhere. I
      // have strong reasons to believe that the races are due to inherent problems in these tests,
      // not in the actual DynamicSpawnScheduler implementation. So whatever. I'll revisit these
      // later as a new set of tests once I'm less tired^W^W^W the legacy spawn scheduler goes away.
      logger.atInfo().log("Skipping test");
      return;
    }

    AtomicBoolean stopLocal = new AtomicBoolean(false);
    CountDownLatch executionCanProceed = new CountDownLatch(2);
    CountDownLatch remoteDone = new CountDownLatch(1);

    MockSpawnStrategy localStrategy =
        new MockSpawnStrategy(
            "MockLocalSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              executionCanProceed.countDown();

              // We cannot use a synchronization primitive to block termination of this thread
              // because we expect to be interrupted by the remote strategy, and even in that case
              // we want to control exactly when this finishes. We could wait for and swallow the
              // interrupt before waiting again on a latch here... but swallowing the interrupt can
              // lead to race conditions.
              while (!stopLocal.get()) {
                Uninterruptibles.sleepUninterruptibly(1, TimeUnit.MILLISECONDS);
              }
              throw new InterruptedException("Local stopped");
            },
            DoExec.NOTHING);

    MockSpawnStrategy remoteStrategy =
        new MockSpawnStrategy(
            "MockRemoteSpawnStrategy",
            (self, spawn, actionExecutionContext) -> {
              try {
                // Wait until the local branch has started so that our completion causes it to be
                // interrupted in a known location.
                countDownAndWait(executionCanProceed);

                if (executionFails) {
                  self.failExecution(actionExecutionContext);
                  throw new AssertionError("Not reachable");
                }
              } finally {
                remoteDone.countDown();
              }
            },
            DoExec.NOTHING);

    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);
    TestThread testThread =
        new TestThread(
            () -> {
              try {
                Spawn spawn = newDynamicSpawn();
                strategyAndContext.exec(spawn);
                checkExecResult.check(null);
              } catch (Exception e) {
                checkExecResult.check(e);
              }
            });
    testThread.start();
    try {
      remoteDone.await();
      // At this point, the remote branch is done and the local branch is waiting until we allow it
      // to complete later on. This is necessary to let us assert the state of the thread's
      // liveliness.
      //
      // However, note that "done" just means that our DoExec hook for remoteStrategy finished.
      // Any exception raised from within it may still be propagating up, so the interrupt below
      // races with that (and thus an InterruptedException can "win" over our own exception). There
      // is no way to handle this condition in the test other than having to acknowledge that it may
      // happen.

      if (interruptThread) {
        testThread.interrupt();
      }

      // The thread running the exec via the strategy must still be alive regardless of our
      // interrupt request (because the local branch is stuck). Wait for a little bit to ensure
      // this is true; any multi-second wait should be sufficient to catch the majority of the
      // bugs.
      testThread.join(2000);
      assertThat(testThread.isAlive()).isTrue();
    } finally {
      // Unblocking the local branch allows the strategy to collect its result and then unblock the
      // thread.
      stopLocal.set(true);
      testThread.joinAndAssertState(TestUtils.WAIT_TIMEOUT_MILLISECONDS);
    }
  }

  @Test
  public void strategyWaitsForBothSpawnsToFinish() throws Exception {
    assertThatStrategyWaitsForBothSpawnsToFinish(
        /*executionFails=*/ false,
        /*interruptThread=*/ false,
        (e) -> {
          if (e != null) {
            throw new IllegalStateException("Expected exec to finish successfully", e);
          }
        });
  }

  @Test
  public void strategyWaitsForBothSpawnsToFinishEvenIfInterrupted() throws Exception {
    assertThatStrategyWaitsForBothSpawnsToFinish(
        /*executionFails=*/ false,
        /*interruptThread=*/ true,
        (e) -> {
          if (e == null) {
            fail("No exception raised");
          } else if (e instanceof InterruptedException) {
            assertThat(Thread.currentThread().isInterrupted()).isFalse();
          } else {
            throw e;
          }
        });
  }

  @Test
  public void strategyWaitsForBothSpawnsToFinishOnFailure() throws Exception {
    assertThatStrategyWaitsForBothSpawnsToFinish(
        /*executionFails=*/ true,
        /*interruptThread=*/ false,
        (e) -> {
          if (e == null) {
            fail("No exception raised");
          } else if (e instanceof ExecException) {
            assertThat(Thread.currentThread().isInterrupted()).isFalse();
          } else {
            throw e;
          }
        });
  }

  @Test
  public void strategyWaitsForBothSpawnsToFinishOnFailureEvenIfInterrupted() throws Exception {
    assertThatStrategyWaitsForBothSpawnsToFinish(
        /*executionFails=*/ true,
        /*interruptThread=*/ true,
        (e) -> {
          if (e == null) {
            fail("No exception raised");
          } else if (e instanceof InterruptedException) {
            // See comment in strategyWaitsForBothSpawnsToFinish regarding the race between the
            // exception we raise on failure and the interrupt. We have to handle this case even
            // though it is supposedly rare.
          } else if (e instanceof ExecException) {
            assertThat(Thread.currentThread().isInterrupted()).isTrue();
          } else {
            throw e;
          }
        });
  }

  private void assertThatStrategyPropagatesException(
      DoExec localExec, DoExec remoteExec, Exception expectedException) throws Exception {
    checkArgument(
        !(expectedException instanceof IllegalStateException),
        "Using an IllegalStateException for testing is fragile because we use that exception "
            + "internally in the DynamicSpawnScheduler and we cannot distinguish it from the "
            + "test's own exception");

    MockSpawnStrategy localStrategy =
        new MockSpawnStrategy("MockLocalSpawnStrategy", localExec, DoExec.NOTHING);
    MockSpawnStrategy remoteStrategy =
        new MockSpawnStrategy("MockRemoteSpawnStrategy", remoteExec, DoExec.NOTHING);
    StrategyAndContext strategyAndContext = createSpawnStrategy(localStrategy, remoteStrategy);

    Spawn spawn = newDynamicSpawn();
    Exception e = assertThrows(expectedException.getClass(), () -> strategyAndContext.exec(spawn));
    assertThat(e).hasMessageThat().matches(expectedException.getMessage());

    Spawn executedSpawn = localStrategy.getExecutedSpawn();
    executedSpawn = executedSpawn == null ? remoteStrategy.getExecutedSpawn() : executedSpawn;
    assertThat(executedSpawn).isEqualTo(spawn);
    assertThat(localStrategy.succeeded()).isFalse();
    assertThat(remoteStrategy.succeeded()).isFalse();
  }

  @Test
  public void strategyPropagatesFasterLocalException() throws Exception {
    RuntimeException e = new IllegalArgumentException("Local spawn execution exception");
    DoExec localExec =
        (self, spawn, actionExecutionContext) -> {
          throw e;
        };

    DoExec remoteExec =
        (self, spawn, actionExecutionContext) -> {
          Thread.sleep(TestUtils.WAIT_TIMEOUT_MILLISECONDS);
          throw new AssertionError("Not reachable");
        };

    assertThatStrategyPropagatesException(
        localExec, remoteExec, legacyBehavior ? new UserExecException(e) : e);
  }

  @Test
  public void strategyPropagatesFasterRemoteException() throws Exception {
    DoExec localExec =
        (self, spawn, actionExecutionContext) -> {
          Thread.sleep(TestUtils.WAIT_TIMEOUT_MILLISECONDS);
          throw new AssertionError("Not reachable");
        };

    RuntimeException e = new IllegalArgumentException("Remote spawn execution exception");
    DoExec remoteExec =
        (self, spawn, actionExecutionContext) -> {
          throw e;
        };

    assertThatStrategyPropagatesException(
        localExec, remoteExec, legacyBehavior ? new UserExecException(e) : e);
  }

  @AutoValue
  abstract static class StrategyAndContext {
    abstract SpawnStrategy strategy();

    abstract ActionExecutionContext context();

    void exec(Spawn spawn) throws ExecException, InterruptedException {
      strategy().exec(spawn, context());
    }
  }
}