Loosen visibility of NodeEntrySubject and make ChainedFunction take an arbitrary Runnable for its notifying parts, rather than a CountDownLatch.

PiperOrigin-RevId: 229187454

src/test/java/com/google/devtools/build/skyframe/ChainedFunction.java

diff --git a/src/test/java/com/google/devtools/build/skyframe/ChainedFunction.java b/src/test/java/com/google/devtools/build/skyframe/ChainedFunction.java
index 477d63f..37233f0 100644
--- a/src/test/java/com/google/devtools/build/skyframe/ChainedFunction.java
+++ b/src/test/java/com/google/devtools/build/skyframe/ChainedFunction.java
@@ -26,9 +26,9 @@
  */
 public final class ChainedFunction implements SkyFunction {
   @Nullable private final SkyValue value;
-  @Nullable private final CountDownLatch notifyStart;
+  private final Runnable notifyStart;
   @Nullable private final CountDownLatch waitToFinish;
-  @Nullable private final CountDownLatch notifyFinish;
+  private final Runnable notifyFinish;
   private final boolean waitForException;
   private final Iterable<SkyKey> deps;
 
@@ -40,6 +40,22 @@
       boolean waitForException,
       @Nullable SkyValue value,
       Iterable<SkyKey> deps) {
+    this(
+        makeRunnable(notifyStart),
+        waitToFinish,
+        makeRunnable(notifyFinish),
+        waitForException,
+        value,
+        deps);
+  }
+
+  private ChainedFunction(
+      Runnable notifyStart,
+      @Nullable CountDownLatch waitToFinish,
+      Runnable notifyFinish,
+      boolean waitForException,
+      @Nullable SkyValue value,
+      Iterable<SkyKey> deps) {
     this.notifyStart = notifyStart;
     this.waitToFinish = waitToFinish;
     this.notifyFinish = notifyFinish;
@@ -53,9 +69,7 @@
   public SkyValue compute(SkyKey key, SkyFunction.Environment env) throws GenericFunctionException,
       InterruptedException {
     try {
-      if (notifyStart != null) {
-        notifyStart.countDown();
-      }
+      notifyStart.run();
       if (waitToFinish != null) {
         TrackingAwaiter.INSTANCE.awaitLatchAndTrackExceptions(
             waitToFinish, key + " timed out waiting to finish");
@@ -77,18 +91,20 @@
       }
       return value;
     } finally {
-      if (notifyFinish != null) {
-        notifyFinish.countDown();
-      }
+      notifyFinish.run();
     }
   }
 
+  private static Runnable makeRunnable(@Nullable CountDownLatch latch) {
+    return latch != null ? latch::countDown : () -> {};
+  }
+
   /** Builder for {@link ChainedFunction} objects. */
   public static class Builder {
     @Nullable private SkyValue value;
-    @Nullable private CountDownLatch notifyStart;
+    Runnable notifyStart = makeRunnable(null);
     @Nullable private CountDownLatch waitToFinish;
-    @Nullable private CountDownLatch notifyFinish;
+    private Runnable notifyFinish = makeRunnable(null);
     private boolean waitForException;
     private Iterable<SkyKey> deps = ImmutableList.of();
 
@@ -97,7 +113,7 @@
       return this;
     }
 
-    public Builder setNotifyStart(CountDownLatch notifyStart) {
+    public Builder setNotifyStart(Runnable notifyStart) {
       this.notifyStart = notifyStart;
       return this;
     }
@@ -107,7 +123,7 @@
       return this;
     }
 
-    public Builder setNotifyFinish(CountDownLatch notifyFinish) {
+    public Builder setNotifyFinish(Runnable notifyFinish) {
       this.notifyFinish = notifyFinish;
       return this;
     }

src/test/java/com/google/devtools/build/skyframe/NodeEntrySubject.java

diff --git a/src/test/java/com/google/devtools/build/skyframe/NodeEntrySubject.java b/src/test/java/com/google/devtools/build/skyframe/NodeEntrySubject.java
index 2856298..7b4c6eb 100644
--- a/src/test/java/com/google/devtools/build/skyframe/NodeEntrySubject.java
+++ b/src/test/java/com/google/devtools/build/skyframe/NodeEntrySubject.java
@@ -24,24 +24,23 @@
 import javax.annotation.Nullable;
 
 /**
- * {@link Subject} for {@link NodeEntry}. Please add to this class if you need more
- * functionality!
+ * {@link Subject} for {@link NodeEntry}. Please add to this class if you need more functionality!
  */
-class NodeEntrySubject extends Subject<NodeEntrySubject, NodeEntry> {
+public class NodeEntrySubject extends Subject<NodeEntrySubject, NodeEntry> {
   NodeEntrySubject(FailureMetadata failureMetadata, NodeEntry nodeEntry) {
     super(failureMetadata, nodeEntry);
   }
 
-  DefaultSubject hasVersionThat() {
+  public DefaultSubject hasVersionThat() {
     return assertThat(getSubject().getVersion()).named(detail("Version"));
   }
 
-  IterableSubject hasTemporaryDirectDepsThat() {
+  public IterableSubject hasTemporaryDirectDepsThat() {
     return assertThat(Iterables.concat(getSubject().getTemporaryDirectDeps()))
         .named(detail("TemporaryDirectDeps"));
   }
 
-  ComparableSubject<?, NodeEntry.DependencyState> addReverseDepAndCheckIfDone(
+  public ComparableSubject<?, NodeEntry.DependencyState> addReverseDepAndCheckIfDone(
       @Nullable SkyKey reverseDep) throws InterruptedException {
     return assertThat(getSubject().addReverseDepAndCheckIfDone(reverseDep))
         .named(detail("AddReverseDepAndCheckIfDone"));