Fix crash bug in AbstractExceptionalParallelEvaluator#doMutatingEvaluation in a very specific window of time inbetween enqueueing one top-level node for evaluation and checking if another top-level node is done. See the added unit test for details.

RELNOTES: None
PiperOrigin-RevId: 205718683

4 files changed, 129 insertions, 21 deletions

diff --git a/src/main/java/com/google/devtools/build/skyframe/AbstractExceptionalParallelEvaluator.java b/src/main/java/com/google/devtools/build/skyframe/AbstractExceptionalParallelEvaluator.java
index af658f32d5..315aa8458c 100644
--- a/src/main/java/com/google/devtools/build/skyframe/AbstractExceptionalParallelEvaluator.java
+++ b/src/main/java/com/google/devtools/build/skyframe/AbstractExceptionalParallelEvaluator.java
@@ -260,17 +260,24 @@ public abstract class AbstractExceptionalParallelEvaluator<E extends Exception>
             throw new IllegalStateException(entry + " for " + skyKey + " in unknown state");
         }
       }
-    } catch (InterruptedException e) {
+    } catch (InterruptedException ie) {
       // When multiple keys are being evaluated, it's possible that a key may get queued before
       // an InterruptedException is thrown from either #addReverseDepAndCheckIfDone or
       // #informProgressReceiverThatValueIsDone on a different key. Therefore we have to make sure
       // all evaluation threads are properly interrupted and shut down, if main thread (current
       // thread) is interrupted.
       Thread.currentThread().interrupt();
-      evaluatorContext.getVisitor().waitForCompletion();
+      try {
+        evaluatorContext.getVisitor().waitForCompletion();
+      } catch (SchedulerException se) {
+        // A SchedulerException due to a SkyFunction observing the interrupt is completely expected.
+        if (!(se.getCause() instanceof InterruptedException)) {
+          throw se;
+        }
+      }
 
       // Rethrow the InterruptedException to avoid proceeding to construct the result.
-      throw e;
+      throw ie;
     }
 
     return waitForCompletionAndConstructResult(skyKeys);
diff --git a/src/main/java/com/google/devtools/build/skyframe/InMemoryGraphImpl.java b/src/main/java/com/google/devtools/build/skyframe/InMemoryGraphImpl.java
index 079ddba175..12a2ba3fd4 100644
--- a/src/main/java/com/google/devtools/build/skyframe/InMemoryGraphImpl.java
+++ b/src/main/java/com/google/devtools/build/skyframe/InMemoryGraphImpl.java
@@ -14,7 +14,6 @@
 package com.google.devtools.build.skyframe;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Function;
 import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -35,7 +34,7 @@ import javax.annotation.Nullable;
  */
 public class InMemoryGraphImpl implements InMemoryGraph {
 
-  protected final ConcurrentMap<SkyKey, InMemoryNodeEntry> nodeMap = new ConcurrentHashMap<>(1024);
+  protected final ConcurrentMap<SkyKey, NodeEntry> nodeMap = new ConcurrentHashMap<>(1024);
   private final boolean keepEdges;
 
   InMemoryGraphImpl() {
@@ -52,7 +51,7 @@ public class InMemoryGraphImpl implements InMemoryGraph {
   }
 
   @Override
-  public InMemoryNodeEntry get(@Nullable SkyKey requestor, Reason reason, SkyKey skyKey) {
+  public NodeEntry get(@Nullable SkyKey requestor, Reason reason, SkyKey skyKey) {
     return nodeMap.get(skyKey);
   }
 
@@ -63,7 +62,7 @@ public class InMemoryGraphImpl implements InMemoryGraph {
     // and ImmutableMap.Builder does not tolerate duplicates. The map will be thrown away shortly.
     HashMap<SkyKey, NodeEntry> result = new HashMap<>();
     for (SkyKey key : keys) {
-      InMemoryNodeEntry entry = get(null, Reason.OTHER, key);
+      NodeEntry entry = get(null, Reason.OTHER, key);
       if (entry != null) {
         result.put(key, entry);
       }
@@ -71,17 +70,20 @@ public class InMemoryGraphImpl implements InMemoryGraph {
     return result;
   }
 
-  protected InMemoryNodeEntry createIfAbsent(SkyKey key) {
-    InMemoryNodeEntry newval =
-        keepEdges ? new InMemoryNodeEntry() : new EdgelessInMemoryNodeEntry();
-    InMemoryNodeEntry oldval = nodeMap.putIfAbsent(key, newval);
+  protected NodeEntry newNodeEntry(SkyKey key) {
+    return keepEdges ? new InMemoryNodeEntry() : new EdgelessInMemoryNodeEntry();
+  }
+
+  protected NodeEntry createIfAbsent(SkyKey key) {
+    NodeEntry newval = newNodeEntry(key);
+    NodeEntry oldval = nodeMap.putIfAbsent(key, newval);
     return oldval == null ? newval : oldval;
   }
 
   @Override
-  public Map<SkyKey, InMemoryNodeEntry> createIfAbsentBatch(
+  public Map<SkyKey, NodeEntry> createIfAbsentBatch(
       @Nullable SkyKey requestor, Reason reason, Iterable<SkyKey> keys) {
-    ImmutableMap.Builder<SkyKey, InMemoryNodeEntry> builder = ImmutableMap.builder();
+    ImmutableMap.Builder<SkyKey, NodeEntry> builder = ImmutableMap.builder();
     for (SkyKey key : keys) {
       builder.put(key, createIfAbsent(key));
     }
@@ -98,10 +100,11 @@ public class InMemoryGraphImpl implements InMemoryGraph {
     return Collections.unmodifiableMap(
         Maps.transformValues(
             nodeMap,
-            new Function<InMemoryNodeEntry, SkyValue>() {
-              @Override
-              public SkyValue apply(InMemoryNodeEntry entry) {
+            entry -> {
+              try {
                 return entry.toValue();
+              } catch (InterruptedException e) {
+                throw new IllegalStateException(e);
               }
             }));
   }
@@ -112,17 +115,21 @@ public class InMemoryGraphImpl implements InMemoryGraph {
         Maps.filterValues(
             Maps.transformValues(
                 nodeMap,
-                new Function<InMemoryNodeEntry, SkyValue>() {
-                  @Override
-                  public SkyValue apply(InMemoryNodeEntry entry) {
-                    return entry.isDone() ? entry.getValue() : null;
+                entry -> {
+                  if (!entry.isDone()) {
+                    return null;
+                  }
+                  try {
+                    return entry.getValue();
+                  } catch (InterruptedException e) {
+                    throw new IllegalStateException(e);
                   }
                 }),
             Predicates.notNull()));
   }
 
   @Override
-  public Map<SkyKey, InMemoryNodeEntry> getAllValues() {
+  public Map<SkyKey, NodeEntry> getAllValues() {
     return Collections.unmodifiableMap(nodeMap);
   }
 
diff --git a/src/test/java/com/google/devtools/build/skyframe/BUILD b/src/test/java/com/google/devtools/build/skyframe/BUILD
index b3808d29c5..d43082d6db 100644
--- a/src/test/java/com/google/devtools/build/skyframe/BUILD
+++ b/src/test/java/com/google/devtools/build/skyframe/BUILD
@@ -63,6 +63,7 @@ java_test(
         "//third_party:guava-testlib",
         "//third_party:jsr305",
         "//third_party:junit4",
+        "//third_party:mockito",
         "//third_party:truth",
     ],
 )
diff --git a/src/test/java/com/google/devtools/build/skyframe/ParallelEvaluatorTest.java b/src/test/java/com/google/devtools/build/skyframe/ParallelEvaluatorTest.java
index 51687de2ea..fad52068cd 100644
--- a/src/test/java/com/google/devtools/build/skyframe/ParallelEvaluatorTest.java
+++ b/src/test/java/com/google/devtools/build/skyframe/ParallelEvaluatorTest.java
@@ -61,6 +61,7 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
 
 /**
  * Tests for {@link ParallelEvaluator}.
@@ -229,6 +230,98 @@ public class ParallelEvaluatorTest {
         });
   }
 
+  @Test
+  public void interruptedEvaluatorThreadAfterEnqueueBeforeWaitForCompletionAndConstructResult()
+      throws InterruptedException {
+    // This is a regression test for a crash bug in
+    // AbstractExceptionalParallelEvaluator#doMutatingEvaluation in a very specific window of time
+    // inbetween enqueueing one top-level node for evaluation and checking if another top-level node
+    // is done.
+
+    // When we have two top-level nodes, A and B,
+    SkyKey keyA = GraphTester.toSkyKey("a");
+    SkyKey keyB = GraphTester.toSkyKey("b");
+
+    // And rig the graph and node entries, such that B's addReverseDepAndCheckIfDone waits for A to
+    // start computing and then tries to observe an interrupt (which will happen on the calling
+    // thread, aka the main Skyframe evaluation thread),
+    CountDownLatch keyAStartedComputingLatch = new CountDownLatch(1);
+    CountDownLatch keyBAddReverseDepAndCheckIfDoneLatch = new CountDownLatch(1);
+    NodeEntry nodeEntryB = Mockito.mock(NodeEntry.class);
+    AtomicBoolean keyBAddReverseDepAndCheckIfDoneInterrupted = new AtomicBoolean(false);
+    Mockito.doAnswer(
+            invocation -> {
+              keyAStartedComputingLatch.await();
+              keyBAddReverseDepAndCheckIfDoneLatch.countDown();
+              try {
+                Thread.sleep(TestUtils.WAIT_TIMEOUT_MILLISECONDS);
+                throw new IllegalStateException("shouldn't get here");
+              } catch (InterruptedException e) {
+                keyBAddReverseDepAndCheckIfDoneInterrupted.set(true);
+                throw e;
+              }
+            })
+        .when(nodeEntryB)
+        .addReverseDepAndCheckIfDone(Mockito.eq(null));
+    graph = new InMemoryGraphImpl() {
+      @Override
+      protected NodeEntry newNodeEntry(SkyKey key) {
+        return key.equals(keyB) ? nodeEntryB : super.newNodeEntry(key);
+      }
+    };
+    // And A's SkyFunction tries to observe an interrupt after it starts computing,
+    AtomicBoolean keyAComputeInterrupted = new AtomicBoolean(false);
+    tester.getOrCreate(keyA).setBuilder(new SkyFunction() {
+      @Override
+      public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedException {
+        keyAStartedComputingLatch.countDown();
+        try {
+          Thread.sleep(TestUtils.WAIT_TIMEOUT_MILLISECONDS);
+          throw new IllegalStateException("shouldn't get here");
+        } catch (InterruptedException e) {
+          keyAComputeInterrupted.set(true);
+          throw e;
+        }
+      }
+
+      @Override
+      public String extractTag(SkyKey skyKey) {
+        return null;
+      }
+    });
+
+    // And we have a dedicated thread that kicks off the evaluation of A and B together (in that
+    // order).
+    TestThread evalThread = new TestThread() {
+      @Override
+      public void runTest() throws Exception {
+        try {
+          eval(/*keepGoing=*/true, keyA, keyB);
+          fail();
+        } catch (InterruptedException e) {
+          // Expected.
+        }
+      }
+    };
+
+    // Then when we start that thread,
+    evalThread.start();
+    // We (the thread running the test) are able to observe that B's addReverseDepAndCheckIfDone has
+    // just been called (implying that A has started to be computed).
+    assertThat(
+            keyBAddReverseDepAndCheckIfDoneLatch.await(
+                TestUtils.WAIT_TIMEOUT_SECONDS, TimeUnit.SECONDS))
+        .isTrue();
+    // Then when we interrupt the evaluation thread,
+    evalThread.interrupt();
+    // The evaluation thread eventually terminates.
+    evalThread.joinAndAssertState(TestUtils.WAIT_TIMEOUT_MILLISECONDS);
+    // And we are able to verify both that A's SkyFunction had observed an interrupt,
+    assertThat(keyAComputeInterrupted.get()).isTrue();
+    // And also that B's addReverseDepAndCheckIfDoneInterrupted had observed an interrupt.
+    assertThat(keyBAddReverseDepAndCheckIfDoneInterrupted.get()).isTrue();
+  }
+
   private void runPartialResultOnInterruption(boolean buildFastFirst) throws Exception {
     graph = new InMemoryGraphImpl();
     // Two runs for fastKey's builder and one for the start of waitKey's builder.