Add "out" field to actions in JSON profile.

This provides a way to link the JSON profile to aquery's [text]proto output via the actions' primary output.

Also added `primary_output_id` to analysis_v2.proto.

RELNOTES: Add actions' primary output details in JSON profile and analysis_v2.proto.
PiperOrigin-RevId: 295940635
diff --git a/src/main/java/com/google/devtools/build/lib/profiler/Profiler.java b/src/main/java/com/google/devtools/build/lib/profiler/Profiler.java
index dc20823..70c8070 100644
--- a/src/main/java/com/google/devtools/build/lib/profiler/Profiler.java
+++ b/src/main/java/com/google/devtools/build/lib/profiler/Profiler.java
@@ -187,14 +187,13 @@
   }
 
   /**
-   * Container for the single task record.
-   * Should never be instantiated directly - use TaskStack.create() instead.
+   * Container for the single task record. Should never be instantiated directly - use
+   * TaskStack.create() instead.
    *
-   * Class itself is not thread safe, but all access to it from Profiler
-   * methods is.
+   * <p>Class itself is not thread safe, but all access to it from Profiler methods is.
    */
   @ThreadCompatible
-  private static final class TaskData {
+  private static class TaskData {
     final long threadId;
     final long startTimeNanos;
     final int id;
@@ -244,6 +243,21 @@
     }
   }
 
+  private static final class ActionTaskData extends TaskData {
+    final String primaryOutputPath;
+
+    ActionTaskData(
+        int id,
+        long startTimeNanos,
+        TaskData parent,
+        ProfilerTask eventType,
+        String description,
+        String primaryOutputPath) {
+      super(id, startTimeNanos, parent, eventType, description);
+      this.primaryOutputPath = primaryOutputPath;
+    }
+  }
+
   /**
    * Tracks nested tasks for each thread.
    *
@@ -285,6 +299,16 @@
       return new TaskData(taskId.incrementAndGet(), startTimeNanos, peek(), eventType, description);
     }
 
+    public void pushActionTask(ProfilerTask eventType, String description, String primaryOutput) {
+      get().add(createActionTask(clock.nanoTime(), eventType, description, primaryOutput));
+    }
+
+    public ActionTaskData createActionTask(
+        long startTimeNanos, ProfilerTask eventType, String description, String primaryOutput) {
+      return new ActionTaskData(
+          taskId.incrementAndGet(), startTimeNanos, peek(), eventType, description, primaryOutput);
+    }
+
     @Override
     public String toString() {
       StringBuilder builder = new StringBuilder(
@@ -493,7 +517,8 @@
       long execStartTimeNanos,
       boolean enabledCpuUsageProfiling,
       boolean slimProfile,
-      boolean enableActionCountProfile)
+      boolean enableActionCountProfile,
+      boolean includePrimaryOutput)
       throws IOException {
     Preconditions.checkState(!isActive(), "Profiler already active");
     initHistograms();
@@ -520,7 +545,13 @@
       switch (format) {
         case JSON_TRACE_FILE_FORMAT:
           writer =
-              new JsonTraceFileWriter(stream, execStartTimeNanos, slimProfile, outputBase, buildID);
+              new JsonTraceFileWriter(
+                  stream,
+                  execStartTimeNanos,
+                  slimProfile,
+                  outputBase,
+                  buildID,
+                  includePrimaryOutput);
           break;
         case JSON_TRACE_FILE_COMPRESSED_FORMAT:
           writer =
@@ -529,7 +560,8 @@
                   execStartTimeNanos,
                   slimProfile,
                   outputBase,
-                  buildID);
+                  buildID,
+                  includePrimaryOutput);
       }
       writer.start();
     }
@@ -793,6 +825,22 @@
   }
 
   /**
+   * Similar to {@link #profile}, but specific to action-related events. Takes an extra argument:
+   * primaryOutput.
+   */
+  public SilentCloseable profileAction(
+      ProfilerTask type, String description, String primaryOutput) {
+
+    Preconditions.checkNotNull(description);
+    if (isActive() && isProfiling(type)) {
+      taskStack.pushActionTask(type, description, primaryOutput);
+      return () -> completeTask(type);
+    } else {
+      return () -> {};
+    }
+  }
+
+  /**
    * Records the beginning of a task as specified, and returns a {@link SilentCloseable} instance
    * that ends the task. This lets the system do the work of ending the task, with the compiler
    * giving a warning if the returned instance is not closed.
@@ -913,6 +961,7 @@
     private final ThreadLocal<Boolean> metadataPosted =
         ThreadLocal.withInitial(() -> Boolean.FALSE);
     private final boolean slimProfile;
+    private final boolean includePrimaryOutput;
     private final UUID buildID;
     private final String outputBase;
 
@@ -928,12 +977,14 @@
         long profileStartTimeNanos,
         boolean slimProfile,
         String outputBase,
-        UUID buildID) {
+        UUID buildID,
+        boolean includePrimaryOutput) {
       this.outStream = outStream;
       this.profileStartTimeNanos = profileStartTimeNanos;
       this.slimProfile = slimProfile;
       this.buildID = buildID;
       this.outputBase = outputBase;
+      this.includePrimaryOutput = includePrimaryOutput;
     }
 
     @Override
@@ -1035,6 +1086,11 @@
         writer.name("dur").value(TimeUnit.NANOSECONDS.toMicros(data.duration));
       }
       writer.name("pid").value(1);
+
+      // Primary outputs are non-mergeable, thus incompatible with slim profiles.
+      if (includePrimaryOutput && data instanceof ActionTaskData) {
+        writer.name("out").value(((ActionTaskData) data).primaryOutputPath);
+      }
       long threadId =
           data.type == ProfilerTask.CRITICAL_PATH_COMPONENT
               ? CRITICAL_PATH_THREAD_ID
diff --git a/src/main/java/com/google/devtools/build/lib/profiler/TraceEvent.java b/src/main/java/com/google/devtools/build/lib/profiler/TraceEvent.java
index 009ec93..c2558c1 100644
--- a/src/main/java/com/google/devtools/build/lib/profiler/TraceEvent.java
+++ b/src/main/java/com/google/devtools/build/lib/profiler/TraceEvent.java
@@ -35,8 +35,20 @@
       String name,
       @Nullable Duration timestamp,
       @Nullable Duration duration,
+      long threadId,
+      @Nullable String primaryOutputPath) {
+    return new AutoValue_TraceEvent(
+        category, name, timestamp, duration, threadId, primaryOutputPath);
+  }
+
+  public static TraceEvent create(
+      @Nullable String category,
+      String name,
+      @Nullable Duration timestamp,
+      @Nullable Duration duration,
       long threadId) {
-    return new AutoValue_TraceEvent(category, name, timestamp, duration, threadId);
+    return new AutoValue_TraceEvent(
+        category, name, timestamp, duration, threadId, /* primaryOutputPath= */ null);
   }
 
   @Nullable
@@ -52,12 +64,17 @@
 
   public abstract long threadId();
 
+  // Only applicable to action-related TraceEvents.
+  @Nullable
+  public abstract String primaryOutputPath();
+
   private static TraceEvent createFromJsonReader(JsonReader reader) throws IOException {
     String category = null;
     String name = null;
     Duration timestamp = null;
     Duration duration = null;
     long threadId = -1;
+    String primaryOutputPath = null;
 
     reader.beginObject();
     while (reader.hasNext()) {
@@ -78,12 +95,15 @@
         case "tid":
           threadId = reader.nextLong();
           break;
+        case "out":
+          primaryOutputPath = reader.nextString();
+          break;
         default:
           reader.skipValue();
       }
     }
     reader.endObject();
-    return TraceEvent.create(category, name, timestamp, duration, threadId);
+    return TraceEvent.create(category, name, timestamp, duration, threadId, primaryOutputPath);
   }
 
   public static List<TraceEvent> parseTraceEvents(JsonReader reader) throws IOException {
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
index 837a6d1..44bec37 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
@@ -357,6 +357,13 @@
       }
       ImmutableSet<ProfilerTask> profiledTasks = profiledTasksBuilder.build();
       if (!profiledTasks.isEmpty()) {
+        if (options.enableJsonProfileDiet && options.includePrimaryOutput) {
+          eventHandler.handle(
+              Event.warn(
+                  "Enabling both --experimental_slim_json_profile and"
+                      + " --experimental_include_primary_output: the \"out\" field"
+                      + " will be omitted in merged actions."));
+        }
         Profiler profiler = Profiler.instance();
         profiler.start(
             profiledTasks,
@@ -369,7 +376,8 @@
             execStartTimeNanos,
             options.enableCpuUsageProfiling,
             options.enableJsonProfileDiet,
-            options.enableActionCountProfile);
+            options.enableActionCountProfile,
+            options.includePrimaryOutput);
         // Instead of logEvent() we're calling the low level function to pass the timings we took in
         // the launcher. We're setting the INIT phase marker so that it follows immediately the
         // LAUNCH phase.
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/CommonCommandOptions.java b/src/main/java/com/google/devtools/build/lib/runtime/CommonCommandOptions.java
index 0659fe5..218dd2d 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/CommonCommandOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/CommonCommandOptions.java
@@ -313,6 +313,16 @@
   public boolean enableJsonProfileDiet;
 
   @Option(
+      name = "experimental_include_primary_output",
+      defaultValue = "false",
+      documentationCategory = OptionDocumentationCategory.LOGGING,
+      effectTags = {OptionEffectTag.AFFECTS_OUTPUTS, OptionEffectTag.BAZEL_MONITORING},
+      help =
+          "Includes the extra \"out\" attribute in action events that contains the exec path "
+              + "to the action's primary output.")
+  public boolean includePrimaryOutput;
+
+  @Option(
       name = "experimental_announce_profile_path",
       defaultValue = "false",
       documentationCategory = OptionDocumentationCategory.LOGGING,
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
index d64509e..2789396 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
@@ -1010,7 +1010,11 @@
       //
       // It is also unclear why we are posting anything directly to reporter. That probably
       // shouldn't happen.
-      try (SilentCloseable c = profiler.profile(ProfilerTask.ACTION, action.describe())) {
+      try (SilentCloseable c =
+          profiler.profileAction(
+              ProfilerTask.ACTION,
+              action.describe(),
+              action.getPrimaryOutput().getExecPathString())) {
         String message = action.getProgressMessage();
         if (message != null) {
           reporter.startTask(null, prependExecPhaseStats(message));
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java b/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java
index 8530c5b..0301bf6 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java
@@ -160,12 +160,11 @@
       // environment as well.
       Map<String, String> fixedEnvironment = spawnAction.getEnvironment().getFixedEnv().toMap();
       for (Map.Entry<String, String> environmentVariable : fixedEnvironment.entrySet()) {
-        AnalysisProtosV2.KeyValuePair.Builder keyValuePairBuilder =
-            AnalysisProtosV2.KeyValuePair.newBuilder();
-        keyValuePairBuilder
-            .setKey(environmentVariable.getKey())
-            .setValue(environmentVariable.getValue());
-        actionBuilder.addEnvironmentVariables(keyValuePairBuilder.build());
+        actionBuilder.addEnvironmentVariables(
+            AnalysisProtosV2.KeyValuePair.newBuilder()
+                .setKey(environmentVariable.getKey())
+                .setValue(environmentVariable.getValue())
+                .build());
       }
     }
 
@@ -221,7 +220,7 @@
       NestedSet<Artifact> inputs = action.getInputs();
       NestedSetView<Artifact> nestedSetView = new NestedSetView<>(inputs);
 
-      if (nestedSetView.directs().size() > 0 || nestedSetView.transitives().size() > 0) {
+      if (!nestedSetView.directs().isEmpty() || !nestedSetView.transitives().isEmpty()) {
         actionBuilder.addInputDepSetIds(
             knownNestedSets.dataToIdAndStreamOutputProto(nestedSetView));
       }
@@ -230,6 +229,9 @@
       for (Artifact artifact : action.getOutputs()) {
         actionBuilder.addOutputIds(knownArtifacts.dataToIdAndStreamOutputProto(artifact));
       }
+
+      actionBuilder.setPrimaryOutputId(
+          knownArtifacts.dataToIdAndStreamOutputProto(action.getPrimaryOutput()));
     }
 
     aqueryOutputHandler.outputAction(actionBuilder.build());
diff --git a/src/main/protobuf/analysis_v2.proto b/src/main/protobuf/analysis_v2.proto
index edc4f1f..03c282a 100644
--- a/src/main/protobuf/analysis_v2.proto
+++ b/src/main/protobuf/analysis_v2.proto
@@ -97,6 +97,9 @@
 
   // The list of param files. This will be only set if explicitly requested.
   repeated ParamFile param_files = 12;
+
+  // The id to an Artifact that is the primary output of this action.
+  uint32 primary_output_id = 13;
 }
 
 // Represents a single target (without configuration information) that is
diff --git a/src/test/java/com/google/devtools/build/lib/buildtool/util/BlazeRuntimeWrapper.java b/src/test/java/com/google/devtools/build/lib/buildtool/util/BlazeRuntimeWrapper.java
index 59d30e4..b75366a 100644
--- a/src/test/java/com/google/devtools/build/lib/buildtool/util/BlazeRuntimeWrapper.java
+++ b/src/test/java/com/google/devtools/build/lib/buildtool/util/BlazeRuntimeWrapper.java
@@ -284,7 +284,8 @@
               /* execStartTimeNanos= */ 42,
               /* enabledCpuUsageProfiling= */ false,
               /* slimProfile= */ false,
-              /* enableActionCountProfile= */ false);
+              /* enableActionCountProfile= */ false,
+              /* includePrimaryOutput= */ false);
       OutErr outErr = env.getReporter().getOutErr();
       System.setOut(new PrintStream(outErr.getOutputStream(), /*autoflush=*/ true));
       System.setErr(new PrintStream(outErr.getErrorStream(), /*autoflush=*/ true));
diff --git a/src/test/java/com/google/devtools/build/lib/profiler/AutoProfilerBenchmark.java b/src/test/java/com/google/devtools/build/lib/profiler/AutoProfilerBenchmark.java
index 72f9a77..f02b2ff 100644
--- a/src/test/java/com/google/devtools/build/lib/profiler/AutoProfilerBenchmark.java
+++ b/src/test/java/com/google/devtools/build/lib/profiler/AutoProfilerBenchmark.java
@@ -40,7 +40,8 @@
             BlazeClock.instance().nanoTime(),
             /* enabledCpuUsageProfiling= */ false,
             /* slimProfile= */ false,
-            /* enableActionCountProfile= */ false);
+            /* enableActionCountProfile= */ false,
+            /* includePrimaryOutput= */ false);
   }
 
   @BeforeExperiment
diff --git a/src/test/java/com/google/devtools/build/lib/profiler/ProfilerTest.java b/src/test/java/com/google/devtools/build/lib/profiler/ProfilerTest.java
index 51615cd..84811ae 100644
--- a/src/test/java/com/google/devtools/build/lib/profiler/ProfilerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/profiler/ProfilerTest.java
@@ -97,7 +97,8 @@
         BlazeClock.nanoTime(),
         /* enabledCpuUsageProfiling= */ false,
         /* slimProfile= */ false,
-        /* enableActionCountProfile= */ false);
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ false);
     return buffer;
   }
 
@@ -113,7 +114,8 @@
         BlazeClock.nanoTime(),
         /* enabledCpuUsageProfiling= */ false,
         /* slimProfile= */ false,
-        /* enableActionCountProfile= */ false);
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ false);
   }
 
   @Test
@@ -205,7 +207,8 @@
         clock.nanoTime(),
         /* enabledCpuUsageProfiling= */ false,
         /* slimProfile= */ false,
-        /* enableActionCountProfile= */ false);
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ false);
     try (SilentCloseable c = profiler.profile(ProfilerTask.ACTION, "action task")) {
       // Next task takes less than 10 ms but should be recorded anyway.
       long before = clock.nanoTime();
@@ -242,7 +245,8 @@
         BlazeClock.instance().nanoTime(),
         /* enabledCpuUsageProfiling= */ false,
         /* slimProfile= */ false,
-        /* enableActionCountProfile= */ false);
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ false);
     profiler.logSimpleTask(10000, 20000, ProfilerTask.VFS_STAT, "stat");
     // Unlike the VFS_STAT event above, the remote execution event will not be recorded since we
     // don't record the slowest remote exec events (see ProfilerTask.java).
@@ -358,7 +362,8 @@
         BlazeClock.instance().nanoTime(),
         /* enabledCpuUsageProfiling= */ false,
         /* slimProfile= */ false,
-        /* enableActionCountProfile= */ false);
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ false);
     profiler.logSimpleTask(10000, 20000, ProfilerTask.VFS_STAT, "stat");
 
     assertThat(ProfilerTask.VFS_STAT.collectsSlowestInstances()).isTrue();
@@ -539,7 +544,8 @@
         initialNanoTime,
         /* enabledCpuUsageProfiling= */ false,
         /* slimProfile= */ false,
-        /* enableActionCountProfile= */ false);
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ false);
     profiler.logSimpleTask(badClock.nanoTime(), ProfilerTask.INFO, "some task");
     profiler.stop();
   }
@@ -593,7 +599,8 @@
         BlazeClock.instance().nanoTime(),
         /* enabledCpuUsageProfiling= */ false,
         /* slimProfile= */ false,
-        /* enableActionCountProfile= */ false);
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ false);
     profiler.logSimpleTaskDuration(
         Profiler.nanoTimeMaybe(), Duration.ofSeconds(10), ProfilerTask.INFO, "foo");
     IOException expected = assertThrows(IOException.class, () -> profiler.stop());
@@ -619,13 +626,45 @@
         BlazeClock.instance().nanoTime(),
         /* enabledCpuUsageProfiling= */ false,
         /* slimProfile= */ false,
-        /* enableActionCountProfile= */ false);
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ false);
     profiler.logSimpleTaskDuration(
         Profiler.nanoTimeMaybe(), Duration.ofSeconds(10), ProfilerTask.INFO, "foo");
     IOException expected = assertThrows(IOException.class, () -> profiler.stop());
     assertThat(expected).hasMessageThat().isEqualTo("Expected failure.");
   }
 
+  @Test
+  public void testPrimaryOutputForAction() throws Exception {
+    ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+
+    profiler.start(
+        getAllProfilerTasks(),
+        buffer,
+        JSON_TRACE_FILE_FORMAT,
+        "dummy_output_base",
+        UUID.randomUUID(),
+        true,
+        clock,
+        clock.nanoTime(),
+        /* enabledCpuUsageProfiling= */ false,
+        /* slimProfile= */ false,
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ true);
+    try (SilentCloseable c = profiler.profileAction(ProfilerTask.ACTION, "test", "foo.out")) {
+      profiler.logEvent(ProfilerTask.PHASE, "event1");
+    }
+    profiler.stop();
+
+    JsonProfile jsonProfile = new JsonProfile(new ByteArrayInputStream(buffer.toByteArray()));
+
+    assertThat(
+            jsonProfile.getTraceEvents().stream()
+                .filter(traceEvent -> "foo.out".equals(traceEvent.primaryOutputPath()))
+                .collect(Collectors.toList()))
+        .hasSize(1);
+  }
+
   private ByteArrayOutputStream getJsonProfileOutputStream(boolean slimProfile) throws IOException {
     ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
     profiler.start(
@@ -639,7 +678,8 @@
         BlazeClock.instance().nanoTime(),
         /* enabledCpuUsageProfiling= */ false,
         slimProfile,
-        /* enableActionCountProfile= */ false);
+        /* enableActionCountProfile= */ false,
+        /* includePrimaryOutput= */ false);
     long curTime = Profiler.nanoTimeMaybe();
     for (int i = 0; i < 100_000; i++) {
       Duration duration;
diff --git a/src/test/shell/integration/aquery_test.sh b/src/test/shell/integration/aquery_test.sh
index 53d15e9..6dcb959 100755
--- a/src/test/shell/integration/aquery_test.sh
+++ b/src/test/shell/integration/aquery_test.sh
@@ -1128,7 +1128,7 @@
 
   # Verify that paths are broken down to path fragments.
   assert_contains "path_fragments {" output
-
+  assert_contains "primary_output_id" output
   # Verify that the appropriate action was included.
   assert_contains "label: \"dummy.txt\"" output
   assert_contains "mnemonic: \"Genrule\"" output