xref: /aosp_15_r20/external/pigweed/pw_system/py/pw_system/trace_client.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1#!/usr/bin/env python3
2# Copyright 2023 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 trace service.
18
19Example usage:
20python pw_console/py/pw_console/trace_client.py
21  -o trace.json
22  -t out/host_device_simulator.speed_optimized/obj/pw_system/bin/system_example
23"""
24
25import argparse
26import logging
27import sys
28from pathlib import Path
29from types import ModuleType
30
31
32from pw_transfer import transfer_pb2
33from pw_log.proto import log_pb2
34from pw_trace_protos import trace_service_pb2
35from pw_trace import trace
36from pw_trace_tokenized import trace_tokenized
37import pw_transfer
38from pw_file import file_pb2
39from pw_hdlc import rpc
40from pw_stream import stream_readers
41from pw_system import device_tracing
42from pw_tokenizer import detokenize
43from pw_console import socket_client
44
45
46_LOG = logging.getLogger(__package__)
47_LOG.level = logging.DEBUG
48_LOG.addHandler(logging.StreamHandler(sys.stdout))
49
50
51def start_tracing_on_device(client):
52    """Start tracing on the device"""
53    service = client.rpcs.pw.trace.proto.TraceService
54    service.Start()
55
56
57def stop_tracing_on_device(client):
58    """Stop tracing on the device"""
59    service = client.rpcs.pw.trace.proto.TraceService
60    return service.Stop()
61
62
63def list_files_on_device(client):
64    """List files on the device"""
65    service = client.rpcs.pw.file.FileSystem
66    return service.List()
67
68
69def delete_file_on_device(client, path):
70    """Delete a file on the device"""
71    service = client.rpcs.pw.file.FileSystem
72    req = file_pb2.DeleteRequest(path=path)
73    return service.Delete(req)
74
75
76def _parse_args():
77    """Parse and return command line arguments."""
78
79    parser = argparse.ArgumentParser(
80        description=__doc__,
81        formatter_class=argparse.RawDescriptionHelpFormatter,
82    )
83    group = parser.add_mutually_exclusive_group(required=False)
84    group.add_argument('-d', '--device', help='the serial port to use')
85    parser.add_argument(
86        '-b',
87        '--baudrate',
88        type=int,
89        default=115200,
90        help='The baud rate to use',
91    )
92    group.add_argument(
93        '-s',
94        '--socket-addr',
95        type=str,
96        default='default',
97        help=(
98            'Use socket to connect to server, type default for '
99            'localhost:33000, or manually input the server address:port'
100        ),
101    )
102    parser.add_argument(
103        '-o',
104        '--trace_output',
105        dest='trace_output_file',
106        help='The json file to which to write the output.',
107    )
108    parser.add_argument(
109        '-t',
110        '--trace_token_database',
111        help='Databases (ELF, binary, or CSV) to use to lookup trace tokens.',
112    )
113    parser.add_argument(
114        '-f',
115        '--ticks_per_second',
116        type=int,
117        dest='ticks_per_second',
118        help='The clock rate of the trace events.',
119    )
120    parser.add_argument(
121        '--time_offset',
122        type=int,
123        dest='time_offset',
124        default=0,
125        help='Time offset (us) of the trace events (Default 0).',
126    )
127    parser.add_argument(
128        '--channel-id',
129        type=int,
130        dest='channel_id',
131        default=rpc.DEFAULT_CHANNEL_ID,
132        help='Channel ID used in RPC communications.',
133    )
134    return parser.parse_args()
135
136
137def _main(args) -> int:
138    detokenizer = detokenize.AutoUpdatingDetokenizer(
139        args.trace_token_database + '#trace'
140    )
141    detokenizer.show_errors = True
142
143    socket_impl = socket_client.SocketClient
144    try:
145        socket_device = socket_impl(args.socket_addr)
146        reader = stream_readers.SelectableReader(socket_device)
147        write = socket_device.write
148    except ValueError:
149        _LOG.exception('Failed to initialize socket at %s', args.socket_addr)
150        return 1
151
152    protos: list[ModuleType | Path] = [
153        log_pb2,
154        file_pb2,
155        transfer_pb2,
156        trace_service_pb2,
157    ]
158
159    with reader:
160        device_client = device_tracing.DeviceWithTracing(
161            args.channel_id,
162            reader,
163            write,
164            protos,
165            detokenizer=detokenizer,
166            timestamp_decoder=None,
167            rpc_timeout_s=5,
168            use_rpc_logging=True,
169            use_hdlc_encoding=True,
170            ticks_per_second=args.ticks_per_second,
171        )
172
173        with device_client:
174            _LOG.info('Starting tracing')
175            start_tracing_on_device(device_client)
176
177            _LOG.info('Stopping tracing')
178            file_id = stop_tracing_on_device(device_client)
179            _LOG.info('Trace file id = %d', file_id.response.file_id)
180
181            _LOG.info('Listing Files')
182            stream_response = list_files_on_device(device_client)
183
184            if not stream_response.status.ok():
185                _LOG.error('Failed to list files %s', stream_response.status)
186                return 1
187
188            for list_response in stream_response.responses:
189                for file in list_response.paths:
190                    _LOG.info('Transfering File: %s', file.path)
191                    try:
192                        data = device_client.transfer_manager.read(file.file_id)
193                        events = trace_tokenized.get_trace_events(
194                            [detokenizer.database],
195                            data,
196                            device_client.ticks_per_second,
197                            args.time_offset,
198                        )
199                        json_lines = trace.generate_trace_json(events)
200                        trace_tokenized.save_trace_file(
201                            json_lines, args.trace_output_file
202                        )
203                    except pw_transfer.Error as err:
204                        print('Failed to read:', err.status)
205
206                    _LOG.info('Deleting File: %s', file.path)
207                    delete_file_on_device(device_client, file.path)
208
209            _LOG.info('All trace transfers completed successfully')
210
211    return 0
212
213
214if __name__ == '__main__':
215    _main(_parse_args())
216