Improve cquery speed by caching Skyframe calls

Calling Skyframe adds a performance cost even when hitting the Skyframe
cache. This change adds a local cache for retrieving configurations.

Experiments show about a 7% reduction in query time on `--output=graph` (see code comments).

Followup to
https://github.com/bazelbuild/bazel/pull/12248/files/939e4cadf7a49a316017d207bd4504ac0c10373c#r503552367.

Closes #12309.

Change-Id: I5dad55b3bb3b3e245a90f4d75e740c63754ff61c
PiperOrigin-RevId: 338255283

src/main/java/com/google/devtools/build/lib/query2/cquery/CqueryThreadsafeCallback.java

@@ -20,6 +20,7 @@
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.query2.NamedThreadSafeOutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.TargetAccessor;
+import com.google.devtools.build.lib.skyframe.BuildConfigurationValue;
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor;
 import java.io.BufferedWriter;
 import java.io.IOException;
@@ -28,6 +29,8 @@
 import java.io.Writer;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import javax.annotation.Nullable;
 
 /**
@@ -45,7 +48,11 @@ public abstract class CqueryThreadsafeCallback
   protected final CqueryOptions options;
   protected OutputStream outputStream;
   protected Writer printStream;
-  protected final SkyframeExecutor skyframeExecutor;
+  // Skyframe calls incur a performance cost, even on cache hits. Consider this before exposing
+  // direct executor access to child classes.
+  private final SkyframeExecutor skyframeExecutor;
+  private final Map<BuildConfigurationValue.Key, BuildConfiguration> configCache =
+      new ConcurrentHashMap<>();
   protected final ConfiguredTargetAccessor accessor;
 
   private final List<String> result = new ArrayList<>();
@@ -88,6 +95,17 @@ public void close(boolean failFast) throws InterruptedException, IOException {
     }
   }
 
+  protected BuildConfiguration getConfiguration(BuildConfigurationValue.Key configKey) {
+    // Experiments querying:
+    //     cquery --output=graph "deps(//src:main/java/com/google/devtools/build/lib:runtime)"
+    // 10 times on a warm Blaze instance show 7% less total query time when using this cache vs.
+    // calling Skyframe directly (and relying on Skyframe's cache).
+    if (configKey == null) {
+      return null;
+    }
+    return configCache.computeIfAbsent(
+        configKey, key -> skyframeExecutor.getConfiguration(eventHandler, key));
+  }
   /**
    * Returns a user-friendly configuration identifier as a prefix of <code>fullId</code>.
    *

src/main/java/com/google/devtools/build/lib/query2/cquery/GraphOutputFormatterCallback.java

@@ -61,13 +61,7 @@ public String getLabel(Node<ConfiguredTarget> node) {
           // hashes.
           ConfiguredTarget ct = node.getLabel();
           return String.format(
-              "%s (%s)",
-              ct.getLabel(),
-              // TODO(gregce): Even if getConfiguration is a cache hit this has overhead, especially
-              // when called many times on the same configuration in the same query. Investigate the
-              // performance impact and apply a cache if measurements justify it (also in other
-              // callbacks that do this).
-              shortId(skyframeExecutor.getConfiguration(eventHandler, ct.getConfigurationKey())));
+              "%s (%s)", ct.getLabel(), shortId(getConfiguration(ct.getConfigurationKey())));
         }
 
         @Override

src/main/java/com/google/devtools/build/lib/query2/cquery/LabelAndConfigurationOutputFormatterCallback.java

@@ -15,7 +15,6 @@
 
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.RequiredConfigFragmentsProvider;
-import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.config.CoreOptions.IncludeConfigFragmentsEnum;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.packages.Target;
@@ -46,8 +45,6 @@ public String getName() {
   @Override
   public void processOutput(Iterable<ConfiguredTarget> partialResult) {
     for (ConfiguredTarget configuredTarget : partialResult) {
-      BuildConfiguration config =
-          skyframeExecutor.getConfiguration(eventHandler, configuredTarget.getConfigurationKey());
       StringBuilder output = new StringBuilder();
       if (showKind) {
         Target actualTarget = accessor.getTargetFromConfiguredTarget(configuredTarget);
@@ -57,7 +54,7 @@ public void processOutput(Iterable<ConfiguredTarget> partialResult) {
           output
               .append(configuredTarget.getOriginalLabel())
               .append(" (")
-              .append(shortId(config))
+              .append(shortId(getConfiguration(configuredTarget.getConfigurationKey())))
               .append(")");
 
       if (options.showRequiredConfigFragments != IncludeConfigFragmentsEnum.OFF) {

src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java

@@ -63,6 +63,7 @@ public String formatName() {
 
   private final OutputType outputType;
   private final AspectResolver resolver;
+  private final SkyframeExecutor skyframeExecutor;
   private final JsonFormat.Printer jsonPrinter = JsonFormat.printer();
 
   private AnalysisProtos.CqueryResult.Builder protoResult;
@@ -79,6 +80,7 @@ public String formatName() {
       OutputType outputType) {
     super(eventHandler, options, out, skyframeExecutor, accessor);
     this.outputType = outputType;
+    this.skyframeExecutor = skyframeExecutor;
     this.resolver = resolver;
   }
 

src/main/java/com/google/devtools/build/lib/query2/cquery/StarlarkOutputFormatterCallback.java

@@ -63,8 +63,7 @@ private class CqueryDialectGlobals {
           @Param(name = "target"),
         })
     public Object buildOptions(ConfiguredTarget target) {
-      BuildConfiguration config =
-          skyframeExecutor.getConfiguration(eventHandler, target.getConfigurationKey());
+      BuildConfiguration config = getConfiguration(target.getConfigurationKey());
 
       if (config == null) {
         // config is null for input file configured targets.

src/main/java/com/google/devtools/build/lib/query2/cquery/TransitionsOutputFormatterCallback.java

@@ -106,9 +106,7 @@ public void processOutput(Iterable<ConfiguredTarget> partialResult) throws Inter
                 ct.getOriginalLabel(), accessor.getTargetFromConfiguredTarget(ct)));
     for (ConfiguredTarget configuredTarget : partialResult) {
       Target target = partialResultMap.get(configuredTarget.getOriginalLabel());
-      BuildConfiguration config =
-          skyframeExecutor.getConfiguration(
-              eventHandler, configuredTarget.getConfigurationKey());
+      BuildConfiguration config = getConfiguration(configuredTarget.getConfigurationKey());
       addResult(
           getRuleClassTransition(configuredTarget, target)
               + String.format("%s (%s)", configuredTarget.getOriginalLabel(), shortId(config)));