#!/usr/bin/env python3
# Copyright 2021 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
from collections import defaultdict
import argparse
import itertools
import json
import logging
import os
import re
import sys
import typing
"""This module contains the utilities necessary to read Dtrace result files and
convert them other format for flamegraph generation, such as pprof profiles.
"""
from protos.third_party.pprof import profile_pb2
def pairwise(iterable):
"s -> (s0,s1), (s1,s2), (s2, s3), ..."
a, b = itertools.tee(iterable)
next(b, None)
return zip(a, b)
class ProfileBuilder:
"""Helper to generate a pprof profile."""
def __init__(self):
self._profile = profile_pb2.Profile()
self._locations = {}
self._profile.string_table.append("")
self._strings = {"": 0}
self._signature_id = self.GetStringId('signature')
def GetStringId(self, string: str) -> int:
"""Returns the id for `string` into the string_table, creating new entry
if not already present.
"""
if string in self._strings:
return self._strings[string]
index = len(self._profile.string_table)
self._strings[string] = index
self._profile.string_table.append(string)
return index
def GetSymbolLocation(self, name: str, system_name: str) -> int:
"""Returns the id for a symbol location defined as
`(name, system_name)`, creating new entry if not already present.
"""
if (name, system_name) in self._locations:
return self._locations[(name, system_name)]
function_id = len(self._profile.location) + 1
self._locations[(name, system_name)] = function_id
function = self._profile.function.add()
function.id = function_id
function.name = self.GetStringId(name)
function.system_name = self.GetStringId(system_name)
# These fields are given default values since they aren't stored
# in dtrace.
function.filename = self.GetStringId("")
function.start_line = 0
location = self._profile.location.add()
location.id = function_id
line = location.line.add()
line.function_id = function_id
return function_id
def AddComment(self, comment: str):
self._profile.comment.append(self.GetStringId(comment))
def AddSampleType(self, type: str, unit: str):
"""
Adds a sample type that describe stackframes in this profile.
See protos/third_party/pprof/src/profile.proto for more details.
"""
assert (len(self._profile.sample) == 0)
sample_type = self._profile.sample_type.add()
sample_type.type = self.GetStringId(type)
sample_type.unit = self.GetStringId(unit)
def AddSample(self, locations: typing.List[int], values: typing.List[int],
signature: str):
"""
Adds a sample in the profile, constructed from a list of locations
representing the stack, and a list of values for that stack (as many values
as sample types described by this profile).
"""
assert (len(self._profile.sample_type) == len(values))
sample = self._profile.sample.add()
for value in values:
sample.value.append(value)
for location in locations:
sample.location_id.append(location)
label = sample.label.add()
label.key = self._signature_id
label.str = self.GetStringId(signature)
def SerializeToString(self) -> str:
return self._profile.SerializeToString()
class DTraceParser:
"""Parses and merges chromium Dtrace profiles.
Typical usage example:
parser = DTraceParser()
parser.ParseDir('./samples/')
parser.ExportToPprof(builder)
"""
def __init__(self, sample_type: str = 'cpu_time'):
"""
Args:
output_filename: The path of the file in which results are written.
"""
self._stack_weights = defaultdict(int)
self._signatures = defaultdict(str)
self._stack_frames = {}
self._sample_type = sample_type
self._post_processing_applied = False
def ParseFile(self, stack_file: typing.TextIO):
"""Parses dtrace `stack_file` and adds the data to this profile.
"""
assert (self._post_processing_applied == False)
stack_frames = []
for line, next_line in pairwise(stack_file):
line_content = line.strip()
if not line_content:
continue
# If the next line is non-empty it's not the last in the stack.
if next_line.strip():
# Matches lines like: "0x17e018987e" or "+0x17e018987e"
if line_content.startswith("0x") or line_content.startswith(
"+0x"):
function = line_content
module = "unsymbolized module"
else:
module, function = line_content.split('`', 1)
# Matches lines with offset like: "module`function+0xf6"
if len(function.split('+0x')) == 2:
[function, offset] = function.split('+0x')
stack_frames.append((module, function))
else:
if len(stack_frames) == 0:
continue
weight = int(line_content)
stack_string = ";".join([
f'{module}`{function}' for (module,
function) in stack_frames
])
self._stack_weights[stack_string] += weight
self._stack_frames[stack_string] = stack_frames
stack_frames = []
if not self._stack_frames:
logging.error("No results found, check directory contents")
sys.exit(-1)
def ParseDir(self, stack_dir: str):
"""Parses all dtrace files in `stack_dir` and adds the data to this
profile.
Args:
stack_dir: The directory where Dtrace profile results can be found.
Raises:
SystemExit: When no results are found in stack_dir.
"""
# Define the format for DTrace stacks filenames.
stack_filename_regex = re.compile('[0-9]*_[0-9]*.txt')
# Treat all files that respect the name stack_filename_regex as DTrace
# stacks.
for root, dirs, files in os.walk(stack_dir):
for stack_filename in files:
if stack_filename_regex.match(stack_filename):
logging.info(f"Processing {stack_filename} ...")
with open(
os.path.join(root, stack_filename),
newline='',
encoding="ISO-8859-1") as stack_file:
self.ParseFile(stack_file)
if not self._stack_frames:
logging.error("No results found, check directory contents")
sys.exit(-1)
def ConvertToPprof(self, profile_builder: ProfileBuilder):
"""Converts this profile to pprof by writing to
`profile_builder`.
"""
profile_builder.AddSampleType(self._sample_type, "counts")
for key in self._stack_frames:
frames = self._stack_frames[key]
weight = self._stack_weights[key]
signature = self._signatures[key]
sample_locations = []
for (module, function) in frames:
sample_locations.append(
profile_builder.GetSymbolLocation(function, module))
profile_builder.AddSample(sample_locations, [weight], signature)
def ConvertToCollapse(self, output_filename: str):
"""Converts this profile to the "collapsed stack" format. In contrast
to the Dtrace format full stacks are writtent on a single line. At
first the different are separated by semi-colons and a space separates
the weight associated with the function.
Example:
base::foo;content::bar;biz::baz 17
base::biz;content::boo;biz::bim 23
...
"""
os.makedirs(
f"{os.path.dirname(os.path.abspath(output_filename))}",
exist_ok=True)
with open(output_filename, 'w') as f:
for key in self._stack_frames:
frames = self._stack_frames[key]
frames_string = ';'.join(
[function for (module, function) in reversed(frames)])
weight = self._stack_weights[key]
# Reform the line in stacked format and write it out.
f.write(f"{frames_string} {weight}\n")
def GetSamplesListForTesting(self):
samples = []
for key in self._stack_frames:
frames = self._stack_frames[key]
weight = self._stack_weights[key]
samples.append({"frames": frames, "weight": weight})
return samples
def AddSamplesForTesting(self, samples):
for sample in samples:
stack_frames = sample['frames']
stack_string = ";".join(
[f'{module}`{function}' for (module, function) in stack_frames])
self._stack_frames[stack_string] = stack_frames
self._stack_weights[stack_string] += sample['weight']
def ApplySignatures(self, stack: typing.List[typing.Tuple[str, str]]):
"""Matches and return known signatures to given stackframe.
"""
if len(stack) >= 512:
return '_OVERFLOWED_'
for module, function in stack:
if function.startswith('safe_browsing::(anonymous namespace)::'
'PlaybackOnBackgroundThread'):
return 'safe_browsing:VisualSignatures'
if function.startswith(
'safe_browsing::(anonymous namespace)::OnModelInputCreated'
):
return 'safe_browsing:VisualSignatures'
if function.startswith('ParkableStringImpl::CompressInBackground'):
return 'ParkableString'
return 'unknown'
def PostProcessStackSamples(self):
"""Applies filtering and enhancing to self.samples(). This function can
only be called once.
Raises:
SystemExit: If this function is called twice on the same object.
"""
if self._post_processing_applied:
logging.error("Post processing cannot be applied twice")
sys.exit(-1)
self._post_processing_applied = True
for key in self._stack_frames:
# Signatures are always added since they are non destructive.
self._signatures[key] = self.ApplySignatures(
self._stack_frames[key])
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description='Export DTrace stack files into another format.')
parser.add_argument(
"--data_dir",
help="Top level directory that contains DTrace stacks. "
"The directory will be fully walked to find stacks "
"and metadata.json files",
required=True)
parser.add_argument(
"--output", help="The file to write the collapsed stacks into.")
parser.add_argument(
'--unit',
dest='unit',
choices=["cpu_samples", "wakeups"],
default="cpu_samples",
help="The unit of counts acquired with DTrace")
parser.add_argument(
'--format',
dest='format',
action='store',
choices=["pprof", "collapsed"],
default="pprof",
help="Output format to generate.")
args = parser.parse_args()
logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.INFO)
# Traverse |data_dir| and concatenate all metadata.json into one big comment
# for the pprof.
full_comment = ""
metadata_filename_regex = re.compile('metadata.json')
for root, dirs, files in os.walk(args.data_dir):
for file in files:
if metadata_filename_regex.match(file):
with open(os.path.join(root, file)) as meta_json:
for line in meta_json:
full_comment += line
parser = DTraceParser(args.unit)
parser.ParseDir(args.data_dir)
parser.PostProcessStackSamples()
output_filename = args.output
data_dir = os.path.abspath(os.path.join(args.data_dir, os.pardir))
if args.format == "pprof":
profile_builder = ProfileBuilder()
profile_builder.AddComment(full_comment)
profile_builder.AddComment(f"Unit : {args.unit}")
parser.ConvertToPprof(profile_builder)
if output_filename is None:
output_filename = os.path.join(data_dir, f"profile_{args.unit}.pb")
with open(output_filename, "wb") as output_file:
output_file.write(profile_builder.SerializeToString())
else:
if output_filename is None:
output_filename = os.path.join(data_dir,
f"profile_{args.unit}.collapsed")
parser.ConvertToCollapse(output_filename)
logging.info(f'Outputing profile in {os.path.abspath(output_filename)}')
Codebase Browser
chromium