[WIP] Tag Bazel sandbox cgroups with network class for firewall restrictions

src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java

@@ -300,4 +300,7 @@ public enum WorkerProtocolFormat {
   /** Requires the execution service do NOT share caches across different workspace. */
   public static final String DIFFERENTIATE_WORKSPACE_CACHE =
       "internal-differentiate-workspace-cache";
+
+  /** Disables cgroups for a spawn */
+  public static final String NO_SUPPORTS_CGROUPS = "no-supports-cgroups";
 }

src/main/java/com/google/devtools/build/lib/analysis/BUILD

@@ -1616,13 +1616,15 @@ java_library(
         ":config/invalid_configuration_exception",
         ":config/run_under",
         ":platform_options",
+        ":test/test_configuration",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/buildeventstream",
         "//src/main/java/com/google/devtools/build/lib/buildeventstream/proto:build_event_stream_java_proto",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/events",
+        "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/skyframe:build_configuration",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi",

src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfigurationValue.java

@@ -28,6 +28,7 @@
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.analysis.PlatformOptions;
 import com.google.devtools.build.lib.analysis.config.OutputDirectories.InvalidMnemonicException;
+import com.google.devtools.build.lib.analysis.test.TestConfiguration;
 import com.google.devtools.build.lib.buildeventstream.BuildEventIdUtil;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos.BuildEventId;
@@ -36,6 +37,7 @@
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.events.EventHandler;
+import com.google.devtools.build.lib.packages.TestSize;
 import com.google.devtools.build.lib.skyframe.BuildConfigurationKey;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.starlarkbuildapi.BuildConfigurationApi;
@@ -901,4 +903,16 @@ private BuildConfigurationEvent createBuildEvent() {
   public ImmutableSet<String> getReservedActionMnemonics() {
     return reservedActionMnemonics;
   }
+
+  public Map<String, Float> getTestResources(TestSize size) {
+    if (!buildOptions.contains(TestConfiguration.TestOptions.class)) {
+      return ImmutableMap.of();
+    }
+    return buildOptions
+        .get(TestConfiguration.TestOptions.class)
+        .testResources
+        .stream()
+        .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, e -> e.getValue().get(size)));
+  }
+
 }

src/main/java/com/google/devtools/build/lib/analysis/test/TestConfiguration.java

@@ -14,10 +14,13 @@
 
 package com.google.devtools.build.lib.analysis.test;
 
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.analysis.OptionsDiffPredicate;
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
 import com.google.devtools.build.lib.analysis.config.CoreOptionConverters.LabelConverter;
@@ -28,8 +31,10 @@
 import com.google.devtools.build.lib.analysis.test.CoverageConfiguration.CoverageOptions;
 import com.google.devtools.build.lib.analysis.test.TestShardingStrategy.ShardingStrategyConverter;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.packages.TestSize;
 import com.google.devtools.build.lib.packages.TestTimeout;
 import com.google.devtools.build.lib.util.RegexFilter;
+import com.google.devtools.common.options.Converter;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.OptionDefinition;
 import com.google.devtools.common.options.OptionDocumentationCategory;
@@ -39,7 +44,9 @@
 import com.google.devtools.common.options.OptionsParsingException;
 import com.google.devtools.common.options.TriState;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
 
@@ -84,6 +91,23 @@ public static class TestOptions extends FragmentOptions {
                 + "-1 tells blaze to use its default timeouts for that category.")
     public Map<TestTimeout, Duration> testTimeout;
 
+    @Option(
+        name = "experimental_test_resources",
+        defaultValue = "null",
+        converter = TestResourcesConverter.class,
+        allowMultiple = true,
+        documentationCategory = OptionDocumentationCategory.TESTING,
+        effectTags = {OptionEffectTag.UNKNOWN},
+        help =
+            "Override the default resources amount for tests. The expected format is"
+                + " <resource>=<value>. If a single positive float is specified as <value>"
+                + " it will override the default resources for all test sizes. If 4"
+                + " comma-separated flats are specified, they will override the resource"
+                + " amount for small, medium, large, enourmous (in that order)."
+                + " Multiple resources can be separated by colon and multiple usages"
+                + " are accumulated.")
+    public List<Map.Entry<String, Map<TestSize, Float>>> testResources;
+
     @Option(
       name = "test_filter",
       allowMultiple = false,
@@ -450,4 +474,56 @@ public String getTypeDescription() {
       return "a positive integer or test_regex@runs. This flag may be passed more than once";
     }
   }
+
+  public static class TestResourcesConverter extends Converter.Contextless<List<Map.Entry<String, Map<TestSize, Float>>>> {
+    @Override
+    public String getTypeDescription() {
+      return "a resource name followed by equal and 1 float or 4 float, e.g memory=10,30,60,100";
+    }
+
+    @Override
+    public List<Map.Entry<String, Map<TestSize, Float>>> convert(String input) throws OptionsParsingException {
+      ImmutableList.Builder<Map.Entry<String, Map<TestSize, Float>>> resources = ImmutableList.builder();
+      Map<String, String> requests;
+      char resourcesSep = ':', valueSep = '=';
+
+      try {
+        requests = Splitter.on(resourcesSep).withKeyValueSeparator(valueSep).split(input);
+      } catch (IllegalArgumentException e) {
+        String message = String.format(
+            "%s. Separate resources from each other with '%c' and resources"
+              + " from their values with '%c', e.g. memory=10,20,50,100:cpu=1",
+            e.getMessage(),
+            resourcesSep,
+            valueSep);
+        throw new OptionsParsingException(message, e);
+      }
+      for (Map.Entry<String, String> request:  requests.entrySet()) {
+        List<Float> values = new ArrayList<>();
+        for (String token: Splitter.on(",").omitEmptyStrings().limit(5).split(request.getValue())) {
+          try {
+            values.add(Float.parseFloat(token));
+          } catch (NumberFormatException e) {
+            throw new OptionsParsingException("'" + token + "' is not a float", e);
+          }
+        }
+        EnumMap<TestSize, Float> amounts = Maps.newEnumMap(TestSize.class);
+        if (values.size() == 1) {
+          amounts.put(TestSize.SMALL, values.get(0));
+          amounts.put(TestSize.MEDIUM, values.get(0));
+          amounts.put(TestSize.LARGE, values.get(0));
+          amounts.put(TestSize.ENORMOUS, values.get(0));
+        } else if (values.size() == 4) {
+          amounts.put(TestSize.SMALL, values.get(0));
+          amounts.put(TestSize.MEDIUM, values.get(1));
+          amounts.put(TestSize.LARGE, values.get(2));
+          amounts.put(TestSize.ENORMOUS, values.get(3));
+        } else {
+          throw new OptionsParsingException("Invalid number of comma-separated entries");
+        }
+        resources.add(Maps.immutableEntry(request.getKey(), amounts));
+      }
+      return resources.build();
+    }
+  }
 }

src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java

@@ -16,6 +16,7 @@
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.actions.ExecutionRequirements;
 import com.google.devtools.build.lib.actions.ExecutionRequirements.ParseableRequirement.ValidationException;
@@ -34,6 +35,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * Container for test target properties available to the
@@ -88,9 +90,36 @@ private static ResourceSet getResourceSetFromSize(TestSize size) {
     isFlaky = ruleContext.attributes().get("flaky", Type.BOOLEAN);
     isExternal = TargetUtils.isExternalTestRule(rule);
 
+
     Map<String, String> executionInfo = Maps.newLinkedHashMap();
     executionInfo.putAll(TargetUtils.getExecutionInfo(rule));
 
+    ImmutableSet<String> requestedResources = executionInfo.keySet()
+        .stream()
+        .map(tag -> {
+          try {
+            String value = ExecutionRequirements.RESOURCES.parseIfMatches(tag);
+            if (value != null) {
+              return value.substring(0, value.indexOf(":"));
+            } else if (ExecutionRequirements.CPU.parseIfMatches(tag) != null) {
+              return "cpu";
+            }
+          } catch (ValidationException e) {
+            return null;
+          }
+          return null;
+        })
+        .filter(Objects::nonNull)
+        .collect(ImmutableSet.toImmutableSet());
+
+    Map<String, Float> testResources = ruleContext.getConfiguration().getTestResources(size);
+    for (Map.Entry<String, Float> request: testResources.entrySet()) {
+      if (requestedResources.contains(request.getKey())) {
+        continue;
+      }
+      executionInfo.put(String.format("resources:%s:%f", request.getKey(), request.getValue()), "");
+    }
+
     boolean incompatibleExclusiveTestSandboxed = false;
 
     TestConfiguration testConfiguration = ruleContext.getFragment(TestConfiguration.class);
@@ -119,6 +148,7 @@ private static ResourceSet getResourceSetFromSize(TestSize size) {
       // This will overwrite whatever TargetUtils put there, which might be confusing.
       executionInfo.putAll(executionRequirements.getExecutionInfo());
     }
+
     ruleContext.getConfiguration().modifyExecutionInfo(executionInfo, TestRunnerAction.MNEMONIC);
     this.executionInfo = ImmutableMap.copyOf(executionInfo);
 

src/main/java/com/google/devtools/build/lib/sandbox/BUILD

@@ -31,6 +31,7 @@ java_library(
     deps = [
         "//src/main/java/com/google/devtools/build/lib/actions:localhost_capacity",
         "//src/main/java/com/google/devtools/build/lib/util",
+        "//src/main/java/com/google/devtools/build/lib/util:cpu_resource_converter",
         "//src/main/java/com/google/devtools/build/lib/util:ram_resource_converter",
         "//src/main/java/com/google/devtools/build/lib/util:resource_converter",
         "//src/main/java/com/google/devtools/build/lib/vfs",
@@ -236,6 +237,7 @@ java_library(
         ":sandboxfs_sandboxed_spawn",
         "//src/main/java/com/google/devtools/build/lib:runtime",
         "//src/main/java/com/google/devtools/build/lib/actions",
+        "//src/main/java/com/google/devtools/build/lib/actions:exec_exception",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/actions:execution_requirements",
         "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
@@ -247,6 +249,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/exec/local",
         "//src/main/java/com/google/devtools/build/lib/exec/local:options",
         "//src/main/java/com/google/devtools/build/lib/profiler",
+        "//src/main/java/com/google/devtools/build/lib/sandbox/cgroups",
         "//src/main/java/com/google/devtools/build/lib/shell",
         "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/main/java/com/google/devtools/build/lib/util:pair",

src/main/java/com/google/devtools/build/lib/sandbox/LinuxSandboxCommandLineBuilder.java

@@ -53,7 +53,7 @@ public class LinuxSandboxCommandLineBuilder {
   private boolean enablePseudoterminal = false;
   private boolean useDebugMode = false;
   private boolean sigintSendsSigterm = false;
-  private String cgroupsDir;
+  private Set<Path> cgroupsDirs = ImmutableSet.of();
 
   private LinuxSandboxCommandLineBuilder(Path linuxSandboxPath, List<String> commandArguments) {
     this.linuxSandboxPath = linuxSandboxPath;
@@ -214,8 +214,8 @@ public LinuxSandboxCommandLineBuilder setUseDebugMode(boolean useDebugMode) {
    * this directory, its parent directory, and the cgroup directory for the Bazel process.
    */
   @CanIgnoreReturnValue
-  public LinuxSandboxCommandLineBuilder setCgroupsDir(String cgroupsDir) {
-    this.cgroupsDir = cgroupsDir;
+  public LinuxSandboxCommandLineBuilder setCgroupsDirs(Set<Path> cgroupsDirs) {
+    this.cgroupsDirs = cgroupsDirs;
     return this;
   }
 
@@ -296,8 +296,8 @@ public ImmutableList<String> build() {
     if (persistentProcess) {
       commandLineBuilder.add("-p");
     }
-    if (cgroupsDir != null) {
-      commandLineBuilder.add("-C", cgroupsDir);
+    for (Path dir: cgroupsDirs) {
+      commandLineBuilder.add("-C", dir.toString());
     }
     commandLineBuilder.add("--");
     commandLineBuilder.addAll(commandArguments);