Notify Skyframe listeners about additional node states: CHECK_DIRTY, INIT_ENV, and COMMIT.

PiperOrigin-RevId: 193261989

4 files changed, 67 insertions, 33 deletions

diff --git a/src/main/java/com/google/devtools/build/skyframe/AbstractParallelEvaluator.java b/src/main/java/com/google/devtools/build/skyframe/AbstractParallelEvaluator.java
index 1391d00e1b..6ed51ad3f8 100644
--- a/src/main/java/com/google/devtools/build/skyframe/AbstractParallelEvaluator.java
+++ b/src/main/java/com/google/devtools/build/skyframe/AbstractParallelEvaluator.java
@@ -26,6 +26,7 @@ import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.util.GroupedList.GroupedListHelper;
 import com.google.devtools.build.skyframe.EvaluationProgressReceiver.EvaluationState;
+import com.google.devtools.build.skyframe.EvaluationProgressReceiver.NodeState;
 import com.google.devtools.build.skyframe.MemoizingEvaluator.EmittedEventState;
 import com.google.devtools.build.skyframe.NodeEntry.DependencyState;
 import com.google.devtools.build.skyframe.NodeEntry.DirtyState;
@@ -322,14 +323,26 @@ public abstract class AbstractParallelEvaluator {
         NodeEntry state =
             Preconditions.checkNotNull(graph.get(null, Reason.EVALUATION, skyKey), skyKey);
         Preconditions.checkState(state.isReady(), "%s %s", skyKey, state);
-        if (maybeHandleDirtyNode(state) == DirtyOutcome.ALREADY_PROCESSED) {
-          return;
+        try {
+          evaluatorContext.getProgressReceiver().stateStarting(skyKey, NodeState.CHECK_DIRTY);
+          if (maybeHandleDirtyNode(state) == DirtyOutcome.ALREADY_PROCESSED) {
+            return;
+          }
+        } finally {
+          evaluatorContext.getProgressReceiver().stateEnding(skyKey, NodeState.CHECK_DIRTY, -1);
         }
 
         Set<SkyKey> oldDeps = state.getAllRemainingDirtyDirectDeps();
-        SkyFunctionEnvironment env =
-            new SkyFunctionEnvironment(
-                skyKey, state.getTemporaryDirectDeps(), oldDeps, evaluatorContext);
+        SkyFunctionEnvironment env;
+        try {
+          evaluatorContext.getProgressReceiver().stateStarting(skyKey,
+              NodeState.INITIALIZING_ENVIRONMENT);
+          env = new SkyFunctionEnvironment(skyKey, state.getTemporaryDirectDeps(), oldDeps,
+              evaluatorContext);
+        } finally {
+          evaluatorContext.getProgressReceiver().stateEnding(skyKey,
+              NodeState.INITIALIZING_ENVIRONMENT, -1);
+        }
         SkyFunctionName functionName = skyKey.functionName();
         SkyFunction factory =
             Preconditions.checkNotNull(
@@ -343,12 +356,13 @@ public abstract class AbstractParallelEvaluator {
         long startTime = BlazeClock.instance().nanoTime();
         try {
           try {
-            evaluatorContext.getProgressReceiver().computing(skyKey);
+            evaluatorContext.getProgressReceiver().stateStarting(skyKey, NodeState.COMPUTE);
             value = factory.compute(skyKey, env);
           } finally {
             long elapsedTimeNanos = BlazeClock.instance().nanoTime() - startTime;
             if (elapsedTimeNanos > 0) {
-              evaluatorContext.getProgressReceiver().computed(skyKey, elapsedTimeNanos);
+              evaluatorContext.getProgressReceiver().stateEnding(skyKey, NodeState.COMPUTE,
+                  elapsedTimeNanos);
               Profiler.instance()
                   .logSimpleTaskDuration(
                       startTime,
@@ -440,14 +454,16 @@ public abstract class AbstractParallelEvaluator {
               skyKey,
               newDirectDeps,
               state);
-          env.setValue(value);
-          registerNewlyDiscoveredDepsForDoneEntry(
-              skyKey,
-              state,
-              oldDeps,
-              env,
-              evaluatorContext.keepGoing());
-          env.commit(state, EnqueueParentBehavior.ENQUEUE);
+
+          try {
+            evaluatorContext.getProgressReceiver().stateStarting(skyKey, NodeState.COMMIT);
+            env.setValue(value);
+            registerNewlyDiscoveredDepsForDoneEntry(
+                skyKey, state, oldDeps, env, evaluatorContext.keepGoing());
+            env.commit(state, EnqueueParentBehavior.ENQUEUE);
+          } finally {
+            evaluatorContext.getProgressReceiver().stateEnding(skyKey, NodeState.COMMIT, -1);
+          }
           return;
         }
 
diff --git a/src/main/java/com/google/devtools/build/skyframe/CompoundEvaluationProgressReceiver.java b/src/main/java/com/google/devtools/build/skyframe/CompoundEvaluationProgressReceiver.java
index 21f50cf6a3..894df66fda 100644
--- a/src/main/java/com/google/devtools/build/skyframe/CompoundEvaluationProgressReceiver.java
+++ b/src/main/java/com/google/devtools/build/skyframe/CompoundEvaluationProgressReceiver.java
@@ -47,16 +47,16 @@ public class CompoundEvaluationProgressReceiver implements EvaluationProgressRec
   }
 
   @Override
-  public void computing(SkyKey skyKey) {
+  public void stateStarting(SkyKey skyKey, NodeState state) {
     for (EvaluationProgressReceiver receiver : receivers) {
-      receiver.computing(skyKey);
+      receiver.stateStarting(skyKey, state);
     }
   }
 
   @Override
-  public void computed(SkyKey skyKey, long elapsedTimeNanos) {
+  public void stateEnding(SkyKey skyKey, NodeState state, long elapsedTimeNanos) {
     for (EvaluationProgressReceiver receiver : receivers) {
-      receiver.computed(skyKey, elapsedTimeNanos);
+      receiver.stateEnding(skyKey, state, elapsedTimeNanos);
     }
   }
 
diff --git a/src/main/java/com/google/devtools/build/skyframe/DirtyTrackingProgressReceiver.java b/src/main/java/com/google/devtools/build/skyframe/DirtyTrackingProgressReceiver.java
index f954a8ca54..22f9dfba1d 100644
--- a/src/main/java/com/google/devtools/build/skyframe/DirtyTrackingProgressReceiver.java
+++ b/src/main/java/com/google/devtools/build/skyframe/DirtyTrackingProgressReceiver.java
@@ -92,16 +92,16 @@ public class DirtyTrackingProgressReceiver implements EvaluationProgressReceiver
   }
 
   @Override
-  public void computing(SkyKey skyKey) {
+  public void stateStarting(SkyKey skyKey, NodeState nodeState) {
     if (progressReceiver != null) {
-      progressReceiver.computing(skyKey);
+      progressReceiver.stateStarting(skyKey, nodeState);
     }
   }
 
   @Override
-  public void computed(SkyKey skyKey, long elapsedTimeNanos) {
+  public void stateEnding(SkyKey skyKey, NodeState nodeState, long elapsedTimeNanos) {
     if (progressReceiver != null) {
-      progressReceiver.computed(skyKey, elapsedTimeNanos);
+      progressReceiver.stateEnding(skyKey, nodeState, elapsedTimeNanos);
     }
   }
 
diff --git a/src/main/java/com/google/devtools/build/skyframe/EvaluationProgressReceiver.java b/src/main/java/com/google/devtools/build/skyframe/EvaluationProgressReceiver.java
index 8787c3e60f..2844ec4f45 100644
--- a/src/main/java/com/google/devtools/build/skyframe/EvaluationProgressReceiver.java
+++ b/src/main/java/com/google/devtools/build/skyframe/EvaluationProgressReceiver.java
@@ -40,6 +40,20 @@ public interface EvaluationProgressReceiver {
   }
 
   /**
+   * Overall state of the node while it is being evaluated.
+   */
+  enum NodeState {
+    /** The node is undergoing a dirtiness check and may be re-validated. */
+    CHECK_DIRTY,
+    /** The node is prepping for evaluation. */
+    INITIALIZING_ENVIRONMENT,
+    /** The node is in compute(). */
+    COMPUTE,
+    /** The node is done evaluation and committing the result. */
+    COMMIT,
+  }
+
+  /**
    * Notifies that the node named by {@code key} has been invalidated.
    *
    * <p>{@code state} indicates the new state of the value.
@@ -59,20 +73,24 @@ public interface EvaluationProgressReceiver {
   void enqueueing(SkyKey skyKey);
 
   /**
-   * Notifies that {@code skyFunction.compute(skyKey, ...)} is about to be called, for some
-   * appropriate {@link SkyFunction} {@code skyFunction}.
+   * Notifies that a node corresponding to {@code skyKey} is about to enter the given
+   * {@code nodeState}.
    *
-   * <p>Notably, this includes {@link SkyFunction#compute} calls due to Skyframe restarts.
+   * <p>Notably, this includes {@link SkyFunction#compute} calls due to Skyframe restarts, but also
+   * dirtiness checking and node completion.
    */
-  void computing(SkyKey skyKey);
+  void stateStarting(SkyKey skyKey, NodeState nodeState);
 
   /**
-   * Notifies that {@code skyFunction.compute(skyKey, ...)} has just been called, for some
-   * appropriate {@link SkyFunction} {@code skyFunction}.
+   * Notifies that a node corresponding to {@code skyKey} is about to complete the given
+   * {@code nodeState}.
+   *
+   * <p>Always called symmetrically with {@link #stateStarting(SkyKey, NodeState)}}.
    *
-   * <p>Notably, this includes {@link SkyFunction#compute} calls due to Skyframe restarts.
+   * <p>{@code elapsedTimeNanos} is either the elapsed time in the {@code nodeState} or -1 if the
+   * timing was not recorded.
    */
-  void computed(SkyKey skyKey, long elapsedTimeNanos);
+  void stateEnding(SkyKey skyKey, NodeState nodeState, long elapsedTimeNanos);
 
   /**
    * Notifies that the node for {@code skyKey} has been evaluated.
@@ -95,11 +113,11 @@ public interface EvaluationProgressReceiver {
     }
 
     @Override
-    public void computing(SkyKey skyKey) {
+    public void stateStarting(SkyKey skyKey, NodeState nodeState) {
     }
 
     @Override
-    public void computed(SkyKey skyKey, long elapsedTimeNanos) {
+    public void stateEnding(SkyKey skyKey, NodeState nodeState, long elapsedTimeNanos) {
     }
 
     @Override