Add ability to collect launch information from bazel run (#2774)

By specifying a path with `BAZEL_APPLE_LAUNCH_INFO_PATH`, `bazel run` of
an application will write the following JSON to that file:
* `platform`: The string `device` if running on a device, otherwise the
lldb platform (e.g. `ios-simulator`).
* `udid`: The UDID of the device.
* `pid`: The PID of the launched application.

This information can be used to attach `lldb` to the running
application.

Since `devicectl` waits until the process exists before writing to
`--json_output`, you’ll need to use `BAZEL_DEVICECTL_LAUNCH_FLAGS` to
not use `--console`.

---------

Signed-off-by: Brentley Jones <github@brentleyjones.com>

Author: brentleyjones

apple/internal/templates/apple_device.template.py

@@ -45,6 +45,7 @@
 import sys
 import tempfile
 from typing import Any, Dict, Optional
+from uuid import uuid4
 import zipfile
 
 
@@ -418,19 +419,99 @@ def run_app(
     logger.info(
         "Launching app %s on %s", app_bundle_id, device_identifier
     )
-    args = [
+    launch_args = [
         devicectl_path,
         "device",
         "process",
         "launch",
         *launch_args,
         "--device",
         device_identifier,
-        app_bundle_id,
     ]
     # Append optional launch arguments.
-    args.extend(sys.argv[1:])
-    subprocess.run(args, env=devicectl_launch_environ(), check=False)
+    app_args = [app_bundle_id] + sys.argv[1:]
+    launch_app(
+        launch_args=launch_args,
+        app_args=app_args,
+        env=devicectl_launch_environ(),
+        device_identifier=device_identifier,
+    )
+
+
+def launch_app(
+    *,
+    launch_args: list[str],
+    app_args: list[str],
+    env: Dict[str, str],
+    device_identifier: str,
+) -> None:
+  """Launches an app in a simulator.
+
+  Args:
+    launch_args: The arguments to pass to simctl to launch the app, excluding
+      the bundle id and app arguments.
+    app_args: The bundle id and app arguments to pass to simctl to launch the
+      app.
+    env: The environment variables to pass to simctl.
+    device_identifier: The identifier of the device.
+  """
+  launch_info_path = os.environ.get("BAZEL_APPLE_LAUNCH_INFO_PATH")
+  if not launch_info_path:
+    subprocess.run(launch_args + app_args, env=env, check=True)
+    return
+
+  if "--json-output" in launch_args:
+    idx = launch_args.index("--json-output")
+    json_output_path = launch_args[idx + 1]
+    delete_json_output = False
+  else:
+    json_output_path = os.path.join(
+        tempfile.gettempdir(), f"devicectl_launch_{uuid4().hex}.json",
+    )
+    launch_args = launch_args + ["--json-output", json_output_path]
+    delete_json_output = True
+  args = launch_args + app_args
+
+  proc = subprocess.Popen(
+      args,
+      env=env,
+  )
+
+  exit_code = proc.wait()
+
+  # `devicectl` only writes to `--json-output` after the process has exited. We
+  # use `subprocess.Popen()` instead of `subprocess.run()` to allow us to write
+  # the launch info before reporting process exit.
+  try:
+    with open(json_output_path, "r", encoding="utf-8") as f:
+      data = json.load(f)
+      pid = (
+          data.get("result", {}).get("process", {}).get("processIdentifier")
+      )
+      if pid:
+        os.makedirs(os.path.dirname(launch_info_path), exist_ok=True)
+        with open(launch_info_path, "w", encoding="utf-8") as f:
+          f.write(json.dumps(
+              {
+                  "platform": "device",
+                  "udid": device_identifier,
+                  "pid": pid,
+              },
+              indent=2,
+          ))
+      else:
+          logger.error("Failed to find PID in JSON output")
+  except Exception as e:
+    logger.error("Failed to write launch info to file: %s", e)
+
+  if delete_json_output:
+    try:
+      os.remove(json_output_path)
+    except Exception:
+      pass
+
+  if exit_code != 0:
+    raise subprocess.CalledProcessError(exit_code, args)
 
 
 def main(

apple/internal/templates/apple_simulator.template.py

@@ -52,12 +52,14 @@
 import pathlib
 import platform
 import plistlib
+import pty
+import re
 import shlex
 import shutil
 import subprocess
 import sys
 import tempfile
-from typing import Dict, Optional
+from typing import IO, Dict, Optional, Sequence
 import zipfile
 
 
@@ -78,6 +80,22 @@
   )
 
 
+class BufferFlusher:
+  """Flushes a buffer to a file descriptor.
+
+  This is used to ensure that the buffer is flushed to the file descriptor
+  as soon as possible.
+  """
+
+  def __init__(self, raw: IO[bytes]):
+    self.raw = raw
+
+  def write(self, b: bytes) -> int:
+    n = self.raw.write(b)
+    self.raw.flush()
+    return n
+
+
 class DeviceType(collections.abc.Mapping):
   """Wraps the `devicetype` dictionary from `simctl list -j`.
 
@@ -662,9 +680,10 @@ def run_app_in_simulator(
   root_dir = os.path.dirname(application_output_path)
   register_dsyms(root_dir)
   with extracted_app(application_output_path, app_name) as app_path:
-    logger.debug("Installing app %s to simulator %s", app_path, simulator_udid)
+    logger.info("Installing app %s to simulator %s", app_path, simulator_udid)
     subprocess.run(
-        [simctl_path, "install", simulator_udid, app_path], check=True
+        [simctl_path, "install", simulator_udid, app_path],
+        check=True,
     )
     app_bundle_id = bundle_id(app_path)
     launch_args = shlex.split(
@@ -686,7 +705,71 @@ def run_app_in_simulator(
     ]
     # Append optional launch arguments.
     args.extend(sys.argv[1:])
-    subprocess.run(args, env=simctl_launch_environ(), check=False)
+    launch_app(args, env=simctl_launch_environ(), simulator_udid=simulator_udid)
+
+
+def launch_app(
+    args: Sequence[str],
+    *,
+    env: Dict[str, str],
+    simulator_udid: str,
+) -> None:
+  """Launches an app in a simulator.
+
+  Args:
+    args: The arguments to pass to simctl.
+    env: The environment variables to pass to simctl.
+    simulator_udid: The UDID of the simulator in which to run the app.
+  """
+  launch_info_path = os.environ.get("BAZEL_APPLE_LAUNCH_INFO_PATH")
+  if not launch_info_path:
+    subprocess.run(args, env=env, check=True)
+    return
+
+  # Open a PTY to capture the output of simctl. We need a PTY to ensure that
+  # the PID is written to stdout before the rest of the app output.
+  primary_fd, secondary_fd = pty.openpty()
+
+  proc = subprocess.Popen(
+      args,
+      env=env,
+      stdout=secondary_fd,
+      close_fds=True,
+  )
+
+  # simctl has the fd dup; close ours.
+  os.close(secondary_fd)
+
+  with os.fdopen(primary_fd, "rb", buffering=0) as r:
+    # Grab PID from the first line of output.
+    first_line = r.readline()
+    pid_match = re.search(rb":\s*(\d+)\s*$", first_line)
+    if pid_match:
+      pid = int(pid_match.group(1))
+      try:
+        os.makedirs(os.path.dirname(launch_info_path), exist_ok=True)
+        with open(launch_info_path, "w", encoding="utf-8") as f:
+          f.write(json.dumps(
+              {
+                  "platform": "ios-simulator",
+                  "udid": simulator_udid,
+                  "pid": pid,
+              },
+              indent=2,
+          ))
+      except Exception as e:
+        logger.error("Failed to write launch info to file: %s", e)
+    else:
+      logger.error("Failed to parse PID from output")
+
+    # Stream the rest until simctl exits.
+    sys.stdout.buffer.write(first_line)
+    sys.stdout.flush()
+    shutil.copyfileobj(r, BufferFlusher(sys.stdout.buffer))
+
+  exit_code = proc.wait()
+  if exit_code != 0:
+    raise subprocess.CalledProcessError(exit_code, args)
 
 
 def main(