1#!/usr/bin/python3
2#
3# Copyright 2019, 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"""
18Usage: mkflame.py <jvmti_trace_file>
19"""
20
21import argparse
22import sys
23
24class TraceCollection:
25  def __init__(self, args):
26    self.args = args
27    # A table indexed by number and containing the definition for that number.
28    self.definitions = {}
29    # The "weight" of a stack trace, either 1 for counting or the size of the allocation.
30    self.weights = {}
31    # The count for each individual allocation.
32    self.allocation_count = {}
33
34  def definition(self, index):
35    """
36    Returns the definition for "index".
37    """
38    return self.definitions[index]
39
40  def set_definition(self, index, definition):
41    """
42    Sets the definition for "index".
43    """
44    self.definitions[index] = definition
45
46  def weight(self, index):
47    """
48    Returns the weight for "index".
49    """
50    return self.weights[index]
51
52  def set_weight(self, index, weight):
53    """
54    Sets the weight for "index".
55    """
56    self.weights[index] = weight
57
58  def read_file(self, filename):
59    """
60    Reads a file into a DefinitionTable.
61    """
62    def process_definition(line):
63      """
64      Adds line to the list of definitions in table.
65      """
66      def expand_stack_trace(definition):
67        """
68        Converts a semicolon-separated list of numbers into the text stack trace.
69        """
70        def get_allocation_thread(thread_type_size):
71          """
72          Returns the thread of an allocation from the thread/type/size record.
73          """
74          THREAD_STRING = "thread["
75          THREAD_STRING_LEN = len(THREAD_STRING)
76          thread_string = thread_type_size[thread_type_size.find(THREAD_STRING) +
77                                           THREAD_STRING_LEN:]
78          return thread_string[:thread_string.find("]")]
79
80        def get_allocation_type(thread_type_size):
81          """
82          Returns the type of an allocation from the thread/type/size record.
83          """
84          TYPE_STRING = "jclass["
85          TYPE_STRING_LEN = len(TYPE_STRING)
86          type_string = thread_type_size[thread_type_size.find(TYPE_STRING) + TYPE_STRING_LEN:]
87          return type_string[:type_string.find(" ")]
88
89        def get_allocation_size(thread_type_size):
90          """
91          Returns the size of an allocation from the thread/type/size record.
92          """
93          SIZE_STRING = "size["
94          SIZE_STRING_LEN = len(SIZE_STRING)
95          size_string = thread_type_size[thread_type_size.find(SIZE_STRING) + SIZE_STRING_LEN:]
96          size_string = size_string[:size_string.find(",")]
97          return int(size_string)
98
99        def get_top_and_weight(index):
100          thread_type_size = self.definition(int(tokens[0]))
101          size = get_allocation_size(thread_type_size)
102          if self.args.type_only:
103            thread_type_size = get_allocation_type(thread_type_size)
104          elif self.args.thread_only:
105            thread_type_size = get_allocation_thread(thread_type_size)
106          return (thread_type_size, size)
107
108        tokens = definition.split(";")
109        # The first element (base) of the stack trace is the thread/type/size.
110        # Get the weight (either 1 or the number of bytes allocated).
111        (thread_type_size, weight) = get_top_and_weight(int(tokens[0]))
112        self.set_weight(index, weight)
113        # Remove the thread/type/size from the base of the stack trace.
114        del tokens[0]
115        # Build the stack trace list.
116        expanded_definition = ""
117        for i in range(len(tokens)):
118          if self.args.depth_limit > 0 and i >= self.args.depth_limit:
119            break
120          token = tokens[i]
121          # Replace semicolons by colons in the method entry signatures.
122          method = self.definition(int(token)).replace(";", ":")
123          if len(expanded_definition) > 0:
124            expanded_definition += ";"
125          expanded_definition += method
126        if not self.args.ignore_type:
127          # Add the thread/type/size as the top-most stack frame.
128          if len(expanded_definition) > 0:
129            expanded_definition += ";"
130          expanded_definition += thread_type_size.replace(";", ":")
131        if self.args.reverse_stack:
132          def_list = expanded_definition.split(";")
133          expanded_definition = ";".join(def_list[::-1])
134        return expanded_definition
135
136      # If the line contains a comma, it is of the form [+=]index,definition,
137      # where index is a string containing an integer, and definition is the
138      # value represented by the integer whenever it is used later.
139      # * Lines starting with + are either a thread/type/size record or a single
140      #   stack frame.  These are simply interned in the table.
141      # * Those starting with = are stack traces, and contain a sequence of
142      #   numbers separated by semicolon.  These are "expanded" and then interned.
143      comma_pos = line.find(",")
144      index = int(line[1:comma_pos])
145      definition = line[comma_pos+1:]
146      if line[0:1] == "=":
147        definition = expand_stack_trace(definition)
148      # Intern the definition in the table.
149      #if len(definition) == 0:
150        # Zero length samples are errors and are discarded.
151        #print("ERROR: definition for " + str(index) + " is empty")
152        #return
153      self.set_definition(index, definition)
154
155    def process_trace(index):
156      """
157      Remembers one stack trace in the list of stack traces we have seen.
158      Remembering a stack trace increments a count associated with the trace.
159      """
160      trace = self.definition(index)
161      if self.args.use_size:
162        weight = self.weight(index)
163      else:
164        weight = 1
165      if trace in self.allocation_count:
166        self.allocation_count[trace] = self.allocation_count[trace] + weight
167      else:
168        self.allocation_count[trace] = weight
169
170    # Read the file, processing each line as a definition or stack trace.
171    tracefile = open(filename, "r")
172    current_allocation_trace = ""
173    for line in tracefile:
174      line = line.rstrip("\n")
175      if line[0:1] == "=" or line[0:1] == "+":
176        # definition.
177        process_definition(line)
178      else:
179        # stack trace.
180        process_trace(int(line))
181
182  def dump_flame_graph(self):
183    """
184    Prints out a stack trace format compatible with flame graph creation utilities.
185    """
186    for definition, weight in self.allocation_count.items():
187      print(definition + " " + str(weight))
188
189def parse_options():
190  parser = argparse.ArgumentParser(description="Convert a trace to a form usable for flame graphs.")
191  parser.add_argument("filename", help="The trace file as input", type=str)
192  parser.add_argument("--use_size", help="Count by allocation size", action="store_true",
193                      default=False)
194  parser.add_argument("--ignore_type", help="Ignore type of allocation", action="store_true",
195                      default=False)
196  parser.add_argument("--reverse_stack", help="Reverse root and top of stacks", action="store_true",
197                      default=False)
198  parser.add_argument("--type_only", help="Only consider allocation type", action="store_true",
199                      default=False)
200  parser.add_argument("--thread_only", help="Only consider allocation thread", action="store_true",
201                      default=False)
202  parser.add_argument("--depth_limit", help="Limit the length of a trace", type=int, default=0)
203  args = parser.parse_args()
204  return args
205
206def main(argv):
207  args = parse_options()
208  trace_collection = TraceCollection(args)
209  trace_collection.read_file(args.filename)
210  trace_collection.dump_flame_graph()
211
212if __name__ == '__main__':
213  sys.exit(main(sys.argv))
214