1#!/usr/bin/env python3
2#
3# Copyright (C) 2023 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18import sys
19import os
20import json
21import argparse
22
23def update_timestamp_perfetto_trace(trace_dict, time_offset):
24    if time_offset == 0:
25        return
26
27    for event in trace_dict['traceEvents']:
28        if 'ts' in event:
29            event['ts'] += time_offset
30
31def update_pid_perfetto_trace(trace_dict, start_pid, max_pid):
32    if start_pid >= max_pid:
33        raise SystemExit(f"Error: start_pid {start_pid} should be smaller than max_pid {max_pid}")
34    if start_pid < 0 or max_pid <= 0:
35        raise SystemExit(f"Error: both start_pid {start_pid} and max_pid {max_pid} should be larger than 0")
36
37    pid_dict = {}
38    current_max_pid = start_pid
39    for event in trace_dict['traceEvents']:
40        if 'pid' in event:
41            old_pid = event['pid']
42            new_pid = pid_dict.get(old_pid)
43            if new_pid == None:
44                if current_max_pid < max_pid:
45                    new_pid = current_max_pid
46                    pid_dict[old_pid] = new_pid
47                    current_max_pid += 1
48                else:
49                    raise SystemExit("Error: due to out of range for allocating pids")
50            event['pid'] = new_pid
51
52def update_trace_file(input_file, time_offset, start_pid=(1<<16), max_pid = (1<<32)):
53    try:
54        with open(input_file, 'r') as f:
55            trace_dict = json.loads(f.read())
56    except Exception as e:
57        print(f'Error: update_trace_file open input file: {input_file} : {e}')
58        return False
59
60    update_timestamp_perfetto_trace(trace_dict, time_offset)
61    update_pid_perfetto_trace(trace_dict, start_pid, max_pid)
62
63    # Save the updated trace data to a new JSON file
64    # add '_updated' to the output filename
65    file_path = os.path.splitext(input_file)
66    output_file = f"{file_path[0]}_updated.json"
67    try:
68        with open(output_file, 'w') as f:
69            json.dump(trace_dict, f)
70    except Exception as e:
71        print(f'Error: update_trace_file open output_file {output_file} : {e}')
72        return False
73
74    print(f"Updated trace data saved to {output_file}")
75    return True
76
77def parseArguments():
78    parser = argparse.ArgumentParser(
79        description='Update perfetto trace event timestamp with the given offset.')
80    parser.add_argument('--input_file', required=True,
81                        help='path to the input Perfetto JSON file')
82    parser.add_argument('--time_offset', type=int, required=False, default=0,
83                        help='time offset value in nanoseconds. If it is 0, timestamp will not be updated.')
84    # At default set the start_pid = (2^16)
85    parser.add_argument('--start_pid', type=int, required=False, default=65536,
86                        help='the smallest pid value')
87    # At default set the max_pid = (2^32)
88    parser.add_argument('--max_pid', type=int, required=False, default=4294967296,
89                        help='the largest pid value. If max_pid == start_[id, pid will not be updated.')
90
91    args = parser.parse_args()
92    return args
93
94if __name__ == '__main__':
95    args = parseArguments();
96    update_trace_file(args.input_file, args.time_offset, args.start_pid, args.max_pid)
97