| #!/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 io | |
| import os | |
| import struct | |
| import subprocess | |
| import sys | |
| import tarfile | |
| import tempfile | |
| import zipfile | |
| from typing import Any, IO, List, Optional, Union | |
| from google.protobuf import text_format | |
| from python.generators.diff_tests.testing import (DataPath, Path, | |
| SimpleperfProto, Tar, | |
| TextProto, TraceInjector, Zip) | |
| 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_member( | |
| self, blueprint: Any, member: Union[str, bytes, TextProto, Path, | |
| DataPath]) -> bytes: | |
| """Serializes a single archive member to bytes (see Zip/Tar docs).""" | |
| if isinstance(member, bytes): | |
| return member | |
| if isinstance(member, TextProto): | |
| proto = ProtoManager([self.trace_descriptor_path] + | |
| self.extension_descriptor_paths).create_message( | |
| 'perfetto.protos.Trace')() | |
| text_format.Merge(member.contents, proto) | |
| return proto.SerializeToString() | |
| if isinstance(member, DataPath): | |
| path = os.path.join(blueprint.test_data_dir, member.filename) | |
| with open(path, 'rb') as f: | |
| return f.read() | |
| if isinstance(member, Path): | |
| path = os.path.abspath(os.path.join(blueprint.index_dir, member.filename)) | |
| with open(path, 'rb') as f: | |
| return f.read() | |
| assert isinstance(member, str) | |
| return member.encode('utf-8') | |
| def serialize_zip_trace(self, blueprint: Any, archive: Zip, | |
| out_stream: IO[bytes]): | |
| buf = io.BytesIO() | |
| with zipfile.ZipFile(buf, 'w') as z: | |
| for name, member in archive.members.items(): | |
| z.writestr(name, self.serialize_member(blueprint, member)) | |
| out_stream.write(buf.getvalue()) | |
| out_stream.flush() | |
| def serialize_tar_trace(self, blueprint: Any, archive: Tar, | |
| out_stream: IO[bytes]): | |
| buf = io.BytesIO() | |
| tar_format = tarfile.PAX_FORMAT if archive.macos_style else ( | |
| tarfile.DEFAULT_FORMAT) | |
| with tarfile.open(fileobj=buf, mode='w', format=tar_format) as t: | |
| for name, member in archive.members.items(): | |
| data = self.serialize_member(blueprint, member) | |
| info = tarfile.TarInfo(name) | |
| info.size = len(data) | |
| if archive.macos_style: | |
| # Force emission of a PAX extended header ('x' typeflag) block, as | |
| # macOS/BSD tar does for every member. | |
| info.pax_headers = {'comment': 'perfetto-test'} | |
| t.addfile(info, io.BytesIO(data)) | |
| raw = buf.getvalue() | |
| if archive.macos_style: | |
| raw = self._space_terminate_tar_numeric_fields(raw) | |
| out_stream.write(raw) | |
| out_stream.flush() | |
| @staticmethod | |
| def _space_terminate_tar_numeric_fields(raw: bytes) -> bytes: | |
| """Rewrites the size field of every tar header from NUL- to space- | |
| termination, matching what macOS/BSD tar emits (11 octal digits followed by | |
| a space). Only 512-byte blocks that are ustar headers are touched; the | |
| parser ignores the header checksum so it is left unchanged.""" | |
| ba = bytearray(raw) | |
| for off in range(0, len(ba) - 512 + 1, 512): | |
| if bytes(ba[off + 257:off + 262]) != b'ustar': | |
| continue | |
| field = bytes(ba[off + 124:off + 136]) | |
| digits = field.split(b'\x00')[0].split(b' ')[0] | |
| if not digits: | |
| continue | |
| try: | |
| value = int(digits, 8) | |
| except ValueError: | |
| continue | |
| ba[off + 124:off + 136] = b'%011o ' % value | |
| return bytes(ba) | |
| 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_zip(): | |
| gen_trace_file = tempfile.NamedTemporaryFile(delete=False) | |
| trace_generator.serialize_zip_trace(test_case.blueprint, | |
| test_case.blueprint.trace, | |
| gen_trace_file) | |
| elif test_case.blueprint.is_trace_tar(): | |
| gen_trace_file = tempfile.NamedTemporaryFile(delete=False) | |
| trace_generator.serialize_tar_trace(test_case.blueprint, | |
| test_case.blueprint.trace, | |
| gen_trace_file) | |
| 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() |