async def run_instruments(
self,
stop: asyncio.Event,
template: str,
app_bundle_id: str,
trace_path: str,
post_process_arguments: Optional[List[str]] = None,
env: Optional[Dict[str, str]] = None,
app_args: Optional[List[str]] = None,
started: Optional[asyncio.Event] = None,
timings: Optional[InstrumentsTimings] = None,
) -> str:
trace_path = os.path.realpath(trace_path)
self.logger.info(f"Starting instruments connection, writing to {trace_path}")
async with self.get_stub() as stub, stub.instruments_run.open() as stream:
self.logger.info("Sending instruments request")
await stream.send_message(
InstrumentsRunRequest(
start=InstrumentsRunRequest.Start(
file_path=None,
template_name=template,
app_bundle_id=app_bundle_id,
environment=env,
arguments=app_args,
timings=translate_instruments_timings(timings),
)
)
)
self.logger.info("Starting instruments")
await drain_until_running(stream=stream, logger=self.logger)
if started:
started.set()
self.logger.info("Instruments has started, waiting for stop")
async for response in stop_wrapper(stream=stream, stop=stop):
output = response.log_output
if len(output):
self.logger.info(output.decode())
self.logger.info("Stopping instruments")
await stream.send_message(
InstrumentsRunRequest(
stop=InstrumentsRunRequest.Stop(
post_process_arguments=post_process_arguments
)
)
)
await stream.end()
self.logger.info(f"Writing instruments from tar to {trace_path}")
await drain_untar(
instruments_generate_bytes(stream=stream, logger=self.logger),
output_path=trace_path,
)
self.logger.info(f"Instruments trace written to {trace_path}")
return trace_path
def translate_instruments_timings(
timings: Optional[InstrumentsTimings],
) -> Optional[InstrumentsRunRequest.InstrumentsTimings]:
return (InstrumentsRunRequest.InstrumentsTimings(
terminate_timeout=timings.terminate_timeout,
launch_retry_timeout=timings.launch_retry_timeout,
launch_error_timeout=timings.launch_error_timeout,
operation_duration=timings.operation_duration,
) if timings else None)
async def run_instruments(
self,
stop: asyncio.Event,
trace_basename: str,
template_name: str,
app_bundle_id: str,
app_environment: Optional[Dict[str, str]] = None,
app_arguments: Optional[List[str]] = None,
tool_arguments: Optional[List[str]] = None,
started: Optional[asyncio.Event] = None,
timings: Optional[InstrumentsTimings] = None,
post_process_arguments: Optional[List[str]] = None,
) -> List[str]:
self.logger.info(f"Starting instruments connection")
async with self.get_stub() as stub, stub.instruments_run.open() as stream:
self.logger.info("Sending instruments request")
await stream.send_message(
InstrumentsRunRequest(
start=InstrumentsRunRequest.Start(
template_name=template_name,
app_bundle_id=app_bundle_id,
environment=app_environment,
arguments=app_arguments,
tool_arguments=tool_arguments,
timings=translate_instruments_timings(timings),
)
)
)
self.logger.info("Starting instruments")
await instruments_drain_until_running(stream=stream, logger=self.logger)
if started:
started.set()
self.logger.info("Instruments has started, waiting for stop")
async for response in stop_wrapper(stream=stream, stop=stop):
output = response.log_output
if len(output):
self.logger.info(output.decode())
self.logger.info("Stopping instruments")
await stream.send_message(
InstrumentsRunRequest(
stop=InstrumentsRunRequest.Stop(
post_process_arguments=post_process_arguments
)
)
)
await stream.end()
result = []
with tempfile.TemporaryDirectory() as tmp_trace_dir:
self.logger.info(
f"Writing instruments data from tar to {tmp_trace_dir}"
)
await drain_untar(
instruments_generate_bytes(stream=stream, logger=self.logger),
output_path=tmp_trace_dir,
)
if os.path.exists(
os.path.join(os.path.abspath(tmp_trace_dir), "instrument_data")
):
# tar is an instruments trace (old behavior)
trace_file = f"{trace_basename}.trace"
shutil.copytree(tmp_trace_dir, trace_file)
result.append(trace_file)
self.logger.info(f"Trace written to {trace_file}")
else:
# tar is a folder containing one or more trace files
for file in os.listdir(tmp_trace_dir):
_, file_extension = os.path.splitext(file)
tmp_trace_file = os.path.join(
os.path.abspath(tmp_trace_dir), file
)
trace_file = f"{trace_basename}{file_extension}"
shutil.move(tmp_trace_file, trace_file)
result.append(trace_file)
self.logger.info(f"Trace written to {trace_file}")
return result