DNM: Yet another simulator pool implementation

apple/testing/default_runner/BUILD

@@ -1,5 +1,6 @@
 load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
 load("@rules_python//python:py_binary.bzl", "py_binary")
+load("@rules_python//python:py_library.bzl", "py_library")
 load(
     "//apple/testing/default_runner:ios_test_runner.bzl",
     "ios_test_runner",
@@ -90,6 +91,12 @@ exports_files([
     "xctrunner_entitlements.template.plist",
 ])
 
+py_library(
+    name = "simulator_utils",
+    srcs = ["simulator_utils.py"],
+    visibility = ["//apple:__subpackages__"],
+)
+
 py_binary(
     name = "simulator_creator",
     srcs = ["simulator_creator.py"],
@@ -99,6 +106,9 @@ py_binary(
     # should be considered an implementation detail of the rules and
     # not used by other things.
     visibility = ["//visibility:public"],
+    deps = [
+        ":simulator_utils",
+    ],
 )
 
 ios_test_runner(

apple/testing/default_runner/ios_xctestrun_runner.bzl

@@ -25,7 +25,9 @@ def _get_template_substitutions(
         xctrunner_entitlements_template,
         pre_action_binary,
         post_action_binary,
-        post_action_determines_exit_code):
+        post_action_determines_exit_code,
+        simulator_pool_server_port,
+        simulator_pool_client):
     substitutions = {
         "device_type": device_type,
         "os_version": os_version,
@@ -43,6 +45,8 @@ def _get_template_substitutions(
         "pre_action_binary": pre_action_binary,
         "post_action_binary": post_action_binary,
         "post_action_determines_exit_code": post_action_determines_exit_code,
+        "simulator_pool_server_port": simulator_pool_server_port,
+        "simulator_pool_client.py": simulator_pool_client,
     }
 
     return {"%({})s".format(key): value for key, value in substitutions.items()}
@@ -71,6 +75,8 @@ def _ios_xctestrun_runner_impl(ctx):
         ctx.file._xctrunner_entitlements_template,
     ]).merge(ctx.attr._simulator_creator[DefaultInfo].default_runfiles)
 
+    runfiles = runfiles.merge(ctx.attr._simulator_pool_client[DefaultInfo].default_runfiles)
+
     default_action_binary = "/usr/bin/true"
 
     pre_action_binary = default_action_binary
@@ -105,6 +111,8 @@ def _ios_xctestrun_runner_impl(ctx):
             pre_action_binary = pre_action_binary,
             post_action_binary = post_action_binary,
             post_action_determines_exit_code = "true" if post_action_determines_exit_code else "false",
+            simulator_pool_server_port = "" if ctx.attr.simulator_pool_server_port else str(ctx.attr.simulator_pool_server_port),
+            simulator_pool_client = ctx.executable._simulator_pool_client.short_path,
         ),
     )
 
@@ -204,6 +212,14 @@ A binary to run following test execution. Runs after testing but before test res
 When true, the exit code of the test run will be set to the exit code of the post action. This is useful for tests that need to fail the test run based on their own criteria.
 """,
         ),
+        "simulator_pool_server_port": attr.int(
+            doc = "The port of a running simulator pool server. If set, the test runner will connect to the simulator pool server and use the simulators from the pool instead of creating new ones.",
+        ),
+        "_simulator_pool_client": attr.label(
+            default = "//apple/testing/simulator_pool:simulator_pool_client",
+            executable = True,
+            cfg = "exec",
+        ),
         "_simulator_creator": attr.label(
             default = Label(
                 "//apple/testing/default_runner:simulator_creator",

apple/testing/default_runner/ios_xctestrun_runner.template.sh

@@ -62,10 +62,30 @@ basename_without_extension() {
   echo "${filename%.*}"
 }
 
+simulator_pool_client_path="%(simulator_pool_client.py)s"
+simulator_pool_server_port="%(simulator_pool_server_port)s"
+simulator_pool_enabled=false
+if [[ -n "$simulator_pool_server_port" ]] && [[ -n "$device_id" ]] && [[ -n "$simulator_pool_client_path" ]]; then
+  simulator_pool_enabled=true
+fi
+simulator_id=""
+
+return_simulator_to_pool() {
+  if [[ "$simulator_pool_enabled" == true ]]; then
+    "$simulator_pool_client_path" return --udid="$simulator_id" --port "$simulator_pool_server_port"
+  fi
+}
+
+teardown() {
+    return_simulator_to_pool
+    rm -rf "${test_tmp_dir}"
+}
+
 test_tmp_dir="$(mktemp -d "${TEST_TMPDIR:-${TMPDIR:-/tmp}}/test_tmp_dir.XXXXXX")"
 if [[ -z "${NO_CLEAN:-}" ]]; then
-  trap 'rm -rf "${test_tmp_dir}"' EXIT
+  trap 'teardown' EXIT
 else
+  return_simulator_to_pool
   test_tmp_dir="${TMPDIR:-/tmp}/test_tmp_dir"
   rm -rf "$test_tmp_dir"
   mkdir -p "$test_tmp_dir"
@@ -404,8 +424,18 @@ else
   simulator_creator_args+=(--no-reuse-simulator)
 fi
 
-simulator_id="unused"
-if [[ "$build_for_device" == false ]]; then
+if [[ "$simulator_pool_enabled" == true ]]; then
+  request_simulator_args=(
+      --port "$simulator_pool_server_port" \
+      --test-target "$test_bundle_name"
+      --device-type "%(device_type)s"
+      --os-version "%(os_version)s"
+  )
+  if [[ -n "${test_host_path:-}" ]]; then
+      request_simulator_args+=(--test-host "$(basename_without_extension "$test_host_path")")
+  fi
+  simulator_id=$("$simulator_pool_client_path" request "${request_simulator_args[@]}")
+else
   simulator_id="$("./%(simulator_creator.py)s" \
     "${simulator_creator_args[@]}"
   )"

apple/testing/default_runner/simulator_creator.py

@@ -20,67 +20,9 @@
 import subprocess
 import sys
 import time
+import apple.testing.default_runner.simulator_utils as simulator_utils
 from typing import List, Optional
 
-
-def _simctl(extra_args: List[str]) -> str:
-    return subprocess.check_output(["xcrun", "simctl"] + extra_args).decode()
-
-
-def _boot_simulator(simulator_id: str) -> None:
-    # This private command boots the simulator if it isn't already, and waits
-    # for the appropriate amount of time until we can actually run tests
-    try:
-        output = _simctl(["bootstatus", simulator_id, "-b"])
-        print(output, file=sys.stderr)
-    except subprocess.CalledProcessError as e:
-        exit_code = e.returncode
-
-        # When reusing simulators we may encounter the error:
-        # 'Unable to boot device in current state: Booted'.
-        #
-        # This is because the simulator is already booted, and we can ignore it
-        # if we check and the simulator is in fact booted.
-        if exit_code == 149:
-            devices = json.loads(
-                _simctl(["list", "devices", "-j", simulator_id]),
-            )["devices"]
-            device = next(
-                (
-                    blob
-                    for devices_for_os in devices.values()
-                    for blob in devices_for_os
-                    if blob["udid"] == simulator_id
-                ),
-                None
-            )
-            if device and device["state"].lower() == "booted":
-                print(
-                    f"Simulator '{device['name']}' ({simulator_id}) is already booted",
-                    file=sys.stderr,
-                )
-                exit_code = 0
-
-        # Both of these errors translate to strange simulator states that may
-        # end up causing issues, but attempting to actually use the simulator
-        # instead of failing at this point might still succeed
-        #
-        # 164: EBADDEVICE
-        # 165: EBADDEVICESTATE
-        if exit_code in (164, 165):
-            print(
-                f"Ignoring 'simctl bootstatus' exit code {exit_code}",
-                file=sys.stderr,
-            )
-        elif exit_code != 0:
-            print(f"'simctl bootstatus' exit code {exit_code}", file=sys.stderr)
-            raise
-
-    # Add more arbitrary delay before tests run. Even bootstatus doesn't wait
-    # long enough and tests can still fail because the simulator isn't ready
-    time.sleep(3)
-
-
 def _device_name(device_type: str, os_version: str) -> str:
     return f"BAZEL_TEST_{device_type}_{os_version}"
 
@@ -109,7 +51,7 @@ def _build_parser() -> argparse.ArgumentParser:
 
 
 def _main(os_version: str, device_type: str, name: Optional[str], reuse_simulator: bool) -> None:
-    devices = json.loads(_simctl(["list", "devices", "-j"]))["devices"]
+    devices = json.loads(simulator_utils.simctl(["list", "devices", "-j"]))["devices"]
     device_name = name or _device_name(device_type, os_version)
     runtime_identifier = "com.apple.CoreSimulator.SimRuntime.iOS-{}".format(
         os_version.replace(".", "-")
@@ -132,19 +74,19 @@ def _main(os_version: str, device_type: str, name: Optional[str], reuse_simulato
         state = existing_device["state"].lower()
         print(f"Existing simulator '{name}' ({simulator_id}) state is: {state}", file=sys.stderr)
         if state != "booted":
-            _boot_simulator(simulator_id)
+            simulator_utils.boot_simulator(simulator_id)
     else:
         if not reuse_simulator:
             # Simulator reuse is based on device name, therefore we must generate a unique name to
             # prevent unintended reuse.
             device_name_suffix = ''.join(random.choices(string.ascii_letters + string.digits, k=8))
             device_name += f"_{device_name_suffix}"
 
-        simulator_id = _simctl(
+        simulator_id = simulator_utils.simctl(
             ["create", device_name, device_type, runtime_identifier]
         ).strip()
         print(f"Created new simulator '{device_name}' ({simulator_id})", file=sys.stderr)
-        _boot_simulator(simulator_id)
+        simulator_utils.boot_simulator(simulator_id)
 
     print(simulator_id.strip())
 

apple/testing/default_runner/simulator_utils.py

@@ -0,0 +1,88 @@
+#!/usr/bin/python3
+# Copyright 2022 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import sys
+import time
+import subprocess
+from typing import List
+
+
+def simctl(extra_args: List[str]) -> str:
+    """Execute simctl command with the given arguments.
+
+    Args:
+        extra_args: List of additional arguments to pass to simctl
+
+    Returns:
+        The decoded output from the simctl command
+
+    Raises:
+        subprocess.CalledProcessError: If the simctl command fails
+    """
+    return subprocess.check_output(["xcrun", "simctl"] + extra_args).decode()
+
+def boot_simulator(simulator_id: str) -> None:
+    # This private command boots the simulator if it isn't already, and waits
+    # for the appropriate amount of time until we can actually run tests
+    try:
+        output = simctl(["bootstatus", simulator_id, "-b"])
+        print(output, file=sys.stderr)
+    except subprocess.CalledProcessError as e:
+        exit_code = e.returncode
+
+        # When reusing simulators we may encounter the error:
+        # 'Unable to boot device in current state: Booted'.
+        #
+        # This is because the simulator is already booted, and we can ignore it
+        # if we check and the simulator is in fact booted.
+        if exit_code == 149:
+            devices = json.loads(
+                simctl(["list", "devices", "-j", simulator_id]),
+            )["devices"]
+            device = next(
+                (
+                    blob
+                    for devices_for_os in devices.values()
+                    for blob in devices_for_os
+                    if blob["udid"] == simulator_id
+                ),
+                None
+            )
+            if device and device["state"].lower() == "booted":
+                print(
+                    f"Simulator '{device['name']}' ({simulator_id}) is already booted",
+                    file=sys.stderr,
+                )
+                exit_code = 0
+
+        # Both of these errors translate to strange simulator states that may
+        # end up causing issues, but attempting to actually use the simulator
+        # instead of failing at this point might still succeed
+        #
+        # 164: EBADDEVICE
+        # 165: EBADDEVICESTATE
+        if exit_code in (164, 165):
+            print(
+                f"Ignoring 'simctl bootstatus' exit code {exit_code}",
+                file=sys.stderr,
+            )
+        elif exit_code != 0:
+            print(f"'simctl bootstatus' exit code {exit_code}", file=sys.stderr)
+            raise
+
+    # Add more arbitrary delay before tests run. Even bootstatus doesn't wait
+    # long enough and tests can still fail because the simulator isn't ready
+    time.sleep(3)

apple/testing/simulator_pool/BUILD

@@ -0,0 +1,46 @@
+load("@rules_python//python:py_binary.bzl", "py_binary")
+load("//apple/testing/simulator_pool:create_simulator_pool.bzl", "create_simulator_pool")
+
+exports_files(["create_simulator_pool.template.sh"])
+
+create_simulator_pool(
+    name = "create_simulator_pool",
+    device_type = "iPhone 15 Pro",
+    os_version = "18.3",
+    pool_size = 3,
+    server_port = 50051,
+)
+
+py_binary(
+    name = "create_simulator_pool_tool",
+    srcs = [
+        "create_simulator_pool_tool.py",
+    ],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        "//apple/testing/default_runner:simulator_utils",
+    ],
+)
+
+py_binary(
+    name = "simulator_pool_server",
+    srcs = [
+        "simulator_pool_server.py",
+    ],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        "//apple/testing/default_runner:simulator_utils",
+    ],
+)
+
+py_binary(
+    name = "simulator_pool_client",
+    srcs = [
+        "simulator_pool_client.py",
+    ],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+)