xref: /aosp_15_r20/external/pigweed/pw_trace_tokenized/py/pw_trace_tokenized/get_trace.py (revision 61c4878ac05f98d0ceed94b57d316916de578985) 
1#!/usr/bin/env python3
2# Copyright 2021 The Pigweed Authors
3#
4# Licensed under the Apache License, Version 2.0 (the "License"); you may not
5# use this file except in compliance with the License. You may obtain a copy of
6# the License at
7#
8#     https://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13# License for the specific language governing permissions and limitations under
14# the License.
15"""
16Generates json trace files viewable using chrome://tracing using RPCs from a
17connected HdlcRpcClient.
18
19Example usage:
20python pw_trace_tokenized/py/pw_trace_tokenized/get_trace.py -s localhost:33000
21  -o trace.json
22  -t
23  out/pw_strict_host_clang_debug/obj/pw_trace_tokenized/bin/trace_tokenized_example_rpc
24  pw_trace_tokenized/pw_trace_protos/trace_rpc.proto
25"""  # pylint: disable=line-too-long
26# pylint: enable=line-too-long
27
28import argparse
29import glob
30import logging
31from pathlib import Path
32import socket
33import sys
34from typing import Collection, Iterable, Iterator
35
36import serial
37from pw_tokenizer import database
38from pw_trace import trace
39from pw_hdlc.rpc import HdlcRpcClient, default_channels
40from pw_stream import stream_readers
41from pw_trace_tokenized import trace_tokenized
42
43_LOG = logging.getLogger('pw_trace_tokenizer')
44
45PW_RPC_MAX_PACKET_SIZE = 256
46SOCKET_SERVER = 'localhost'
47SOCKET_PORT = 33000
48MKFIFO_MODE = 0o666
49
50
51class SocketClientImpl:
52    def __init__(self, config: str):
53        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
54        socket_server = ''
55        socket_port = 0
56
57        if config == 'default':
58            socket_server = SOCKET_SERVER
59            socket_port = SOCKET_PORT
60        else:
61            socket_server, socket_port_str = config.split(':')
62            socket_port = int(socket_port_str)
63        self.socket.connect((socket_server, socket_port))
64
65    def write(self, data: bytes):
66        self.socket.sendall(data)
67
68    def read(self, num_bytes: int = PW_RPC_MAX_PACKET_SIZE):
69        return self.socket.recv(num_bytes)
70
71
72def _expand_globs(globs: Iterable[str]) -> Iterator[Path]:
73    for pattern in globs:
74        for file in glob.glob(pattern, recursive=True):
75            yield Path(file)
76
77
78def get_hdlc_rpc_client(
79    device: str,
80    baudrate: int,
81    proto_globs: Collection[str],
82    socket_addr: str,
83    **kwargs,
84):
85    """Get the HdlcRpcClient based on arguments."""
86    del kwargs  # ignore
87    if not proto_globs:
88        proto_globs = ['**/*.proto']
89
90    protos = list(_expand_globs(proto_globs))
91
92    if not protos:
93        _LOG.critical(
94            'No .proto files were found with %s', ', '.join(proto_globs)
95        )
96        _LOG.critical('At least one .proto file is required')
97        return 1
98
99    _LOG.debug(
100        'Found %d .proto files found with %s',
101        len(protos),
102        ', '.join(proto_globs),
103    )
104
105    # TODO(rgoliver): When pw has a generalized transport for RPC this should
106    # use it so it isn't specific to HDLC
107    if socket_addr is None:
108        serial_device = serial.Serial(device, baudrate, timeout=1)
109        reader = stream_readers.SerialReader(serial_device)
110        write_function = serial_device.write
111    else:
112        try:
113            socket_device = SocketClientImpl(socket_addr)
114            reader = stream_readers.SocketReader(
115                socket_device.socket, PW_RPC_MAX_PACKET_SIZE
116            )
117            write_function = socket_device.write
118        except ValueError:
119            _LOG.exception('Failed to initialize socket at %s', socket_addr)
120            return 1
121
122    return HdlcRpcClient(reader, protos, default_channels(write_function))
123
124
125def get_trace_data_from_device(client):
126    """Get the trace data using RPC from a Client"""
127    data = b''
128    service = client.client.channel(1).rpcs.pw.trace.TraceService
129    result = service.GetTraceData().responses
130    for streamed_data in result:
131        data = data + bytes([len(streamed_data.data)])
132        data = data + streamed_data.data
133        _LOG.debug(''.join(format(x, '02x') for x in streamed_data.data))
134    return data
135
136
137def _parse_args():
138    """Parse and return command line arguments."""
139
140    parser = argparse.ArgumentParser(
141        description=__doc__,
142        formatter_class=argparse.RawDescriptionHelpFormatter,
143    )
144    group = parser.add_mutually_exclusive_group(required=True)
145    group.add_argument('-d', '--device', help='the serial port to use')
146    parser.add_argument(
147        '-b',
148        '--baudrate',
149        type=int,
150        default=115200,
151        help='the baud rate to use',
152    )
153    group.add_argument(
154        '-s',
155        '--socket-addr',
156        type=str,
157        help='use socket to connect to server, type default for\
158            localhost:33000, or manually input the server address:port',
159    )
160    parser.add_argument(
161        '-o',
162        '--trace_output',
163        dest='trace_output_file',
164        help=('The json file to which to write the output.'),
165    )
166    parser.add_argument(
167        '-t',
168        '--trace_token_database',
169        help='Databases (ELF, binary, or CSV) to use to lookup trace tokens.',
170    )
171    parser.add_argument(
172        'proto_globs', nargs='+', help='glob pattern for .proto files'
173    )
174    parser.add_argument(
175        '-f',
176        '--ticks_per_second',
177        type=int,
178        dest='ticks_per_second',
179        default=1000,
180        help=('The clock rate of the trace events (Default 1000).'),
181    )
182    parser.add_argument(
183        '--time_offset',
184        type=int,
185        dest='time_offset',
186        default=0,
187        help=('Time offset (us) of the trace events (Default 0).'),
188    )
189    return parser.parse_args()
190
191
192def _main(args):
193    token_database = database.load_token_database(
194        args.trace_token_database, domain="trace"
195    )
196    _LOG.info(database.database_summary(token_database))
197    client = get_hdlc_rpc_client(**vars(args))
198    data = get_trace_data_from_device(client)
199    events = trace_tokenized.get_trace_events(
200        [token_database], data, args.ticks_per_second, args.time_offset
201    )
202    json_lines = trace.generate_trace_json(events)
203    trace_tokenized.save_trace_file(json_lines, args.trace_output_file)
204
205
206if __name__ == '__main__':
207    if sys.version_info[0] < 3:
208        sys.exit('ERROR: The detokenizer command line tools require Python 3.')
209    _main(_parse_args())
210