Cache all previously computed values in SkyQueryEnvironment#beforeEvaluateQuery if possible to save on latency for small queries.

This assumes that if the graph is up to date, then the data in SkyQueryEnvironment is also up to date. It also assumes that a ForkJoinPool remains usable until #shutdownNow is called.

--
MOS_MIGRATED_REVID=139386363

9 files changed, 77 insertions, 29 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
index 1ed21a08c7..ad51cd0ccc 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
@@ -136,6 +136,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   private MultisetSemaphore<PackageIdentifier> packageSemaphore;
   protected WalkableGraph graph;
   private InterruptibleSupplier<ImmutableSet<PathFragment>> blacklistPatternsSupplier;
+  private GraphBackedRecursivePackageProvider graphBackedRecursivePackageProvider;
   private ForkJoinPool forkJoinPool;
   private RecursivePackageProviderBackedTargetPatternResolver resolver;
   private final SkyKey universeKey;
@@ -203,26 +204,38 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   }
 
   private void beforeEvaluateQuery() throws InterruptedException {
-    EvaluationResult<SkyValue> result;
-    try (AutoProfiler p = AutoProfiler.logged("evaluation and walkable graph", LOG)) {
-      result = graphFactory.prepareAndGet(universeKey, loadingPhaseThreads, eventHandler);
-    }
-    checkEvaluationResult(result);
+    boolean resolverNeedsRecreation = false;
+    if (graph == null || !graphFactory.isUpToDate(universeKey)) {
+      // If this environment is uninitialized or the graph factory needs to evaluate, do so. We
+      // assume here that this environment cannot be initialized-but-stale if the factory is up
+      // to date.
+      EvaluationResult<SkyValue> result;
+      try (AutoProfiler p = AutoProfiler.logged("evaluation and walkable graph", LOG)) {
+        result = graphFactory.prepareAndGet(universeKey, loadingPhaseThreads, eventHandler);
+      }
+      checkEvaluationResult(result);
 
-    packageSemaphore = makeFreshPackageMultisetSemaphore();
-    graph = result.getWalkableGraph();
-    blacklistPatternsSupplier = InterruptibleSupplier.Memoize.of(new BlacklistSupplier(graph));
+      packageSemaphore = makeFreshPackageMultisetSemaphore();
+      graph = result.getWalkableGraph();
+      blacklistPatternsSupplier = InterruptibleSupplier.Memoize.of(new BlacklistSupplier(graph));
 
-    GraphBackedRecursivePackageProvider graphBackedRecursivePackageProvider =
-        new GraphBackedRecursivePackageProvider(graph, universeTargetPatternKeys, pkgPath);
-    forkJoinPool =
-        NamedForkJoinPool.newNamedPool("QueryEnvironment", queryEvaluationParallelismLevel);
-    resolver =
-        new RecursivePackageProviderBackedTargetPatternResolver(
-            graphBackedRecursivePackageProvider,
-            eventHandler,
-            TargetPatternEvaluator.DEFAULT_FILTERING_POLICY,
-            packageSemaphore);
+      graphBackedRecursivePackageProvider =
+          new GraphBackedRecursivePackageProvider(graph, universeTargetPatternKeys, pkgPath);
+      resolverNeedsRecreation = true;
+    }
+    if (forkJoinPool == null) {
+      forkJoinPool =
+          NamedForkJoinPool.newNamedPool("QueryEnvironment", queryEvaluationParallelismLevel);
+      resolverNeedsRecreation = true;
+    }
+    if (resolverNeedsRecreation) {
+      resolver =
+          new RecursivePackageProviderBackedTargetPatternResolver(
+              graphBackedRecursivePackageProvider,
+              eventHandler,
+              TargetPatternEvaluator.DEFAULT_FILTERING_POLICY,
+              packageSemaphore);
+    }
   }
 
   protected MultisetSemaphore<PackageIdentifier> makeFreshPackageMultisetSemaphore() {
@@ -319,14 +332,21 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   protected void evalTopLevelInternal(
       QueryExpression expr, OutputFormatterCallback<Target> callback)
           throws QueryException, InterruptedException {
+    boolean poolNeedsShutdown = true;
     try {
       super.evalTopLevelInternal(expr, callback);
+      poolNeedsShutdown = false;
     } finally {
-      // Force termination of remaining tasks - if evaluateQuery was successful there should be
-      // none, if it failed abruptly (e.g. was interrupted) we don't want to leave any dangling
-      // threads running tasks.
-      forkJoinPool.shutdownNow();
+      if (poolNeedsShutdown) {
+        // Force termination of remaining tasks if evaluation failed abruptly (e.g. was
+        // interrupted). We don't want to leave any dangling threads running tasks.
+        forkJoinPool.shutdownNow();
+      }
       forkJoinPool.awaitQuiescence(Long.MAX_VALUE, TimeUnit.MILLISECONDS);
+      if (poolNeedsShutdown) {
+        // Signal that pool must be recreated on the next invocation.
+        forkJoinPool = null;
+      }
     }
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValue.java
index d572058eac..8901745c00 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValue.java
@@ -58,7 +58,7 @@ public final class PrecomputedValue implements SkyValue {
     }
 
     void inject(Injectable injectable) {
-      injectable.inject(ImmutableMap.of(precomputed.key, new PrecomputedValue(supplier.get())));
+      injectable.inject(precomputed.key, new PrecomputedValue(supplier.get()));
     }
   }
 
@@ -179,7 +179,7 @@ public final class PrecomputedValue implements SkyValue {
      * Injects a new variable value.
      */
     public void set(Injectable injectable, T value) {
-      injectable.inject(ImmutableMap.of(key, new PrecomputedValue(value)));
+      injectable.inject(key, new PrecomputedValue(value));
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutor.java
index f64ed90fae..55fac95b51 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutor.java
@@ -234,7 +234,7 @@ public final class SequencedSkyframeExecutor extends SkyframeExecutor {
   }
 
   @Override
-  protected BuildDriver newBuildDriver() {
+  protected BuildDriver getBuildDriver() {
     return new SequentialBuildDriver(memoizingEvaluator);
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
index 2c38f6d924..39090f75a9 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
@@ -565,7 +565,7 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
     memoizingEvaluator = evaluatorSupplier.create(
         skyFunctions, evaluatorDiffer(), progressReceiver, emittedEventState,
         hasIncrementalState());
-    buildDriver = newBuildDriver();
+    buildDriver = getBuildDriver();
   }
 
   protected SkyframeProgressReceiver newSkyframeProgressReceiver() {
@@ -588,7 +588,7 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
 
   protected abstract Differencer evaluatorDiffer();
 
-  protected abstract BuildDriver newBuildDriver();
+  protected abstract BuildDriver getBuildDriver();
 
   /**
    * Values whose values are known at startup and guaranteed constant are still wiped from the
@@ -1546,6 +1546,11 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
     return evaluationResult;
   }
 
+  @Override
+  public boolean isUpToDate(SkyKey universeKey) {
+    return buildDriver.alreadyEvaluated(ImmutableList.of(universeKey));
+  }
+
   /**
    * Get metadata related to the prepareAndGet() lookup. Resulting data is specific to the
    * underlying evaluation implementation.
diff --git a/src/main/java/com/google/devtools/build/skyframe/BuildDriver.java b/src/main/java/com/google/devtools/build/skyframe/BuildDriver.java
index 4ebe244ea0..654dfaead2 100644
--- a/src/main/java/com/google/devtools/build/skyframe/BuildDriver.java
+++ b/src/main/java/com/google/devtools/build/skyframe/BuildDriver.java
@@ -36,6 +36,12 @@ public interface BuildDriver {
   String meta(Iterable<SkyKey> roots, OptionsClassProvider options)
       throws AbruptExitException, InterruptedException;
 
+  /**
+   * Returns true if this {@link BuildDriver} instance has already been used to {@link #evaluate}
+   * the given {@code roots} at the Version that would be passed along to the next call to {@link
+   * MemoizingEvaluator#evaluate} in {@link #evaluate}.
+   */
+  boolean alreadyEvaluated(Iterable<SkyKey> roots);
 
   MemoizingEvaluator getGraphForTesting();
 
diff --git a/src/main/java/com/google/devtools/build/skyframe/Injectable.java b/src/main/java/com/google/devtools/build/skyframe/Injectable.java
index 99d6687caf..9c5032c723 100644
--- a/src/main/java/com/google/devtools/build/skyframe/Injectable.java
+++ b/src/main/java/com/google/devtools/build/skyframe/Injectable.java
@@ -20,4 +20,5 @@ import java.util.Map;
  */
 public interface Injectable {
   void inject(Map<SkyKey, ? extends SkyValue> values);
+  void inject(SkyKey key, SkyValue value);
 }
diff --git a/src/main/java/com/google/devtools/build/skyframe/RecordingDifferencer.java b/src/main/java/com/google/devtools/build/skyframe/RecordingDifferencer.java
index 4d25e0e65a..a4d768d043 100644
--- a/src/main/java/com/google/devtools/build/skyframe/RecordingDifferencer.java
+++ b/src/main/java/com/google/devtools/build/skyframe/RecordingDifferencer.java
@@ -16,7 +16,6 @@ package com.google.devtools.build.skyframe;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.concurrent.ThreadSafety;
-
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -73,4 +72,9 @@ public class RecordingDifferencer implements Differencer, Injectable {
   public void inject(Map<SkyKey, ? extends SkyValue> values) {
     valuesToInject.putAll(values);
   }
+
+  @Override
+  public void inject(SkyKey key, SkyValue value) {
+    valuesToInject.put(key, value);
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/skyframe/SequentialBuildDriver.java b/src/main/java/com/google/devtools/build/skyframe/SequentialBuildDriver.java
index 905c25fc65..f28a800f15 100644
--- a/src/main/java/com/google/devtools/build/skyframe/SequentialBuildDriver.java
+++ b/src/main/java/com/google/devtools/build/skyframe/SequentialBuildDriver.java
@@ -46,7 +46,12 @@ public class SequentialBuildDriver implements BuildDriver {
    return "";
  }
 
- @Override
+  @Override
+  public boolean alreadyEvaluated(Iterable<SkyKey> roots) {
+    return false;
+  }
+
+  @Override
   public MemoizingEvaluator getGraphForTesting() {
     return memoizingEvaluator;
   }
diff --git a/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java b/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java
index 50de277be9..3d7796b3cc 100644
--- a/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java
+++ b/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java
@@ -95,6 +95,13 @@ public interface WalkableGraph {
     EvaluationResult<SkyValue> prepareAndGet(
         SkyKey universeKey, int numThreads, EventHandler eventHandler) throws InterruptedException;
 
+    /**
+     * Returns true if this instance has already been used to {@link #prepareAndGet} {@code
+     * universeKey}. If so, cached results from {@link #prepareAndGet} can be re-used safely,
+     * potentially saving some processing time.
+     */
+    boolean isUpToDate(SkyKey universeKey);
+
     /** Returns the {@link SkyKey} that defines this universe. */
     SkyKey getUniverseKey(Collection<String> roots, String offset);
   }