1# gecko.py - Convert perf record output to Firefox's gecko profile format
2# SPDX-License-Identifier: GPL-2.0
3#
4# The script converts perf.data to Gecko Profile Format,
5# which can be read by https://profiler.firefox.com/.
6#
7# Usage:
8#
9#     perf record -a -g -F 99 sleep 60
10#     perf script report gecko
11#
12# Combined:
13#
14#     perf script gecko -F 99 -a sleep 60
15
16import os
17import sys
18import time
19import json
20import string
21import random
22import argparse
23import threading
24import webbrowser
25import urllib.parse
26from os import system
27from functools import reduce
28from dataclasses import dataclass, field
29from http.server import HTTPServer, SimpleHTTPRequestHandler, test
30from typing import List, Dict, Optional, NamedTuple, Set, Tuple, Any
31
32# Add the Perf-Trace-Util library to the Python path
33sys.path.append(os.environ['PERF_EXEC_PATH'] + \
34	'/scripts/python/Perf-Trace-Util/lib/Perf/Trace')
35
36from perf_trace_context import *
37from Core import *
38
39StringID = int
40StackID = int
41FrameID = int
42CategoryID = int
43Milliseconds = float
44
45# start_time is intialiazed only once for the all event traces.
46start_time = None
47
48# https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/profile.js#L425
49# Follow Brendan Gregg's Flamegraph convention: orange for kernel and yellow for user space by default.
50CATEGORIES = None
51
52# The product name is used by the profiler UI to show the Operating system and Processor.
53PRODUCT = os.popen('uname -op').read().strip()
54
55# store the output file
56output_file = None
57
58# Here key = tid, value = Thread
59tid_to_thread = dict()
60
61# The HTTP server is used to serve the profile to the profiler UI.
62http_server_thread = None
63
64# The category index is used by the profiler UI to show the color of the flame graph.
65USER_CATEGORY_INDEX = 0
66KERNEL_CATEGORY_INDEX = 1
67
68# https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/gecko-profile.js#L156
69class Frame(NamedTuple):
70	string_id: StringID
71	relevantForJS: bool
72	innerWindowID: int
73	implementation: None
74	optimizations: None
75	line: None
76	column: None
77	category: CategoryID
78	subcategory: int
79
80# https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/gecko-profile.js#L216
81class Stack(NamedTuple):
82	prefix_id: Optional[StackID]
83	frame_id: FrameID
84
85# https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/gecko-profile.js#L90
86class Sample(NamedTuple):
87	stack_id: Optional[StackID]
88	time_ms: Milliseconds
89	responsiveness: int
90
91@dataclass
92class Thread:
93	"""A builder for a profile of the thread.
94
95	Attributes:
96		comm: Thread command-line (name).
97		pid: process ID of containing process.
98		tid: thread ID.
99		samples: Timeline of profile samples.
100		frameTable: interned stack frame ID -> stack frame.
101		stringTable: interned string ID -> string.
102		stringMap: interned string -> string ID.
103		stackTable: interned stack ID -> stack.
104		stackMap: (stack prefix ID, leaf stack frame ID) -> interned Stack ID.
105		frameMap: Stack Frame string -> interned Frame ID.
106		comm: str
107		pid: int
108		tid: int
109		samples: List[Sample] = field(default_factory=list)
110		frameTable: List[Frame] = field(default_factory=list)
111		stringTable: List[str] = field(default_factory=list)
112		stringMap: Dict[str, int] = field(default_factory=dict)
113		stackTable: List[Stack] = field(default_factory=list)
114		stackMap: Dict[Tuple[Optional[int], int], int] = field(default_factory=dict)
115		frameMap: Dict[str, int] = field(default_factory=dict)
116	"""
117	comm: str
118	pid: int
119	tid: int
120	samples: List[Sample] = field(default_factory=list)
121	frameTable: List[Frame] = field(default_factory=list)
122	stringTable: List[str] = field(default_factory=list)
123	stringMap: Dict[str, int] = field(default_factory=dict)
124	stackTable: List[Stack] = field(default_factory=list)
125	stackMap: Dict[Tuple[Optional[int], int], int] = field(default_factory=dict)
126	frameMap: Dict[str, int] = field(default_factory=dict)
127
128	def _intern_stack(self, frame_id: int, prefix_id: Optional[int]) -> int:
129		"""Gets a matching stack, or saves the new stack. Returns a Stack ID."""
130		key = f"{frame_id}" if prefix_id is None else f"{frame_id},{prefix_id}"
131		# key = (prefix_id, frame_id)
132		stack_id = self.stackMap.get(key)
133		if stack_id is None:
134			# return stack_id
135			stack_id = len(self.stackTable)
136			self.stackTable.append(Stack(prefix_id=prefix_id, frame_id=frame_id))
137			self.stackMap[key] = stack_id
138		return stack_id
139
140	def _intern_string(self, string: str) -> int:
141		"""Gets a matching string, or saves the new string. Returns a String ID."""
142		string_id = self.stringMap.get(string)
143		if string_id is not None:
144			return string_id
145		string_id = len(self.stringTable)
146		self.stringTable.append(string)
147		self.stringMap[string] = string_id
148		return string_id
149
150	def _intern_frame(self, frame_str: str) -> int:
151		"""Gets a matching stack frame, or saves the new frame. Returns a Frame ID."""
152		frame_id = self.frameMap.get(frame_str)
153		if frame_id is not None:
154			return frame_id
155		frame_id = len(self.frameTable)
156		self.frameMap[frame_str] = frame_id
157		string_id = self._intern_string(frame_str)
158
159		symbol_name_to_category = KERNEL_CATEGORY_INDEX if frame_str.find('kallsyms') != -1 \
160		or frame_str.find('/vmlinux') != -1 \
161		or frame_str.endswith('.ko)') \
162		else USER_CATEGORY_INDEX
163
164		self.frameTable.append(Frame(
165			string_id=string_id,
166			relevantForJS=False,
167			innerWindowID=0,
168			implementation=None,
169			optimizations=None,
170			line=None,
171			column=None,
172			category=symbol_name_to_category,
173			subcategory=None,
174		))
175		return frame_id
176
177	def _add_sample(self, comm: str, stack: List[str], time_ms: Milliseconds) -> None:
178		"""Add a timestamped stack trace sample to the thread builder.
179		Args:
180			comm: command-line (name) of the thread at this sample
181			stack: sampled stack frames. Root first, leaf last.
182			time_ms: timestamp of sample in milliseconds.
183		"""
184		# Ihreads may not set their names right after they are created.
185		# Instead, they might do it later. In such situations, to use the latest name they have set.
186		if self.comm != comm:
187			self.comm = comm
188
189		prefix_stack_id = reduce(lambda prefix_id, frame: self._intern_stack
190						(self._intern_frame(frame), prefix_id), stack, None)
191		if prefix_stack_id is not None:
192			self.samples.append(Sample(stack_id=prefix_stack_id,
193									time_ms=time_ms,
194									responsiveness=0))
195
196	def _to_json_dict(self) -> Dict:
197		"""Converts current Thread to GeckoThread JSON format."""
198		# Gecko profile format is row-oriented data as List[List],
199		# And a schema for interpreting each index.
200		# Schema:
201		# https://github.com/firefox-devtools/profiler/blob/main/docs-developer/gecko-profile-format.md
202		# https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/gecko-profile.js#L230
203		return {
204			"tid": self.tid,
205			"pid": self.pid,
206			"name": self.comm,
207			# https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/gecko-profile.js#L51
208			"markers": {
209				"schema": {
210					"name": 0,
211					"startTime": 1,
212					"endTime": 2,
213					"phase": 3,
214					"category": 4,
215					"data": 5,
216				},
217				"data": [],
218			},
219
220			# https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/gecko-profile.js#L90
221			"samples": {
222				"schema": {
223					"stack": 0,
224					"time": 1,
225					"responsiveness": 2,
226				},
227				"data": self.samples
228			},
229
230			# https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/gecko-profile.js#L156
231			"frameTable": {
232				"schema": {
233					"location": 0,
234					"relevantForJS": 1,
235					"innerWindowID": 2,
236					"implementation": 3,
237					"optimizations": 4,
238					"line": 5,
239					"column": 6,
240					"category": 7,
241					"subcategory": 8,
242				},
243				"data": self.frameTable,
244			},
245
246			# https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/gecko-profile.js#L216
247			"stackTable": {
248				"schema": {
249					"prefix": 0,
250					"frame": 1,
251				},
252				"data": self.stackTable,
253			},
254			"stringTable": self.stringTable,
255			"registerTime": 0,
256			"unregisterTime": None,
257			"processType": "default",
258		}
259
260# Uses perf script python interface to parse each
261# event and store the data in the thread builder.
262def process_event(param_dict: Dict) -> None:
263	global start_time
264	global tid_to_thread
265	time_stamp = (param_dict['sample']['time'] // 1000) / 1000
266	pid = param_dict['sample']['pid']
267	tid = param_dict['sample']['tid']
268	comm = param_dict['comm']
269
270	# Start time is the time of the first sample
271	if not start_time:
272		start_time = time_stamp
273
274	# Parse and append the callchain of the current sample into a stack.
275	stack = []
276	if param_dict['callchain']:
277		for call in param_dict['callchain']:
278			if 'sym' not in call:
279				continue
280			stack.append(f'{call["sym"]["name"]} (in {call["dso"]})')
281		if len(stack) != 0:
282			# Reverse the stack, as root come first and the leaf at the end.
283			stack = stack[::-1]
284
285	# During perf record if -g is not used, the callchain is not available.
286	# In that case, the symbol and dso are available in the event parameters.
287	else:
288		func = param_dict['symbol'] if 'symbol' in param_dict else '[unknown]'
289		dso = param_dict['dso'] if 'dso' in param_dict else '[unknown]'
290		stack.append(f'{func} (in {dso})')
291
292	# Add sample to the specific thread.
293	thread = tid_to_thread.get(tid)
294	if thread is None:
295		thread = Thread(comm=comm, pid=pid, tid=tid)
296		tid_to_thread[tid] = thread
297	thread._add_sample(comm=comm, stack=stack, time_ms=time_stamp)
298
299def trace_begin() -> None:
300	global output_file
301	if (output_file is None):
302		print("Staring Firefox Profiler on your default browser...")
303		global http_server_thread
304		http_server_thread = threading.Thread(target=test, args=(CORSRequestHandler, HTTPServer,))
305		http_server_thread.daemon = True
306		http_server_thread.start()
307
308# Trace_end runs at the end and will be used to aggregate
309# the data into the final json object and print it out to stdout.
310def trace_end() -> None:
311	global output_file
312	threads = [thread._to_json_dict() for thread in tid_to_thread.values()]
313
314	# Schema: https://github.com/firefox-devtools/profiler/blob/53970305b51b9b472e26d7457fee1d66cd4e2737/src/types/gecko-profile.js#L305
315	gecko_profile_with_meta = {
316		"meta": {
317			"interval": 1,
318			"processType": 0,
319			"product": PRODUCT,
320			"stackwalk": 1,
321			"debug": 0,
322			"gcpoison": 0,
323			"asyncstack": 1,
324			"startTime": start_time,
325			"shutdownTime": None,
326			"version": 24,
327			"presymbolicated": True,
328			"categories": CATEGORIES,
329			"markerSchema": [],
330			},
331		"libs": [],
332		"threads": threads,
333		"processes": [],
334		"pausedRanges": [],
335	}
336	# launch the profiler on local host if not specified --save-only args, otherwise print to file
337	if (output_file is None):
338		output_file = 'gecko_profile.json'
339		with open(output_file, 'w') as f:
340			json.dump(gecko_profile_with_meta, f, indent=2)
341		launchFirefox(output_file)
342		time.sleep(1)
343		print(f'[ perf gecko: Captured and wrote into {output_file} ]')
344	else:
345		print(f'[ perf gecko: Captured and wrote into {output_file} ]')
346		with open(output_file, 'w') as f:
347			json.dump(gecko_profile_with_meta, f, indent=2)
348
349# Used to enable Cross-Origin Resource Sharing (CORS) for requests coming from 'https://profiler.firefox.com', allowing it to access resources from this server.
350class CORSRequestHandler(SimpleHTTPRequestHandler):
351	def end_headers (self):
352		self.send_header('Access-Control-Allow-Origin', 'https://profiler.firefox.com')
353		SimpleHTTPRequestHandler.end_headers(self)
354
355# start a local server to serve the gecko_profile.json file to the profiler.firefox.com
356def launchFirefox(file):
357	safe_string = urllib.parse.quote_plus(f'http://localhost:8000/{file}')
358	url = 'https://profiler.firefox.com/from-url/' + safe_string
359	webbrowser.open(f'{url}')
360
361def main() -> None:
362	global output_file
363	global CATEGORIES
364	parser = argparse.ArgumentParser(description="Convert perf.data to Firefox\'s Gecko Profile format which can be uploaded to profiler.firefox.com for visualization")
365
366	# Add the command-line options
367	# Colors must be defined according to this:
368	# https://github.com/firefox-devtools/profiler/blob/50124adbfa488adba6e2674a8f2618cf34b59cd2/res/css/categories.css
369	parser.add_argument('--user-color', default='yellow', help='Color for the User category', choices=['yellow', 'blue', 'purple', 'green', 'orange', 'red', 'grey', 'magenta'])
370	parser.add_argument('--kernel-color', default='orange', help='Color for the Kernel category', choices=['yellow', 'blue', 'purple', 'green', 'orange', 'red', 'grey', 'magenta'])
371	# If --save-only is specified, the output will be saved to a file instead of opening Firefox's profiler directly.
372	parser.add_argument('--save-only', help='Save the output to a file instead of opening Firefox\'s profiler')
373
374	# Parse the command-line arguments
375	args = parser.parse_args()
376	# Access the values provided by the user
377	user_color = args.user_color
378	kernel_color = args.kernel_color
379	output_file = args.save_only
380
381	CATEGORIES = [
382		{
383			"name": 'User',
384			"color": user_color,
385			"subcategories": ['Other']
386		},
387		{
388			"name": 'Kernel',
389			"color": kernel_color,
390			"subcategories": ['Other']
391		},
392	]
393
394if __name__ == '__main__':
395	main()
396