[rbe] Support --platform override/rewrite

Setting `rewrapper --platform` would clobber the value
that comes from `--cfg fuchsia-rewrapper.cfg`.
Now we can effectively dictionary-merge the platform value
from both the cfg and the command-line.
The --platform flag when passed through our remote_action.py
wrapper will be re-written with a merged value.

Other: move some code+test to cl_utils.py for better reuse.

Bug: 342692553
Change-Id: Ic00e8b448aa7979e4a580ac38fea8630453c5fa1
Reviewed-on: https://fuchsia-review.googlesource.com/c/fuchsia/+/1054556
Reviewed-by: David Turner <[email protected]>
Commit-Queue: David Fang <[email protected]>
Fuchsia-Auto-Submit: David Fang <[email protected]>

Author: fangism

build/rbe/action_diff.py

@@ -353,7 +353,7 @@ def main(argv: Sequence[str]) -> int:
     # len(logs) == 2, enforced by _MAIN_ARG_PARSER.
 
     cfg_path = remotetool._REPROXY_CFG
-    reproxy_cfg = remotetool.read_config_file_lines(
+    reproxy_cfg = cl_util.read_config_file_lines(
         cfg_path.read_text().splitlines()
     )
 

build/rbe/cl_utils.py

@@ -367,6 +367,51 @@ def expand_paths_from_files(files: Iterable[Path]) -> Iterable[Path]:
                         yield Path(p.replace(" ", "\\ "))  # preserve escape
 
 
+def read_config_file_lines(lines: Iterable[str]) -> Dict[str, str]:
+    """Parser for reading RBE config files.
+
+    RBE config files are text files with lines of "VAR=VALUE"
+    (ignoring whole-line #-comments and blank lines).
+    Spec details can be found at:
+    https://github.com/bazelbuild/reclient/blob/main/internal/pkg/rbeflag/rbeflag.go
+
+    Args:
+      lines: lines of config from a file
+
+    Returns:
+      dictionary of key-value pairs read from the config file.
+    """
+    result = {}
+    for line in lines:
+        stripped = line.strip()
+        if stripped and not stripped.startswith("#"):
+            key, sep, value = stripped.partition("=")
+            if sep == "=":
+                result[key] = value
+    return result
+
+
+def values_dict_to_config_value(
+    values: Dict[str, str], eq: str = "=", sep: str = ","
+) -> str:
+    """Return a string representation of a dictionary for config lines.
+
+    It is assumed that the order among different keys is inconsequential,
+    so we sort the output by keys for determinism.
+
+    Args:
+      values: arbitrary string dictionary of key-values
+      eq: string to join keys and values
+      sep: string to join key-value pairs
+
+    Returns:
+      Joined string representation, ordered by keys (deterministic).
+    """
+    # Keys are unique, so this effectively sorts (key, value) pairs
+    # by key without regard to value.
+    return sep.join(k + eq + v for k, v in sorted(values.items()))
+
+
 def keyed_flags_to_values_dict(
     flags: Iterable[str], convert_type: Callable[[str], Any] = None
 ) -> Dict[str, Sequence[str]]:

build/rbe/cl_utils_test.py

@@ -522,6 +522,55 @@ def test_fuse_repeated(self):
         )
 
 
+class ReadConfigFileLinesTests(unittest.TestCase):
+    def test_empty_file(self):
+        self.assertEqual(cl_utils.read_config_file_lines([]), dict())
+
+    def test_ignore_blank_lines(self):
+        self.assertEqual(
+            cl_utils.read_config_file_lines(["", "\t", "\n"]), dict()
+        )
+
+    def test_ignore_comments(self):
+        self.assertEqual(
+            cl_utils.read_config_file_lines(["####", "# comment"]), dict()
+        )
+
+    def test_ignore_non_key_value_pairs(self):
+        self.assertEqual(
+            cl_utils.read_config_file_lines(["value-only"]), dict()
+        )
+
+    def test_key_value(self):
+        self.assertEqual(
+            cl_utils.read_config_file_lines(["key=value"]), {"key": "value"}
+        )
+
+    def test_last_wins(self):
+        self.assertEqual(
+            cl_utils.read_config_file_lines(["key=value-1", "key=value-2"]),
+            {"key": "value-2"},
+        )
+
+
+class ValuesDictToConfigValueTests(unittest.TestCase):
+    def test_empty(self):
+        self.assertEqual(cl_utils.values_dict_to_config_value({}), "")
+
+    def test_one_value(self):
+        self.assertEqual(
+            cl_utils.values_dict_to_config_value({"a": "bb"}), "a=bb"
+        )
+
+    def test_bunch_of_values_must_be_key_sorted(self):
+        self.assertEqual(
+            cl_utils.values_dict_to_config_value(
+                {"a": "yy", "zz": "aa", "pp": "qqq"}
+            ),
+            "a=yy,pp=qqq,zz=aa",
+        )
+
+
 class KeyedFlagsToValuesDictTests(unittest.TestCase):
     def test_empty(self):
         self.assertEqual(

build/rbe/remote_action.py

@@ -909,6 +909,7 @@ def __init__(
         input_list_paths: Sequence[Path] = None,
         output_files: Sequence[Path] = None,
         output_dirs: Sequence[Path] = None,
+        platform: str = None,
         disable: bool = False,
         verbose: bool = False,
         save_temps: bool = False,
@@ -944,6 +945,7 @@ def __init__(
             current working dir.
           output_dirs: directories to be fetched after remote execution, relative to the
             current working dir.
+          platform: rewrapper --platform, containing remote execution parameters
           disable: if true, execute locally.
           verbose: if true, print more information about what is happening.
           compare_with_local: if true, also run locally and compare outputs.
@@ -1012,6 +1014,9 @@ def __init__(
         # (which can be verified with --compare [compare_with_local]).
         self._local_only_command = local_only_command or command
 
+        # platform is handled by specially, by merging with cfg.
+        self._platform = platform
+
         # Detect some known rewrapper options
         (
             self._rewrapper_known_options,
@@ -1170,6 +1175,14 @@ def _generate_options(self) -> Iterable[str]:
         if self.exec_strategy:
             yield f"--exec_strategy={self.exec_strategy}"
 
+        if self.platform:
+            # Then merge the value from --cfg and --platform to override
+            # the value from the --cfg.  Yield the rewritten flag.
+            platform_value = cl_utils.values_dict_to_config_value(
+                self.merged_platform
+            )
+            yield f"--platform={platform_value}"
+
         yield from self._options
 
     @property
@@ -1193,6 +1206,43 @@ def download_regex(self) -> Optional[str]:
     def download_outputs(self) -> bool:
         return self._rewrapper_known_options.download_outputs
 
+    @property
+    def platform(self) -> Optional[str]:
+        return self._platform
+
+    @property
+    def merged_platform(self) -> Dict[str, str]:
+        """Combined platform values from --cfg and --platform."""
+        merged_values: Dict[str, str] = {}
+
+        def take_dict_last_values(key_values: str) -> Dict[str, str]:
+            return {
+                k: cl_utils.last_value_or_default(v, "")
+                for k, v in cl_utils.keyed_flags_to_values_dict(
+                    key_values.split(",")
+                ).items()
+            }
+
+        cfg = self.config
+        if cfg:
+            rewrapper_cfg: Dict[str, str] = cl_utils.read_config_file_lines(
+                cfg.read_text().splitlines()
+            )
+            cfg_platform = rewrapper_cfg.get("platform", "")
+            if cfg_platform:
+                # in case of multiple/conflicting values, take the last one
+                merged_values.update(take_dict_last_values(cfg_platform))
+
+        # TODO(b/342692553): merge RBE_platform environment variable
+
+        cl_platform = self.platform
+        if cl_platform:
+            # override earlier values
+            # in case of multiple/conflicting values, take the last one
+            merged_values.update(take_dict_last_values(cl_platform))
+
+        return merged_values
+
     @property
     def need_download_stub_predicate(self) -> Callable[[Path], bool]:
         """Return a function that indicates whether an output path will require a download stub."""
@@ -2346,6 +2396,11 @@ def inherit_main_arg_parser_flags(
         metavar="DIRS",
         help="Specify additional remote output directories, comma-separated, relative to the current working dir (repeatable, cumulative).  Note: This is different than `rewrapper --output_directories`, which expects exec_root-relative paths.",
     )
+    rewrapper_group.add_argument(
+        "--platform",
+        type=str,
+        help="The rewrapper platform variable specifies remote execution parameters, such as OS type, image, worker constraints, etc.  The value specified in this wrapper is intercepted and merged with the value found in the --cfg.",
+    )
 
     main_group = parser.add_argument_group(
         title="common",
@@ -2504,6 +2559,7 @@ def remote_action_from_args(
         input_list_paths=input_list_paths,
         output_files=output_files,
         output_dirs=output_dirs,
+        platform=main_args.platform,
         disable=main_args.local,
         verbose=main_args.verbose,
         label=main_args.label,

build/rbe/remote_action_test.py

@@ -834,10 +834,56 @@ def test_cfg(self):
         cfg = Path("other.cfg")
         main_args, other = p.parse_known_args([f"--cfg={cfg}", "--", "echo"])
         self.assertEqual(main_args.cfg, cfg)
+        self.assertEqual(other, [])
         action = remote_action.remote_action_from_args(main_args)
+        self.assertEqual(action.config, cfg)
         self.assertEqual(action.local_only_command, ["echo"])
         self.assertEqual(action.options, ["--cfg", str(cfg)])
 
+    def test_platform_merge_override(self):
+        p = self._make_main_parser()
+        cfg = Path("other.cfg")
+        platform_value = "foo=zoo,alice=bob"
+        # Test both styles of flags.
+        test_flag_variants = (
+            [f"--platform={platform_value}"],
+            ["--platform", platform_value],
+        )
+        for flags in test_flag_variants:
+            main_args, remote_options = p.parse_known_args(
+                [f"--cfg={cfg}"] + flags + ["--", "echo"]
+            )
+            self.assertEqual(main_args.cfg, cfg)
+            self.assertEqual(main_args.platform, platform_value)
+            action = remote_action.remote_action_from_args(
+                main_args=main_args,
+                remote_options=remote_options,
+            )
+            self.assertEqual(action.config, cfg)
+            self.assertEqual(action.platform, platform_value)
+            self.assertEqual(action.local_only_command, ["echo"])
+            with mock.patch.object(
+                Path,
+                "read_text",
+                return_value="\n".join(
+                    [
+                        "parameter_this=1",
+                        "parameter_that=do_not_care",
+                        "platform=foo=bar,baz=quux",
+                    ]
+                ),
+            ) as mock_read_cfg:
+                self.assertEqual(
+                    action.options,
+                    [
+                        "--cfg",
+                        str(cfg),
+                        "--platform=alice=bob,baz=quux,foo=zoo",
+                    ],
+                )
+
+            mock_read_cfg.assert_called_once_with()
+
     def test_bindir(self):
         p = self._make_main_parser()
         bindir = Path("/usr/local/bin")

build/rbe/remotetool.py

@@ -230,30 +230,6 @@ def parse_show_action_output(lines: Iterable[str]) -> ShowActionResult:
     )
 
 
-def read_config_file_lines(lines: Iterable[str]) -> Dict[str, str]:
-    """Parser for reading RBE config files.
-
-    RBE config files are text files with lines of "VAR=VALUE"
-    (ignoring whole-line #-comments and blank lines).
-    Spec details can be found at:
-    https://github.com/bazelbuild/reclient/blob/main/internal/pkg/rbeflag/rbeflag.go
-
-    Args:
-      lines: lines of config from a file
-
-    Returns:
-      dictionary of key-value pairs read from the config file.
-    """
-    result = {}
-    for line in lines:
-        stripped = line.strip()
-        if stripped and not stripped.startswith("#"):
-            key, sep, value = stripped.partition("=")
-            if sep == "=":
-                result[key] = value
-    return result
-
-
 class RemoteTool(object):
     def __init__(self, reproxy_cfg: Path):
         self._reproxy_cfg = reproxy_cfg
@@ -409,7 +385,9 @@ def download_dir(
 
 
 def configure_remotetool(cfg: Path) -> RemoteTool:
-    return RemoteTool(read_config_file_lines(cfg.read_text().splitlines()))
+    return RemoteTool(
+        cl_utils.read_config_file_lines(cfg.read_text().splitlines())
+    )
 
 
 def main(argv: Sequence[str]) -> int:

build/rbe/remotetool_test.py

@@ -224,37 +224,6 @@ def test_missing_section(self):
             remotetool.parse_show_action_output([])
 
 
-class ReadConfigFileLinesTests(unittest.TestCase):
-    def test_empty_file(self):
-        self.assertEqual(remotetool.read_config_file_lines([]), dict())
-
-    def test_ignore_blank_lines(self):
-        self.assertEqual(
-            remotetool.read_config_file_lines(["", "\t", "\n"]), dict()
-        )
-
-    def test_ignore_comments(self):
-        self.assertEqual(
-            remotetool.read_config_file_lines(["####", "# comment"]), dict()
-        )
-
-    def test_ignore_non_key_value_pairs(self):
-        self.assertEqual(
-            remotetool.read_config_file_lines(["value-only"]), dict()
-        )
-
-    def test_key_value(self):
-        self.assertEqual(
-            remotetool.read_config_file_lines(["key=value"]), {"key": "value"}
-        )
-
-    def test_last_wins(self):
-        self.assertEqual(
-            remotetool.read_config_file_lines(["key=value-1", "key=value-2"]),
-            {"key": "value-2"},
-        )
-
-
 _TEST_CFG = {
     "service": "other.remote.service.com:443",
     "instance": "projects/your-project/instance/default",