| #!/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)}') |
