src/test/java/com/google/devtools/build/lib/remote/RemoteRewoundActionSynchronizerTest.java

// Copyright 2026 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.remote;

import static com.google.common.truth.Truth.assertThat;
import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
import static org.junit.Assert.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.SettableFuture;
import com.google.devtools.build.lib.actions.Action;
import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
import com.google.devtools.build.lib.actions.ActionLookupData;
import com.google.devtools.build.lib.actions.ActionLookupKey;
import com.google.devtools.build.lib.actions.ActionLookupValue;
import com.google.devtools.build.lib.actions.ActionTemplate;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.actions.Artifact.DerivedArtifact;
import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
import com.google.devtools.build.lib.actions.Artifact.SpecialArtifactType;
import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
import com.google.devtools.build.lib.actions.ArtifactRoot;
import com.google.devtools.build.lib.actions.ArtifactRoot.RootType;
import com.google.devtools.build.lib.actions.FileArtifactValue;
import com.google.devtools.build.lib.actions.InputMetadataProvider;
import com.google.devtools.build.lib.actions.RunfilesArtifactValue;
import com.google.devtools.build.lib.actions.RunfilesTree;
import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
import com.google.devtools.build.lib.collect.nestedset.Order;
import com.google.devtools.build.lib.profiler.SilentCloseable;
import com.google.devtools.build.lib.remote.options.RemoteOutputsMode;
import com.google.devtools.build.lib.skyframe.ActionTemplateExpansionValue;
import com.google.devtools.build.lib.skyframe.ActionTemplateExpansionValue.ActionTemplateExpansionKey;
import com.google.devtools.build.lib.testutil.TestThread;
import com.google.devtools.build.lib.vfs.DigestHashFunction;
import com.google.devtools.build.lib.vfs.FileSystem;
import com.google.devtools.build.lib.vfs.PathFragment;
import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
import com.google.devtools.build.skyframe.WalkableGraph;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.InOrder;

/** Tests for {@link RemoteRewoundActionSynchronizer}. */
@RunWith(JUnit4.class)
public final class RemoteRewoundActionSynchronizerTest {
  private static final long DEADLOCK_TIMEOUT_MILLIS = 10_000;

  private RemoteActionInputFetcher actionInputFetcher;
  private WalkableGraph graph;
  private RemoteRewoundActionSynchronizer synchronizer;

  @Before
  public void setUp() {
    actionInputFetcher = mock(RemoteActionInputFetcher.class);
    graph = mock(WalkableGraph.class);
    synchronizer = new RemoteRewoundActionSynchronizer(actionInputFetcher, graph);
  }

  @Test
  public void rewind_cancelsAllRegisteredTasksOnce() throws Exception {
    Action action = newAction();
    var first = mock(RemoteRewoundActionSynchronizer.Cancellable.class);
    var second = mock(RemoteRewoundActionSynchronizer.Cancellable.class);
    var unusedFirst = synchronizer.registerOutputUploadTask(action, first);
    var unusedSecond = synchronizer.registerOutputUploadTask(action, second);

    rewind(action);
    rewind(action);

    verify(first, times(1)).requestCancellation();
    verify(first, times(1)).awaitCompletion();
    verify(second, times(1)).requestCancellation();
    verify(second, times(1)).awaitCompletion();
  }

  @Test
  public void rewind_interruptedWhileAwaiting_cancelsAndAwaitsEveryTask() throws Exception {
    Action action = newAction();
    var first = mock(RemoteRewoundActionSynchronizer.Cancellable.class);
    var second = mock(RemoteRewoundActionSynchronizer.Cancellable.class);
    doThrow(new InterruptedException()).doNothing().when(first).awaitCompletion();
    var unusedFirst = synchronizer.registerOutputUploadTask(action, first);
    var unusedSecond = synchronizer.registerOutputUploadTask(action, second);

    assertThrows(InterruptedException.class, () -> rewind(action));

    InOrder order = inOrder(first, second);
    order.verify(first).requestCancellation();
    order.verify(second).requestCancellation();
    order.verify(first, times(2)).awaitCompletion();
    order.verify(second).awaitCompletion();
    verify(actionInputFetcher, never()).handleRewoundActionOutputs(any());
  }

  @Test
  public void unregisterHandle_preventsCancellation() throws Exception {
    Action action = newAction();
    var task = mock(RemoteRewoundActionSynchronizer.Cancellable.class);
    Runnable unregister = synchronizer.registerOutputUploadTask(action, task);

    unregister.run();
    rewind(action);

    verify(task, never()).requestCancellation();
    verify(task, never()).awaitCompletion();
    assertThat(synchronizer.hasRegisteredOutputUploadTasks(action)).isFalse();
  }

  @Test
  public void predecessorUnregisterHandle_doesNotRemoveReplacementTask() throws Exception {
    Action action = newAction();
    var predecessor = mock(RemoteRewoundActionSynchronizer.Cancellable.class);
    Runnable unregisterPredecessor = synchronizer.registerOutputUploadTask(action, predecessor);
    rewind(action);
    var replacement = mock(RemoteRewoundActionSynchronizer.Cancellable.class);
    var unused = synchronizer.registerOutputUploadTask(action, replacement);

    unregisterPredecessor.run();
    rewind(action);

    verify(predecessor, times(1)).requestCancellation();
    verify(predecessor, times(1)).awaitCompletion();
    verify(replacement, times(1)).requestCancellation();
    verify(replacement, times(1)).awaitCompletion();
  }

  @Test
  public void unregisterHandle_usesIdentityComparison() throws Exception {
    Action action = newAction();
    var first = new EqualCancellable();
    var second = new EqualCancellable();
    Runnable unregisterFirst = synchronizer.registerOutputUploadTask(action, first);
    var unused = synchronizer.registerOutputUploadTask(action, second);

    unregisterFirst.run();
    rewind(action);

    assertThat(first.cancellations.get()).isEqualTo(0);
    assertThat(second.cancellations.get()).isEqualTo(1);
  }

  /**
   * Shared actions are equivalent actions of different owners generating the same outputs, which
   * SkyframeActionExecutor executes as one. A consumer of one owner's output excludes the rewinding
   * of the other owner's action, which deletes the same file, and that rewinding cancels the
   * uploads registered by either owner.
   */
  @Test
  public void sharedAction_rewoundUnderOtherOwner_waitsForConsumerAndCancelsUpload()
      throws Exception {
    FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256);
    ArtifactRoot root = ArtifactRoot.asDerivedRoot(fs.getPath("/exec"), RootType.OUTPUT, "out");
    var owner = ActionsTestUtil.NULL_ARTIFACT_OWNER;
    var otherOwner = ActionsTestUtil.YET_ANOTHER_NULL_ARTIFACT_OWNER;
    PathFragment sharedExecPath = root.getExecPath().getRelative("shared");
    DerivedArtifact output = DerivedArtifact.create(root, sharedExecPath, owner);
    output.setGeneratingActionKey(ActionLookupData.create(owner, 0));
    Action action = newAction(ImmutableList.of(output), ImmutableList.of());
    DerivedArtifact otherOutput = DerivedArtifact.create(root, sharedExecPath, otherOwner);
    otherOutput.setGeneratingActionKey(ActionLookupData.create(otherOwner, 0));
    Action otherAction = newAction(ImmutableList.of(otherOutput), ImmutableList.of());
    DerivedArtifact consumerOutput =
        (DerivedArtifact) ActionsTestUtil.createArtifact(root, "consumer.out");
    consumerOutput.setGeneratingActionKey(ActionLookupData.create(owner, 1));
    Action consumer = newAction(ImmutableList.of(consumerOutput), ImmutableList.of(output));
    mockActions(owner, ImmutableList.of(action, consumer));
    mockActions(otherOwner, ImmutableList.of(otherAction));
    var upload = mock(RemoteRewoundActionSynchronizer.Cancellable.class);
    var unused = synchronizer.registerOutputUploadTask(action, upload);

    // Switch to the fine locks with an unrelated rewound action first: the coarse lock would
    // exclude every rewound action regardless of which keys guard the consumed output.
    rewind(newAction());

    var otherPreparation = new TestThread(() -> rewind(otherAction));
    try (SilentCloseable execution =
        synchronizer.enterActionExecution(
            consumer, /* wasRewound= */ false, mock(InputMetadataProvider.class))) {
      otherPreparation.start();
      waitUntilBlocked(otherPreparation);
    }
    otherPreparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
    verify(upload).requestCancellation();
    verify(upload).awaitCompletion();
    assertThat(synchronizer.hasRegisteredOutputUploadTasks(action)).isFalse();
  }

  /**
   * Input discovery may read scheduling dependencies (e.g. headers reachable via {@code #include})
   * that are not declared inputs, so their producers must not be rewound while it is running.
   */
  @Test
  public void inputDiscovery_locksProducersOfSchedulingDependencies() throws Exception {
    FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256);
    ArtifactRoot root = ArtifactRoot.asDerivedRoot(fs.getPath("/exec"), RootType.OUTPUT, "out");
    var owner = ActionsTestUtil.NULL_ARTIFACT_OWNER;
    DerivedArtifact header = (DerivedArtifact) ActionsTestUtil.createArtifact(root, "header.h");
    header.setGeneratingActionKey(ActionLookupData.create(owner, 0));
    Action headerProducer = newAction(ImmutableList.of(header), ImmutableList.of());
    DerivedArtifact source = (DerivedArtifact) ActionsTestUtil.createArtifact(root, "source.cc");
    source.setGeneratingActionKey(ActionLookupData.create(owner, 1));
    Action sourceProducer = newAction(ImmutableList.of(source), ImmutableList.of());
    DerivedArtifact object = (DerivedArtifact) ActionsTestUtil.createArtifact(root, "object.o");
    object.setGeneratingActionKey(ActionLookupData.create(owner, 2));
    Action compile =
        newAction(
            ImmutableList.of(object),
            /* inputs= */ ImmutableList.of(source),
            /* schedulingDependencies= */ ImmutableList.of(header));
    mockActions(owner, ImmutableList.of(headerProducer, sourceProducer, compile));
    InputMetadataProvider metadataProvider = mock(InputMetadataProvider.class);

    // Switch to the fine locks with an unrelated rewound action first: the coarse lock would
    // exclude every rewound action regardless of which keys guard the discovery.
    rewind(newAction());

    var headerProducerPreparation = new TestThread(() -> rewind(headerProducer));
    try (SilentCloseable discovery =
        synchronizer.enterInputDiscovery(
            compile, NestedSetBuilder.create(Order.STABLE_ORDER, source), metadataProvider)) {
      headerProducerPreparation.start();
      waitUntilBlocked(headerProducerPreparation);
    }
    headerProducerPreparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
  }

  /**
   * An action consuming a runfiles tree is guarded by the keys of the actions generating the
   * artifacts the tree contains, which are taken from its metadata rather than from all runfiles
   * trees known to the metadata provider. The action generating the runfiles tree itself isn't
   * excluded, as it doesn't write to disk.
   */
  @Test
  public void runfilesTreeConsumer_excludesOnlyProducersOfItsRunfiles() throws Exception {
    FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256);
    ArtifactRoot root = ArtifactRoot.asDerivedRoot(fs.getPath("/exec"), RootType.OUTPUT, "out");
    var owner = ActionsTestUtil.NULL_ARTIFACT_OWNER;
    DerivedArtifact file = (DerivedArtifact) ActionsTestUtil.createArtifact(root, "file");
    file.setGeneratingActionKey(ActionLookupData.create(owner, 0));
    Action producer = newAction(ImmutableList.of(file), ImmutableList.of());
    DerivedArtifact unrelatedFile =
        (DerivedArtifact) ActionsTestUtil.createArtifact(root, "unrelated");
    unrelatedFile.setGeneratingActionKey(ActionLookupData.create(owner, 1));
    Action unrelatedProducer = newAction(ImmutableList.of(unrelatedFile), ImmutableList.of());
    SpecialArtifact runfiles = ActionsTestUtil.createRunfilesArtifact(root, "out/runfiles");
    runfiles.setGeneratingActionKey(ActionLookupData.create(owner, 2));
    Action runfilesAction = newAction(ImmutableList.of(runfiles), ImmutableList.of(file));
    DerivedArtifact consumerOutput =
        (DerivedArtifact) ActionsTestUtil.createArtifact(root, "consumer.out");
    consumerOutput.setGeneratingActionKey(ActionLookupData.create(owner, 3));
    Action consumer = newAction(ImmutableList.of(consumerOutput), ImmutableList.of(runfiles));
    mockActions(owner, ImmutableList.of(producer, unrelatedProducer, runfilesAction, consumer));

    RunfilesTree runfilesTree = mock(RunfilesTree.class);
    when(runfilesTree.getArtifacts()).thenReturn(NestedSetBuilder.create(Order.STABLE_ORDER, file));
    RunfilesTree unrelatedRunfilesTree = mock(RunfilesTree.class);
    when(unrelatedRunfilesTree.getArtifacts())
        .thenReturn(NestedSetBuilder.create(Order.STABLE_ORDER, unrelatedFile));
    InputMetadataProvider metadataProvider = mock(InputMetadataProvider.class);
    when(metadataProvider.getRunfilesTrees())
        .thenReturn(ImmutableList.of(runfilesTree, unrelatedRunfilesTree));
    var runfilesValue =
        new RunfilesArtifactValue(
            runfilesTree,
            ImmutableList.of(file),
            ImmutableList.of(FileArtifactValue.createForNormalFile(new byte[32], null, 0)),
            ImmutableList.of(),
            ImmutableList.of(),
            ImmutableList.of(),
            ImmutableList.of());
    when(metadataProvider.getRunfilesMetadata(runfiles)).thenReturn(runfilesValue);

    // Switch to the fine locks with an unrelated rewound action first: the coarse lock would
    // exclude every rewound action regardless of which keys guard the runfiles tree.
    rewind(newAction());

    var preparation = new TestThread(() -> rewind(producer));
    var unrelatedPreparation = new TestThread(() -> rewind(unrelatedProducer));
    var runfilesPreparation = new TestThread(() -> rewind(runfilesAction));
    try (SilentCloseable execution =
        synchronizer.enterActionExecution(consumer, /* wasRewound= */ false, metadataProvider)) {
      preparation.start();
      waitUntilBlocked(preparation);
      unrelatedPreparation.start();
      unrelatedPreparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
      runfilesPreparation.start();
      runfilesPreparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
    }
    preparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
  }

  /**
   * Runfiles trees are hidden top-level outputs and thus only reach {@link
   * RemoteImportantOutputHandler} through the metadata provider, but the rewinding of the actions
   * generating their artifacts must still be excluded while the handler processes them.
   */
  @Test
  public void outputProcessing_waitsForRewoundRunfilesProducer() throws Exception {
    FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256);
    ArtifactRoot root = ArtifactRoot.asDerivedRoot(fs.getPath("/exec"), RootType.OUTPUT, "out");
    var owner = ActionsTestUtil.NULL_ARTIFACT_OWNER;
    DerivedArtifact runfile = (DerivedArtifact) ActionsTestUtil.createArtifact(root, "runfile");
    runfile.setGeneratingActionKey(ActionLookupData.create(owner, 0));
    Action producer = newAction(ImmutableList.of(runfile), ImmutableList.of());
    mockActions(owner, ImmutableList.of(producer));
    InputMetadataProvider metadataProvider =
        mockTopLevelMetadataProvider(
            /* runfiles= */ ImmutableList.of(runfile),
            /* remoteArtifacts= */ ImmutableList.of(runfile));
    when(actionInputFetcher.prefetchFiles(any(), any(), any(), any(), any(), any()))
        .thenReturn(immediateVoidFuture());
    RemoteImportantOutputHandler handler = newImportantOutputHandler(runfile);

    // Switch to the fine locks with an unrelated rewound action first: the coarse lock would
    // exclude every rewound action regardless of which keys guard the processed outputs.
    rewind(newAction());

    var processing =
        new TestThread(
            () ->
                assertThat(
                        handler
                            .processOutputsAndGetLostArtifacts(
                                /* importantOutputs= */ ImmutableList.of(), metadataProvider)
                            .isEmpty())
                    .isTrue());
    try (SilentCloseable preparation =
        synchronizer.enterActionPreparation(producer, /* wasRewound= */ true)) {
      processing.start();
      waitUntilBlocked(processing);
    }
    processing.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
  }

  @Test
  public void outputProcessing_excludesOnlyRewindingOfProducersOfOutputsAndRunfiles()
      throws Exception {
    FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256);
    ArtifactRoot root = ArtifactRoot.asDerivedRoot(fs.getPath("/exec"), RootType.OUTPUT, "out");
    var owner = ActionsTestUtil.NULL_ARTIFACT_OWNER;
    DerivedArtifact importantOutput =
        (DerivedArtifact) ActionsTestUtil.createArtifact(root, "important");
    importantOutput.setGeneratingActionKey(ActionLookupData.create(owner, 0));
    Action importantOutputProducer =
        newAction(ImmutableList.of(importantOutput), ImmutableList.of());
    DerivedArtifact runfile = (DerivedArtifact) ActionsTestUtil.createArtifact(root, "runfile");
    runfile.setGeneratingActionKey(ActionLookupData.create(owner, 1));
    Action runfileProducer = newAction(ImmutableList.of(runfile), ImmutableList.of());
    DerivedArtifact unrelated = (DerivedArtifact) ActionsTestUtil.createArtifact(root, "unrelated");
    unrelated.setGeneratingActionKey(ActionLookupData.create(owner, 2));
    Action unrelatedProducer = newAction(ImmutableList.of(unrelated), ImmutableList.of());
    mockActions(
        owner, ImmutableList.of(importantOutputProducer, runfileProducer, unrelatedProducer));
    InputMetadataProvider metadataProvider =
        mockTopLevelMetadataProvider(
            /* runfiles= */ ImmutableList.of(runfile),
            /* remoteArtifacts= */ ImmutableList.of(importantOutput, runfile));
    // The processing holds its locks while the download is pending.
    SettableFuture<Void> download = SettableFuture.create();
    when(actionInputFetcher.prefetchFiles(any(), any(), any(), any(), any(), any()))
        .thenReturn(download);
    RemoteImportantOutputHandler handler = newImportantOutputHandler(importantOutput, runfile);

    // Switch to the fine locks with an unrelated rewound action first: the coarse lock would
    // exclude every rewound action regardless of which keys guard the processed outputs.
    rewind(newAction());

    var processing =
        new TestThread(
            () ->
                assertThat(
                        handler
                            .processOutputsAndGetLostArtifacts(
                                ImmutableList.of(importantOutput), metadataProvider)
                            .isEmpty())
                    .isTrue());
    processing.start();
    waitUntilBlocked(processing);

    var importantOutputPreparation = new TestThread(() -> rewind(importantOutputProducer));
    importantOutputPreparation.start();
    waitUntilBlocked(importantOutputPreparation);
    var runfilePreparation = new TestThread(() -> rewind(runfileProducer));
    runfilePreparation.start();
    waitUntilBlocked(runfilePreparation);
    var unrelatedPreparation = new TestThread(() -> rewind(unrelatedProducer));
    unrelatedPreparation.start();
    unrelatedPreparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);

    download.set(null);
    processing.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
    importantOutputPreparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
    runfilePreparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
  }

  /**
   * A consumer of a tree artifact declared by an {@link ActionTemplate} excludes the rewinding of
   * every expanded action populating it, including one whose only output in the tree is an empty
   * subdirectory and which thus doesn't generate any of the tree's files.
   */
  @Test
  public void expandedActionRewound_emptySubdirectoryProducer_waitsForTreeConsumer()
      throws Exception {
    FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256);
    ArtifactRoot root = ArtifactRoot.asDerivedRoot(fs.getPath("/exec"), RootType.OUTPUT, "out");
    var owner = ActionsTestUtil.NULL_ARTIFACT_OWNER;
    SpecialArtifact tree = newTreeArtifact(root, "tree", ActionLookupData.create(owner, 0));
    ActionTemplateExpansionKey expansion = ActionTemplateExpansionValue.key(owner, 0);
    SpecialArtifact subdirectory =
        SpecialArtifact.createSubTreeArtifact(tree, PathFragment.create("empty"), expansion);
    subdirectory.setGeneratingActionKey(ActionLookupData.create(expansion, 0));
    Action producer = newAction(ImmutableList.of(subdirectory), ImmutableList.of());
    DerivedArtifact consumerOutput =
        (DerivedArtifact) ActionsTestUtil.createArtifact(root, "consumer.out");
    consumerOutput.setGeneratingActionKey(ActionLookupData.create(owner, 1));
    Action consumer = newAction(ImmutableList.of(consumerOutput), ImmutableList.of(tree));

    InputMetadataProvider metadataProvider = mock(InputMetadataProvider.class);
    mockTemplateExpansion(owner, ImmutableList.of(producer));

    // Switch to the fine locks with an unrelated rewound action first: the coarse lock would
    // exclude the producer regardless of which keys guard the tree.
    rewind(newAction());

    var preparation = new TestThread(() -> rewind(producer));
    try (SilentCloseable execution =
        synchronizer.enterActionExecution(consumer, /* wasRewound= */ false, metadataProvider)) {
      preparation.start();
      waitUntilBlocked(preparation);
    }
    preparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
  }

  /**
   * A consumer of one output tree of an {@link ActionTemplate} doesn't depend on expanded actions
   * that only populate another output tree of the same template and thus must not wait for their
   * rewinding, whereas it does exclude the rewinding of the actions populating its own tree.
   */
  @Test
  public void expandedActionRewound_treeConsumer_excludesOnlyProducersOfItsTree() throws Exception {
    FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256);
    ArtifactRoot root = ArtifactRoot.asDerivedRoot(fs.getPath("/exec"), RootType.OUTPUT, "out");
    var owner = ActionsTestUtil.NULL_ARTIFACT_OWNER;
    var templateKey = ActionLookupData.create(owner, 0);
    ActionTemplateExpansionKey expansion = ActionTemplateExpansionValue.key(owner, 0);
    SpecialArtifact firstTree = newTreeArtifact(root, "first", templateKey);
    SpecialArtifact secondTree = newTreeArtifact(root, "second", templateKey);
    TreeFileArtifact firstFile =
        TreeFileArtifact.createTemplateExpansionOutput(firstTree, "file", expansion);
    firstFile.setGeneratingActionKey(ActionLookupData.create(expansion, 0));
    TreeFileArtifact secondFile =
        TreeFileArtifact.createTemplateExpansionOutput(secondTree, "file", expansion);
    secondFile.setGeneratingActionKey(ActionLookupData.create(expansion, 1));
    Action firstProducer = newAction(ImmutableList.of(firstFile), ImmutableList.of());
    Action secondProducer = newAction(ImmutableList.of(secondFile), ImmutableList.of());
    DerivedArtifact consumerOutput =
        (DerivedArtifact) ActionsTestUtil.createArtifact(root, "consumer.out");
    consumerOutput.setGeneratingActionKey(ActionLookupData.create(owner, 1));
    Action consumer = newAction(ImmutableList.of(consumerOutput), ImmutableList.of(firstTree));

    InputMetadataProvider metadataProvider = mock(InputMetadataProvider.class);
    mockTemplateExpansion(owner, ImmutableList.of(firstProducer, secondProducer));

    // Switch to the fine locks with an unrelated rewound action first: the coarse lock would
    // exclude the consumer regardless of which keys guard the trees.
    rewind(newAction());

    try (SilentCloseable secondPreparation =
        synchronizer.enterActionPreparation(secondProducer, /* wasRewound= */ true)) {
      // The consumer of the first tree is admitted while the producer of the second tree holds its
      // write lock.
      var consumerExecution =
          new TestThread(
              () -> {
                try (SilentCloseable execution =
                    synchronizer.enterActionExecution(
                        consumer, /* wasRewound= */ false, metadataProvider)) {}
              });
      consumerExecution.start();
      consumerExecution.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);

      // The producer of the first tree still waits for the consumer.
      var firstPreparation = new TestThread(() -> rewind(firstProducer));
      try (SilentCloseable execution =
          synchronizer.enterActionExecution(consumer, /* wasRewound= */ false, metadataProvider)) {
        firstPreparation.start();
        waitUntilBlocked(firstPreparation);
      }
      firstPreparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
    }
  }

  /**
   * Regression test for a lock cycle between the rewound expanded action P of an action template, a
   * rewound consumer C of the tree artifact P populates, and an action D expanded from a downstream
   * template that consumes both P's file and C's output:
   *
   * <ol>
   *   <li>D holds the read lock of P and waits for the read lock of C, which C holds for writing.
   *   <li>P waits for the write lock of its own key, which D holds for reading.
   *   <li>C, which depends on P but not on D, requests the read lock of P.
   * </ol>
   *
   * <p>C must be admitted even though P is waiting for the same lock as a writer. Any lock that
   * queues readers behind waiting writers would make C wait for P and close the cycle. {@link
   * java.util.concurrent.locks.ReentrantReadWriteLock} does so in both fairness modes and {@link
   * java.util.concurrent.locks.StampedLock} does so in {@code readLock}, but not in {@code
   * readLockInterruptibly}, which only waits for a writer that holds the lock.
   */
  @Test
  public void expandedActionRewound_treeConsumerNotQueuedBehindWaitingProducer() throws Exception {
    FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256);
    ArtifactRoot root = ArtifactRoot.asDerivedRoot(fs.getPath("/exec"), RootType.OUTPUT, "out");
    var owner = ActionsTestUtil.NULL_ARTIFACT_OWNER;

    // The upstream template (action 0 of the owner) declares a tree artifact that its single
    // expanded action P populates with a file.
    SpecialArtifact upstreamTree =
        newTreeArtifact(root, "upstream", ActionLookupData.create(owner, 0));
    ActionTemplateExpansionKey upstreamExpansion = ActionTemplateExpansionValue.key(owner, 0);
    TreeFileArtifact upstreamFile =
        TreeFileArtifact.createTemplateExpansionOutput(upstreamTree, "file", upstreamExpansion);
    upstreamFile.setGeneratingActionKey(ActionLookupData.create(upstreamExpansion, 0));
    Action upstreamAction = newAction(ImmutableList.of(upstreamFile), ImmutableList.of());

    // The regular action C (action 1 of the owner) consumes the whole upstream tree artifact.
    DerivedArtifact treeConsumerOutput =
        (DerivedArtifact) ActionsTestUtil.createArtifact(root, "tree_consumer.out");
    treeConsumerOutput.setGeneratingActionKey(ActionLookupData.create(owner, 1));
    Action treeConsumer =
        newAction(ImmutableList.of(treeConsumerOutput), ImmutableList.of(upstreamTree));

    // The downstream template (action 2 of the owner) is expanded over the upstream tree artifact.
    // Its expanded action D consumes the file generated by P as well as the output of C and
    // acquires their read locks in this order.
    SpecialArtifact downstreamTree =
        newTreeArtifact(root, "downstream", ActionLookupData.create(owner, 2));
    ActionTemplateExpansionKey downstreamExpansion = ActionTemplateExpansionValue.key(owner, 2);
    TreeFileArtifact downstreamFile =
        TreeFileArtifact.createTemplateExpansionOutput(downstreamTree, "file", downstreamExpansion);
    downstreamFile.setGeneratingActionKey(ActionLookupData.create(downstreamExpansion, 0));
    Action downstreamAction =
        newAction(
            ImmutableList.of(downstreamFile), ImmutableList.of(upstreamFile, treeConsumerOutput));

    InputMetadataProvider metadataProvider = mock(InputMetadataProvider.class);
    mockActions(
        owner,
        ImmutableList.of(mock(ActionTemplate.class), treeConsumer, mock(ActionTemplate.class)));
    mockTemplateExpansion(owner, /* actionIndex= */ 0, ImmutableList.of(upstreamAction));

    // C is rewound and prepares for its re-execution, which makes it hold the write lock guarding
    // its output until the end of its execution.
    SilentCloseable treeConsumerPreparation =
        synchronizer.enterActionPreparation(treeConsumer, /* wasRewound= */ true);
    // D enters execution, acquires the read lock guarding the upstream file and then blocks on the
    // read lock guarding the output of C.
    var downstreamExecution =
        new TestThread(
            () -> {
              try (SilentCloseable execution =
                  synchronizer.enterActionExecution(
                      downstreamAction, /* wasRewound= */ false, metadataProvider)) {}
            });
    downstreamExecution.start();
    waitUntilBlocked(downstreamExecution);
    // P is rewound and prepares for its re-execution, which blocks on the read lock held by D.
    var upstreamPreparation = new TestThread(() -> rewind(upstreamAction));
    upstreamPreparation.start();
    waitUntilBlocked(upstreamPreparation);
    // C enters execution, which requires the read lock guarding the upstream tree artifact. It must
    // not wait for P, which waits for D, which waits for C.
    var treeConsumerExecution =
        new TestThread(
            () -> {
              try (treeConsumerPreparation;
                  SilentCloseable execution =
                      synchronizer.enterActionExecution(
                          treeConsumer, /* wasRewound= */ true, metadataProvider)) {}
            });
    treeConsumerExecution.start();

    treeConsumerExecution.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
    downstreamExecution.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
    upstreamPreparation.joinAndAssertState(DEADLOCK_TIMEOUT_MILLIS);
  }

  /**
   * Makes the given owner's only action an {@link ActionTemplate} whose expansion consists of the
   * given actions.
   */
  private void mockTemplateExpansion(ActionLookupKey owner, ImmutableList<Action> expandedActions)
      throws InterruptedException {
    mockActions(owner, ImmutableList.of(mock(ActionTemplate.class)));
    mockTemplateExpansion(owner, /* actionIndex= */ 0, expandedActions);
  }

  /** Makes the given actions the expansion of the given owner's template at the given index. */
  private void mockTemplateExpansion(
      ActionLookupKey owner, int actionIndex, ImmutableList<Action> expandedActions)
      throws InterruptedException {
    when(graph.getValue(ActionTemplateExpansionValue.key(owner, actionIndex)))
        .thenReturn(new ActionTemplateExpansionValue(ImmutableList.copyOf(expandedActions)));
  }

  /** Makes the given actions, in order, the actions of the given owner. */
  private void mockActions(
      ActionLookupKey owner, ImmutableList<? extends ActionAnalysisMetadata> actions)
      throws InterruptedException {
    ActionLookupValue ownerValue = mock(ActionLookupValue.class);
    when(ownerValue.getActions()).thenReturn(ImmutableList.<ActionAnalysisMetadata>copyOf(actions));
    when(graph.getValue(owner)).thenReturn(ownerValue);
  }

  /**
   * Returns a metadata provider that, just like the one {@code CompletionFunction} passes to the
   * {@link RemoteImportantOutputHandler}, provides a runfiles tree containing the given runfiles
   * and remote metadata for the given artifacts.
   */
  private static InputMetadataProvider mockTopLevelMetadataProvider(
      ImmutableList<Artifact> runfiles, ImmutableList<Artifact> remoteArtifacts) throws Exception {
    RunfilesTree runfilesTree = mock(RunfilesTree.class);
    when(runfilesTree.getArtifacts())
        .thenReturn(NestedSetBuilder.wrap(Order.STABLE_ORDER, runfiles));
    InputMetadataProvider metadataProvider = mock(InputMetadataProvider.class);
    when(metadataProvider.getRunfilesTrees()).thenReturn(ImmutableList.of(runfilesTree));
    FileArtifactValue remoteMetadata = FileArtifactValue.createForRemoteFile(new byte[32], 1, 1);
    for (Artifact artifact : remoteArtifacts) {
      when(metadataProvider.getInputMetadata(artifact)).thenReturn(remoteMetadata);
    }
    return metadataProvider;
  }

  /**
   * Returns a handler for a build that downloads the given top-level outputs, which are registered
   * with the output checker just like the outputs and runfiles of top-level targets are after their
   * analysis.
   */
  private RemoteImportantOutputHandler newImportantOutputHandler(Artifact... topLevelOutputs) {
    var remoteOutputChecker =
        new RemoteOutputChecker("build", RemoteOutputsMode.TOPLEVEL, ImmutableList.of());
    for (Artifact output : topLevelOutputs) {
      remoteOutputChecker.addOutputToDownload(output);
    }
    return new RemoteImportantOutputHandler(
        graph, remoteOutputChecker, actionInputFetcher, synchronizer);
  }

  @SuppressWarnings("ThreadPriorityCheck")
  private static void waitUntilBlocked(Thread thread) {
    Thread.State state;
    while ((state = thread.getState()) != Thread.State.WAITING) {
      assertThat(state).isNotEqualTo(Thread.State.TERMINATED);
      Thread.yield();
    }
  }

  private void rewind(Action action) throws Exception {
    try (SilentCloseable ignored = synchronizer.enterActionPreparation(action, true)) {
      // Cancellation happens while entering preparation.
    }
  }

  private static Action newAction() {
    FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256);
    ArtifactRoot outputRoot =
        ArtifactRoot.asDerivedRoot(fs.getPath("/exec"), RootType.OUTPUT, "out");
    DerivedArtifact output = (DerivedArtifact) ActionsTestUtil.createArtifact(outputRoot, "output");
    output.setGeneratingActionKey(ActionsTestUtil.NULL_ACTION_LOOKUP_DATA);
    return newAction(ImmutableList.of(output), ImmutableList.of());
  }

  private static Action newAction(
      ImmutableList<? extends Artifact> outputs, ImmutableList<? extends Artifact> inputs) {
    return newAction(outputs, inputs, /* schedulingDependencies= */ ImmutableList.of());
  }

  private static Action newAction(
      ImmutableList<? extends Artifact> outputs,
      ImmutableList<? extends Artifact> inputs,
      ImmutableList<? extends Artifact> schedulingDependencies) {
    Action action = mock(Action.class);
    when(action.getPrimaryOutput()).thenReturn(outputs.get(0));
    when(action.getOutputs()).thenReturn(ImmutableList.copyOf(outputs));
    when(action.getInputs()).thenReturn(NestedSetBuilder.wrap(Order.STABLE_ORDER, inputs));
    when(action.getSchedulingDependencies())
        .thenReturn(NestedSetBuilder.wrap(Order.STABLE_ORDER, schedulingDependencies));
    return action;
  }

  private static SpecialArtifact newTreeArtifact(
      ArtifactRoot root, String name, ActionLookupData generatingActionKey) {
    SpecialArtifact tree =
        SpecialArtifact.create(
            root,
            root.getExecPath().getRelative(name),
            generatingActionKey.getActionLookupKey(),
            SpecialArtifactType.TREE);
    tree.setGeneratingActionKey(generatingActionKey);
    return tree;
  }

  private static final class EqualCancellable
      implements RemoteRewoundActionSynchronizer.Cancellable {
    private final AtomicInteger cancellations = new AtomicInteger();

    @Override
    public void requestCancellation() {
      cancellations.incrementAndGet();
    }

    @Override
    public void awaitCompletion() {}

    @Override
    public boolean equals(Object obj) {
      return obj instanceof EqualCancellable;
    }

    @Override
    public int hashCode() {
      return 1;
    }
  }
}