Only set up the ExecutionProgressReceiver after the first

SomeExecutionStartedEvent. This only concerns Skymeld. PiperOrigin-RevId: 460687994 Change-Id: If09af05787196d8bb90de2e2415ea1572ec68e5a
aranguyen · Jul 20, 2022 · bdbd739 · bdbd739
1 parent be1bde6
commit bdbd739
Showing 1 changed file with 56 additions and 19 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
@@ -24,6 +24,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
+import com.google.common.eventbus.Subscribe;
 import com.google.common.flogger.GoogleLogger;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionCacheChecker;
@@ -91,6 +92,7 @@
 import com.google.devtools.build.lib.skyframe.ConfiguredTargetKey;
 import com.google.devtools.build.lib.skyframe.PackageRootsNoSymlinkCreation;
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor;
+import com.google.devtools.build.lib.skyframe.TopLevelStatusEvents.SomeExecutionStartedEvent;
 import com.google.devtools.build.lib.util.AbruptExitException;
 import com.google.devtools.build.lib.util.DetailedExitCode;
 import com.google.devtools.build.lib.vfs.ModifiedFileSet;
@@ -112,6 +114,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import javax.annotation.Nullable;
 
 /**
@@ -327,26 +330,10 @@ public void prepareForExecution(
           env.getReporter(), executor, request, skyframeBuilder.getActionCacheChecker());
     }
 
-    // Note that executionProgressReceiver accesses builtTargets concurrently (after wrapping in a
-    // synchronized collection), so unsynchronized access to this variable is unsafe while it runs.
-    // TODO(leba): count test actions
-    ExecutionProgressReceiver executionProgressReceiver =
-        new ExecutionProgressReceiver(
-            Preconditions.checkNotNull(builtTargets),
-            Preconditions.checkNotNull(builtAspects),
-            /*exclusiveTestsCount=*/ 0,
-            env.getEventBus());
-    skyframeExecutor
-        .getEventBus()
-        .post(new ExecutionProgressReceiverAvailableEvent(executionProgressReceiver));
-
-    ActionExecutionStatusReporter statusReporter =
-        ActionExecutionStatusReporter.create(env.getReporter(), skyframeExecutor.getEventBus());
+    env.getEventBus()
+        .register(
+            new ExecutionProgressReceiverSetup(skyframeExecutor, env, builtTargets, builtAspects));
     // TODO(leba): Add watchdog support.
-    skyframeExecutor.setActionExecutionProgressReportingObjects(
-        executionProgressReceiver, executionProgressReceiver, statusReporter);
-    skyframeExecutor.setExecutionProgressReceiver(executionProgressReceiver);
-
     for (ExecutorLifecycleListener executorLifecycleListener : executorLifecycleListeners) {
       try (SilentCloseable c =
           Profiler.instance().profile(executorLifecycleListener + ".executionPhaseStarting")) {
@@ -994,4 +981,54 @@ private static AbruptExitException createExitException(
                 .build()),
         e);
   }
+
+  /**
+   * A listener that prepares the ExecutionProgressReceiver upon receiving the first
+   * SomeExecutionStartedEvent. Only activated once a build.
+   */
+  private static class ExecutionProgressReceiverSetup {
+    private final SkyframeExecutor skyframeExecutor;
+    private final CommandEnvironment env;
+    private final Set<ConfiguredTargetKey> builtTargets;
+    private final Set<AspectKey> builtAspects;
+    private final AtomicBoolean activated = new AtomicBoolean(false);
+
+    ExecutionProgressReceiverSetup(
+        SkyframeExecutor skyframeExecutor,
+        CommandEnvironment env,
+        Set<ConfiguredTargetKey> builtTargets,
+        Set<AspectKey> builtAspects) {
+      this.skyframeExecutor = skyframeExecutor;
+      this.env = env;
+      // TODO(b/227138583) Remove these.
+      this.builtTargets = builtTargets;
+      this.builtAspects = builtAspects;
+    }
+
+    @Subscribe
+    public void setupExecutionProgressReceiver(SomeExecutionStartedEvent unused) {
+      if (activated.compareAndSet(false, true)) {
+        // Note that executionProgressReceiver accesses builtTargets concurrently (after wrapping in
+        // a synchronized collection), so unsynchronized access to this variable is unsafe while it
+        // runs.
+        // TODO(leba): count test actions
+        ExecutionProgressReceiver executionProgressReceiver =
+            new ExecutionProgressReceiver(
+                Preconditions.checkNotNull(builtTargets),
+                Preconditions.checkNotNull(builtAspects),
+                /*exclusiveTestsCount=*/ 0,
+                env.getEventBus());
+        env.getEventBus()
+            .post(new ExecutionProgressReceiverAvailableEvent(executionProgressReceiver));
+
+        ActionExecutionStatusReporter statusReporter =
+            ActionExecutionStatusReporter.create(env.getReporter(), skyframeExecutor.getEventBus());
+        skyframeExecutor.setActionExecutionProgressReportingObjects(
+            executionProgressReceiver, executionProgressReceiver, statusReporter);
+        skyframeExecutor.setExecutionProgressReceiver(executionProgressReceiver);
+
+        env.getEventBus().unregister(this);
+      }
+    }
+  }
 }