| #!/usr/bin/env python3 | |
| # Copyright (C) 2023 The Android Open Source Project | |
| # | |
| # 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 os | |
| import struct | |
| import subprocess | |
| import sys | |
| import tempfile | |
| from typing import Any, IO, List, Optional | |
| from google.protobuf import text_format | |
| from python.generators.diff_tests.testing import (DataPath, Path, | |
| SimpleperfProto, TextProto, | |
| TraceInjector) | |
| from python.generators.diff_tests.utils import ProtoManager | |
| ROOT_DIR = os.path.dirname( | |
| os.path.dirname( | |
| os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) | |
| class TraceGenerator: | |
| """A helper class for generating trace files.""" | |
| def __init__(self, trace_descriptor_path: str, | |
| extension_descriptor_paths: List[str]): | |
| self.trace_descriptor_path = trace_descriptor_path | |
| self.extension_descriptor_paths = extension_descriptor_paths | |
| def serialize_textproto_trace(self, text_proto_path: str, | |
| out_stream: IO[bytes]): | |
| proto = ProtoManager([self.trace_descriptor_path] + | |
| self.extension_descriptor_paths).create_message( | |
| 'perfetto.protos.Trace')() | |
| with open(text_proto_path, 'r') as text_proto_file: | |
| text_format.Merge(text_proto_file.read(), proto) | |
| out_stream.write(proto.SerializeToString()) | |
| out_stream.flush() | |
| def serialize_python_trace(self, root_dir: str, python_trace_path: str, | |
| out_stream: IO[bytes]): | |
| python_cmd = [ | |
| sys.executable, | |
| python_trace_path, | |
| self.trace_descriptor_path, | |
| ] | |
| python_cmd.extend(self.extension_descriptor_paths) | |
| # Add the test dir to the PYTHONPATH to allow synth_common to be found. | |
| env = os.environ.copy() | |
| if 'PYTHONPATH' in env: | |
| env['PYTHONPATH'] = "{}:{}".format( | |
| os.path.join(root_dir, 'test'), env['PYTHONPATH']) | |
| else: | |
| env['PYTHONPATH'] = os.path.join(root_dir, 'test') | |
| # Use start_new_session to put trace generation in a separate process group. | |
| # This prevents SIGINT from Ctrl+C reaching these processes and causing | |
| # protobuf to crash mid-operation. | |
| subprocess.check_call( | |
| python_cmd, env=env, stdout=out_stream, start_new_session=True) | |
| def serialize_simpleperf_proto_trace(self, simpleperf_trace: SimpleperfProto, | |
| out_stream: IO[bytes]): | |
| # Write simpleperf_proto header | |
| # Magic: "SIMPLEPERF" (10 bytes, null-padded) | |
| magic = b"SIMPLEPERF" | |
| out_stream.write(magic) | |
| # Version: LittleEndian16 = 1 | |
| out_stream.write(struct.pack('<H', 1)) | |
| # Get the simpleperf Record proto message type | |
| proto_manager = ProtoManager([self.trace_descriptor_path] + | |
| self.extension_descriptor_paths) | |
| record_proto_class = proto_manager.create_message( | |
| 'perfetto.third_party.simpleperf.proto.Record') | |
| # Write each record | |
| for record_textproto in simpleperf_trace.records: | |
| record = record_proto_class() | |
| text_format.Merge(record_textproto, record) | |
| record_bytes = record.SerializeToString() | |
| # Write record size (LittleEndian32) | |
| out_stream.write(struct.pack('<I', len(record_bytes))) | |
| # Write record data | |
| out_stream.write(record_bytes) | |
| # End marker: LittleEndian32(0) | |
| out_stream.write(struct.pack('<I', 0)) | |
| out_stream.flush() | |
| def generate_trace_file(test_case: Any, trace_descriptor_path: str, | |
| extension_descriptor_paths: List[str], | |
| simpleperf_descriptor: str) -> Optional[Any]: | |
| # We can't use delete=True here. When using that on Windows, the | |
| # resulting file is opened in exclusive mode (in turn that's a subtle | |
| # side-effect of the underlying CreateFile(FILE_ATTRIBUTE_TEMPORARY)) | |
| # and TP fails to open the passed path. | |
| gen_trace_file = None | |
| trace_generator = TraceGenerator(trace_descriptor_path, | |
| extension_descriptor_paths) | |
| if test_case.blueprint.is_trace_file(): | |
| assert test_case.trace_path | |
| if test_case.trace_path.endswith('.py'): | |
| gen_trace_file = tempfile.NamedTemporaryFile(delete=False) | |
| trace_generator.serialize_python_trace(ROOT_DIR, test_case.trace_path, | |
| gen_trace_file) | |
| elif test_case.trace_path.endswith('.textproto'): | |
| gen_trace_file = tempfile.NamedTemporaryFile(delete=False) | |
| trace_generator.serialize_textproto_trace(test_case.trace_path, | |
| gen_trace_file) | |
| elif test_case.blueprint.is_trace_textproto(): | |
| gen_trace_file = tempfile.NamedTemporaryFile(delete=False) | |
| proto = ProtoManager([trace_descriptor_path] + | |
| extension_descriptor_paths).create_message( | |
| 'perfetto.protos.Trace')() | |
| assert isinstance(test_case.blueprint.trace, TextProto) | |
| text_format.Merge(test_case.blueprint.trace.contents, proto) | |
| gen_trace_file.write(proto.SerializeToString()) | |
| gen_trace_file.flush() | |
| elif test_case.blueprint.is_trace_simpleperf_proto(): | |
| gen_trace_file = tempfile.NamedTemporaryFile(delete=False) | |
| # Simpleperf is a separate format, so use a dedicated generator with | |
| # the simpleperf descriptor instead of the general Perfetto trace extensions. | |
| simpleperf_generator = TraceGenerator(trace_descriptor_path, | |
| [simpleperf_descriptor]) | |
| simpleperf_generator.serialize_simpleperf_proto_trace( | |
| test_case.blueprint.trace, gen_trace_file) | |
| else: | |
| gen_trace_file = tempfile.NamedTemporaryFile(delete=False) | |
| with open(gen_trace_file.name, 'w') as trace_file: | |
| if not isinstance(test_case.blueprint.trace, (Path, DataPath)): | |
| trace_file.write(test_case.blueprint.trace.contents) | |
| if test_case.blueprint.trace_modifier is not None: | |
| if gen_trace_file: | |
| # Overwrite |gen_trace_file|. | |
| modify_trace(trace_descriptor_path, extension_descriptor_paths, | |
| gen_trace_file.name, gen_trace_file.name, | |
| test_case.blueprint.trace_modifier) | |
| else: | |
| # Create |gen_trace_file| to save the modified trace. | |
| gen_trace_file = tempfile.NamedTemporaryFile(delete=False) | |
| assert test_case.trace_path | |
| modify_trace(trace_descriptor_path, extension_descriptor_paths, | |
| test_case.trace_path, gen_trace_file.name, | |
| test_case.blueprint.trace_modifier) | |
| return gen_trace_file | |
| def modify_trace(trace_descriptor_path: str, | |
| extension_descriptor_paths: List[str], in_trace_path: str, | |
| out_trace_path: str, modifier: TraceInjector): | |
| """Modifies a trace file with the given modifier.""" | |
| trace_proto = ProtoManager([trace_descriptor_path] + | |
| extension_descriptor_paths).create_message( | |
| 'perfetto.protos.Trace')() | |
| with open(in_trace_path, "rb") as f: | |
| # This may raise DecodeError when |in_trace_path| isn't protobuf. | |
| trace_proto.ParseFromString(f.read()) | |
| # Modify the trace proto object with the provided modifier function. | |
| modifier.inject(trace_proto) | |
| with open(out_trace_path, "wb") as f: | |
| f.write(trace_proto.SerializeToString()) | |
| f.flush() |