feat: add a helper for rules to work with resource_sets (#792)

* feat: add a helper for rules to work with resource_sets

This API is poorly designed and needs some help to let rule users pick a value in cases that they aren't also the rule author

* chore: add some cpu resource_set values as well

docs/BUILD.bazel

@@ -157,4 +157,9 @@ stardoc_with_diff_test(
     bzl_library_target = "//lib:bats",
 )
 
+stardoc_with_diff_test(
+    name = "resource_sets",
+    bzl_library_target = "//lib:resource_sets",
+)
+
 update_docs()

docs/resource_sets.md

@@ -0,0 +1,31 @@
+<!-- Generated with Stardoc: http://skydoc.bazel.build -->
+
+Utilities for rules that expose resource_set on ctx.actions.run[_shell]
+
+Workaround for https://github.com/bazelbuild/bazel/issues/15187
+
+Note, this workaround only provides some fixed values for either CPU or Memory.
+
+Rule authors who are ALSO the BUILD author might know better, and can
+write custom resource_set functions for use within their own repository.
+This seems to be the use case that Google engineers imagined.
+
+
+<a id="resource_set"></a>
+
+## resource_set
+
+<pre>
+resource_set(<a href="#resource_set-attr">attr</a>)
+</pre>
+
+
+
+**PARAMETERS**
+
+
+| Name  | Description | Default Value |
+| :------------- | :------------- | :------------- |
+| <a id="resource_set-attr"></a>attr |  <p align="center"> - </p>   |  none |
+
+

lib/BUILD.bazel

@@ -307,3 +307,8 @@ bzl_library(
     srcs = ["bats.bzl"],
     deps = ["//lib/private:bats"],
 )
+
+bzl_library(
+    name = "resource_sets",
+    srcs = ["resource_sets.bzl"],
+)

lib/private/source_toolchains_repo.bzl

@@ -84,7 +84,7 @@ source_toolchains_repo = repository_rule(
     attrs = {
         "toolchain_type": attr.string(doc = "Label to the toolchain_type", mandatory = True),
         "toolchain_rule_load_from": attr.string(doc = "Label to the concrete toolchain rule to load from", mandatory = True),
-        "toolchain_rule": attr.string(doc = "Name of the concerete toolchain rule", mandatory = True),
+        "toolchain_rule": attr.string(doc = "Name of the concrete toolchain rule", mandatory = True),
         "binary": attr.string(doc = "Label to the binary", mandatory = True),
     },
 )

lib/resource_sets.bzl

@@ -0,0 +1,88 @@
+"""Utilities for rules that expose resource_set on ctx.actions.run[_shell]
+
+Workaround for https://github.com/bazelbuild/bazel/issues/15187
+
+Note, this workaround only provides some fixed values for either CPU or Memory.
+
+Rule authors who are ALSO the BUILD author might know better, and can
+write custom resource_set functions for use within their own repository.
+This seems to be the use case that Google engineers imagined.
+"""
+
+resource_set_values = [
+    "cpu_2",
+    "cpu_4",
+    "default",
+    "mem_512m",
+    "mem_1g",
+    "mem_2g",
+    "mem_4g",
+    "mem_8g",
+    "mem_16g",
+    "mem_32g",
+]
+
+def _resource_set_cpu_2(_, __):
+    return {"cpu": 2}
+
+def _resource_set_cpu_4(_, __):
+    return {"cpu": 4}
+
+def _resource_set_mem_512m(_, __):
+    return {"memory": 512}
+
+def _resource_set_mem_1g(_, __):
+    return {"memory": 1024}
+
+def _resource_set_mem_2g(_, __):
+    return {"memory": 2048}
+
+def _resource_set_mem_4g(_, __):
+    return {"memory": 4096}
+
+def _resource_set_mem_8g(_, __):
+    return {"memory": 8192}
+
+def _resource_set_mem_16g(_, __):
+    return {"memory": 16384}
+
+def _resource_set_mem_32g(_, __):
+    return {"memory": 32768}
+
+# buildifier: disable=function-docstring
+def resource_set(attr):
+    if attr.resource_set == "cpu_2":
+        return _resource_set_cpu_2
+    if attr.resource_set == "cpu_4":
+        return _resource_set_cpu_4
+    if attr.resource_set == "default":
+        return None
+    if attr.resource_set == "mem_512m":
+        return _resource_set_mem_512m
+    if attr.resource_set == "mem_1g":
+        return _resource_set_mem_1g
+    if attr.resource_set == "mem_2g":
+        return _resource_set_mem_2g
+    if attr.resource_set == "mem_4g":
+        return _resource_set_mem_4g
+    if attr.resource_set == "mem_8g":
+        return _resource_set_mem_8g
+    if attr.resource_set == "mem_16g":
+        return _resource_set_mem_16g
+    if attr.resource_set == "mem_32g":
+        return _resource_set_mem_32g
+    fail("unknown resource set", attr.resource_set)
+
+resource_set_attr = {
+    "resource_set": attr.string(
+        doc = """A predefined function used as the resource_set for actions.
+
+        Used with --experimental_action_resource_set to reserve more RAM/CPU, preventing Bazel overscheduling resource-intensive actions.
+
+        By default, Bazel allocates 1 CPU and 250M of RAM.
+        https://github.com/bazelbuild/bazel/blob/058f943037e21710837eda9ca2f85b5f8538c8c5/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java#L77
+        """,
+        default = "default",
+        values = resource_set_values,
+    ),
+}