Add a test for rewinding a top-level output owned by an aspect.

PiperOrigin-RevId: 605450336 Change-Id: I48e847f5ca026c180cd980089e327a53a277db51
bazelbuild · Feb 8, 2024 · b157016 · b157016
1 parent 0e04fc2
commit b157016
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 10 deletions.
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/rewinding/BUILD b/src/test/java/com/google/devtools/build/lib/skyframe/rewinding/BUILD
@@ -35,6 +35,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset:artifact_nested_set_key",
         "//src/main/java/com/google/devtools/build/lib/exec:module_action_context_registry",
         "//src/main/java/com/google/devtools/build/lib/exec:spawn_exec_exception",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:aspect_key_creator",
         "//src/main/java/com/google/devtools/build/lib/skyframe/rewinding",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",

diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/rewinding/RewindingTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/rewinding/RewindingTest.java
@@ -257,4 +257,9 @@ public void discoveredCppModuleLost() throws Exception {
   public void topLevelOutputRewound_regularFile() throws Exception {
     helper.runTopLevelOutputRewound_regularFile();
   }
+
+  @Test
+  public void topLevelOutputRewound_aspectOwned() throws Exception {
+    helper.runTopLevelOutputRewound_aspectOwned();
+  }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/rewinding/RewindingTestsHelper.java b/src/test/java/com/google/devtools/build/lib/skyframe/rewinding/RewindingTestsHelper.java
@@ -44,9 +44,11 @@
 import com.google.devtools.build.lib.actions.ActionLookupData;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.BuildFailedException;
+import com.google.devtools.build.lib.actions.EventReportingArtifacts;
 import com.google.devtools.build.lib.actions.LostInputsExecException;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnResult;
+import com.google.devtools.build.lib.analysis.AspectCompleteEvent;
 import com.google.devtools.build.lib.analysis.TargetCompleteEvent;
 import com.google.devtools.build.lib.analysis.config.CoreOptions;
 import com.google.devtools.build.lib.bugreport.BugReporter;
@@ -62,6 +64,7 @@
 import com.google.devtools.build.lib.server.FailureDetails.ActionRewinding;
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
 import com.google.devtools.build.lib.server.FailureDetails.Spawn.Code;
+import com.google.devtools.build.lib.skyframe.AspectKeyCreator.AspectKey;
 import com.google.devtools.build.lib.testutil.ActionEventRecorder;
 import com.google.devtools.build.lib.testutil.ActionEventRecorder.ActionRewindEventAsserter;
 import com.google.devtools.build.lib.testutil.ControllableActionStrategyModule;
@@ -2659,16 +2662,7 @@ public final void runTopLevelOutputRewound_regularFile() throws Exception {
     Label fooLost = Label.parseCanonical("//foo:lost");
     List<SkyKey> rewoundKeys = collectOrderedRewoundKeys();
     Map<Label, TargetCompleteEvent> targetCompleteEvents = recordTargetCompleteEvents();
-
-    testCase.injectListenerAtStartOfNextBuild(
-        (key, type, order, context) -> {
-          if (type == EventType.MARK_DIRTY
-              || (isActionExecutionKey(key, fooLost) && type == EventType.SET_VALUE)) {
-            // The TargetCompleteEvent should not be emitted until after rewinding completes.
-            // Otherwise, we may publish stale artifact URIs to the BEP.
-            assertThat(targetCompleteEvents).isEmpty();
-          }
-        });
+    listenForNoCompletionEventsBeforeRewinding(fooLost, targetCompleteEvents);
 
     testCase.buildTarget("//foo:lost");
 
@@ -2680,6 +2674,48 @@ public final void runTopLevelOutputRewound_regularFile() throws Exception {
     assertThat(targetCompleteEvents.keySet()).containsExactly(fooLost);
   }
 
+  public final void runTopLevelOutputRewound_aspectOwned() throws Exception {
+    testCase.write(
+        "foo/defs.bzl",
+        "def _lost_aspect_impl(target, ctx):",
+        "  out = ctx.actions.declare_file(ctx.rule.attr.name + '_lost_aspect.out')",
+        "  ctx.actions.run_shell(outputs = [out], command = 'echo lost_aspect > %s' % out.path)",
+        "  return [OutputGroupInfo(default = depset([out]))]",
+        "",
+        "lost_aspect = aspect(implementation = _lost_aspect_impl)");
+    testCase.write("foo/BUILD", "sh_library(name = 'lib')");
+    lostOutputsModule.addLostOutput(getExecPath("bin/foo/lib_lost_aspect.out"));
+    Label fooLib = Label.parseCanonical("//foo:lib");
+    List<SkyKey> rewoundKeys = collectOrderedRewoundKeys();
+    Map<Label, AspectCompleteEvent> aspectCompleteEvents = recordAspectCompleteEvents();
+    listenForNoCompletionEventsBeforeRewinding(fooLib, aspectCompleteEvents);
+
+    testCase.addOptions("--aspects=foo/defs.bzl%lost_aspect");
+    testCase.buildTarget("//foo:lib");
+
+    lostOutputsModule.verifyAllLostOutputsConsumed();
+    assertThat(rewoundKeys).hasSize(1);
+    assertThat(rewoundArtifactOwnerLabels(rewoundKeys)).containsExactly("//foo:lib");
+    assertThat(((ActionLookupData) rewoundKeys.get(0)).getActionLookupKey())
+        .isInstanceOf(AspectKey.class);
+    assertThat(ImmutableMultiset.copyOf(getExecutedSpawnDescriptions()))
+        .hasCount("Action foo/lib_lost_aspect.out", 2);
+    assertThat(aspectCompleteEvents.keySet()).containsExactly(fooLib);
+  }
+
+  private void listenForNoCompletionEventsBeforeRewinding(
+      Label lostLabel, Map<Label, ? extends EventReportingArtifacts> events) {
+    testCase.injectListenerAtStartOfNextBuild(
+        (key, type, order, context) -> {
+          if (type == EventType.MARK_DIRTY
+              || (isActionExecutionKey(key, lostLabel) && type == EventType.SET_VALUE)) {
+            // Completion events for lost outputs should not be emitted until after rewinding
+            // completes. Otherwise, we may publish stale artifact URIs to the BEP.
+            assertThat(events).isEmpty();
+          }
+        });
+  }
+
   static boolean isActionExecutionKey(Object key, Label label) {
     return key instanceof ActionLookupData && label.equals(((ActionLookupData) key).getLabel());
   }
@@ -2713,6 +2749,22 @@ public void accept(TargetCompleteEvent event) {
     return targetCompleteEvents;
   }
 
+  private Map<Label, AspectCompleteEvent> recordAspectCompleteEvents() {
+    Map<Label, AspectCompleteEvent> aspectCompleteEvents = new HashMap<>();
+    testCase
+        .getRuntimeWrapper()
+        .registerSubscriber(
+            new Object() {
+              @Subscribe
+              @SuppressWarnings("unused")
+              public void accept(AspectCompleteEvent event) {
+                // If we need to track targets with multiple aspects, we could change the key type.
+                aspectCompleteEvents.put(event.getLabel(), event);
+              }
+            });
+    return aspectCompleteEvents;
+  }
+
   /**
    * Converts a root-relative output path to an exec path, accounting for the top-level
    * configuration's mnemonic and {@link TestConstants#PRODUCT_NAME}.