Update ToolchainResolutionFunction to consider multiple available

execution platforms.

Part of #4442.

Change-Id: I6678d57f4aaadcb19032bf58820606242ba66a25
PiperOrigin-RevId: 184707708

3 files changed, 91 insertions, 31 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunction.java
index 0711230f24..16b055bec0 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunction.java
@@ -14,8 +14,10 @@
 
 package com.google.devtools.build.lib.skyframe;
 
-import com.google.common.annotations.VisibleForTesting;
+import static java.util.stream.Collectors.joining;
+
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.PlatformConfiguration;
 import com.google.devtools.build.lib.analysis.PlatformOptions;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
@@ -34,6 +36,9 @@ import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyFunctionException;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 import javax.annotation.Nullable;
 
 /** {@link SkyFunction} which performs toolchain resolution for a class of rules. */
@@ -69,29 +74,39 @@ public class ToolchainResolutionFunction implements SkyFunction {
 
     // Find the right one.
     boolean debug = configuration.getOptions().get(PlatformOptions.class).toolchainResolutionDebug;
-    DeclaredToolchainInfo toolchain =
+    ImmutableMap<PlatformInfo, Label> resolvedToolchainLabels =
         resolveConstraints(
             key.toolchainType(),
-            key.execPlatform(),
+            key.availableExecutionPlatforms(),
             key.targetPlatform(),
             toolchains.registeredToolchains(),
             debug ? env.getListener() : null);
 
-    if (toolchain == null) {
+    if (resolvedToolchainLabels.isEmpty()) {
       throw new ToolchainResolutionFunctionException(
           new NoToolchainFoundException(key.toolchainType()));
     }
-    return ToolchainResolutionValue.create(toolchain.toolchainLabel());
+
+    return ToolchainResolutionValue.create(resolvedToolchainLabels);
   }
 
-  @VisibleForTesting
-  static DeclaredToolchainInfo resolveConstraints(
+  /**
+   * Given the available execution platforms and toolchains, find the set of platform, toolchain
+   * pairs that are compatible a) with each other, and b) with the toolchain type and target
+   * platform.
+   */
+  private static ImmutableMap<PlatformInfo, Label> resolveConstraints(
       Label toolchainType,
-      PlatformInfo execPlatform,
+      List<PlatformInfo> availableExecutionPlatforms,
       PlatformInfo targetPlatform,
       ImmutableList<DeclaredToolchainInfo> toolchains,
       @Nullable EventHandler eventHandler) {
 
+    // Platforms may exist multiple times in availableExecutionPlatforms. The Set lets this code
+    // check whether a platform has already been seen during processing.
+    Set<PlatformInfo> platformsSeen = new HashSet<>();
+    ImmutableMap.Builder<PlatformInfo, Label> builder = ImmutableMap.builder();
+
     debugMessage(eventHandler, "Looking for toolchain of type %s...", toolchainType);
     for (DeclaredToolchainInfo toolchain : toolchains) {
       // Make sure the type matches.
@@ -99,13 +114,8 @@ public class ToolchainResolutionFunction implements SkyFunction {
         continue;
       }
       debugMessage(eventHandler, "  Considering toolchain %s...", toolchain.toolchainLabel());
-      if (!checkConstraints(eventHandler, toolchain.execConstraints(), "execution", execPlatform)) {
-        debugMessage(
-            eventHandler,
-            "  Rejected toolchain %s, because of execution platform mismatch",
-            toolchain.toolchainLabel());
-        continue;
-      }
+
+      // Make sure the target platform matches.
       if (!checkConstraints(
           eventHandler, toolchain.targetConstraints(), "target", targetPlatform)) {
         debugMessage(
@@ -115,12 +125,36 @@ public class ToolchainResolutionFunction implements SkyFunction {
         continue;
       }
 
-      debugMessage(eventHandler, "  Selected toolchain %s", toolchain.toolchainLabel());
-      return toolchain;
+      // Find the matching execution platforms.
+      for (PlatformInfo executionPlatform : availableExecutionPlatforms) {
+        if (!checkConstraints(
+            eventHandler, toolchain.execConstraints(), "execution", executionPlatform)) {
+          continue;
+        }
+
+        // Only add the toolchains if this is a new platform.
+        if (!platformsSeen.contains(executionPlatform)) {
+          builder.put(executionPlatform, toolchain.toolchainLabel());
+          platformsSeen.add(executionPlatform);
+        }
+      }
     }
 
-    debugMessage(eventHandler, "  No toolchain found");
-    return null;
+    ImmutableMap<PlatformInfo, Label> resolvedToolchainLabels = builder.build();
+    if (resolvedToolchainLabels.isEmpty()) {
+      debugMessage(eventHandler, "  No toolchains found");
+    } else {
+      debugMessage(
+          eventHandler,
+          "  Selected execution platforms and toolchains: {%s}",
+          resolvedToolchainLabels
+              .entrySet()
+              .stream()
+              .map(e -> String.format("%s -> %s", e.getKey().label(), e.getValue()))
+              .collect(joining(", ")));
+    }
+
+    return resolvedToolchainLabels;
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionValue.java
index 01909e6e8e..1f39059adf 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionValue.java
@@ -15,14 +15,22 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import java.util.List;
 
-/** A value which represents the selected toolchain for a specific target and execution platform. */
+/**
+ * A value which represents the map of potential execution platforms and resolved toolchains. This
+ * value only considers a single toolchain type, which allows for a Skyframe cache per toolchain
+ * type. Callers will need to consider all toolchain types that are required and merge the results
+ * together appropriately.
+ */
 @AutoValue
 public abstract class ToolchainResolutionValue implements SkyValue {
 
@@ -31,9 +39,9 @@ public abstract class ToolchainResolutionValue implements SkyValue {
       BuildConfiguration configuration,
       Label toolchainType,
       PlatformInfo targetPlatform,
-      PlatformInfo execPlatform) {
+      List<PlatformInfo> availableExecutionPlatforms) {
     return ToolchainResolutionKey.create(
-        configuration, toolchainType, targetPlatform, execPlatform);
+        configuration, toolchainType, targetPlatform, availableExecutionPlatforms);
   }
 
   /** {@link SkyKey} implementation used for {@link ToolchainResolutionFunction}. */
@@ -50,21 +58,30 @@ public abstract class ToolchainResolutionValue implements SkyValue {
 
     public abstract PlatformInfo targetPlatform();
 
-    public abstract PlatformInfo execPlatform();
+    public abstract ImmutableList<PlatformInfo> availableExecutionPlatforms();
 
     public static ToolchainResolutionKey create(
         BuildConfiguration configuration,
         Label toolchainType,
         PlatformInfo targetPlatform,
-        PlatformInfo execPlatform) {
+        List<PlatformInfo> availableExecutionPlatforms) {
       return new AutoValue_ToolchainResolutionValue_ToolchainResolutionKey(
-          configuration, toolchainType, targetPlatform, execPlatform);
+          configuration,
+          toolchainType,
+          targetPlatform,
+          ImmutableList.copyOf(availableExecutionPlatforms));
     }
   }
 
-  public static ToolchainResolutionValue create(Label toolchainLabel) {
-    return new AutoValue_ToolchainResolutionValue(toolchainLabel);
+  public static ToolchainResolutionValue create(
+      ImmutableMap<PlatformInfo, Label> availableToolchainLabels) {
+    return new AutoValue_ToolchainResolutionValue(availableToolchainLabels);
   }
 
-  public abstract Label toolchainLabel();
+  /**
+   * Returns the resolved set of toolchain labels (as {@link Label}) for the requested toolchain
+   * type, keyed by the execution platforms (as {@link PlatformInfo}). Ordering is not preserved, if
+   * the caller cares about the order of platforms it must take care of that directly.
+   */
+  public abstract ImmutableMap<PlatformInfo, Label> availableToolchainLabels();
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainUtil.java b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainUtil.java
index cbfcbd4514..abc002ac57 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainUtil.java
@@ -18,6 +18,7 @@ import com.google.auto.value.AutoValue;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.PlatformConfiguration;
 import com.google.devtools.build.lib.analysis.ToolchainContext;
@@ -202,7 +203,7 @@ public class ToolchainUtil {
     for (Label toolchainType : requiredToolchains) {
       registeredToolchainKeys.add(
           ToolchainResolutionValue.key(
-              configuration, toolchainType, targetPlatform, executionPlatform));
+              configuration, toolchainType, targetPlatform, ImmutableList.of(executionPlatform)));
     }
 
     Map<
@@ -238,9 +239,17 @@ public class ToolchainUtil {
         if (valueOrException.get() == null) {
           valuesMissing = true;
         } else {
+          ToolchainResolutionValue toolchainResolutionValue =
+              (ToolchainResolutionValue) valueOrException.get();
+
+          // TODO(https://github.com/bazelbuild/bazel/issues/4442): Handle finding the best
+          // execution platform when multiple are available.
           Label toolchainLabel =
-              ((ToolchainResolutionValue) valueOrException.get()).toolchainLabel();
-          builder.put(requiredToolchainType, toolchainLabel);
+              Iterables.getFirst(
+                  toolchainResolutionValue.availableToolchainLabels().values(), null);
+          if (toolchainLabel != null) {
+            builder.put(requiredToolchainType, toolchainLabel);
+          }
         }
       } catch (NoToolchainFoundException e) {
         // Save the missing type and continue looping to check for more.