Add size weighting to the script

Add the ability to weight the counts by the total size allocated at that
stack trace.  Also puts the thread/type/size at the top of the stack
trace and fixes a bug that erroneously inserted an additional stack
frame at the base.  Some small code cleanups in the agent also.

Bug: none
Test: attach, run app, pull, process, enjoy the flame graph.
Change-Id: I6256a99c2651696b973c25a5955e3bb9251b86a1

1 files changed, 40 insertions, 23 deletions

diff --git a/tools/jvmti-agents/ti-alloc-sample/mkflame.py b/tools/jvmti-agents/ti-alloc-sample/mkflame.py
index 8f1dccfd99..41a29067b8 100755
--- a/tools/jvmti-agents/ti-alloc-sample/mkflame.py
+++ b/tools/jvmti-agents/ti-alloc-sample/mkflame.py
@@ -21,6 +21,14 @@ Usage: mkflame.py <jvmti_trace_file>
 import sys
 
 table = {}
+weights = {}
+
+def get_size(thread_type_size):
+  SIZE_STRING = "size["
+  SIZE_STRING_LEN = len(SIZE_STRING)
+  size_string = thread_type_size[thread_type_size.find(SIZE_STRING) + SIZE_STRING_LEN:]
+  size_string = size_string[:size_string.find(",")]
+  return int(size_string)
 
 def add_definition_to_table(line):
   """
@@ -29,24 +37,25 @@ def add_definition_to_table(line):
   comma_pos = line.find(",")
   index = int(line[1:comma_pos])
   definition = line[comma_pos+1:]
+  expanded_definition = ""
+  weights[index] = 1
   if line[0:1] == "=":
-    # Skip the type/size prefix for flame graphs.
-    semi_pos = definition.find(";")
-    definition = definition[semi_pos + 1:]
-    # Expand stack frame definitions to be a semicolon-separated list of stack
-    # frame methods.
-    expanded_definition = ""
-    while definition != "":
-      semi_pos = definition.find(";")
-      if semi_pos == -1:
-        method_index = int(definition)
-        definition = ""
-      else:
-        method_index = int(definition[:semi_pos])
-        definition = definition[semi_pos + 1:]
+    tokens = definition.split(";")
+    # Pick the thread/type/size off the front (base) of the stack trace.
+    thread_type_size = lookup_definition(int(tokens[0])).replace(";", ":")
+    weights[index] = get_size(thread_type_size)
+    del tokens[0]
+    # Build the stack trace list.
+    for token in tokens:
       # Replace semicolons by colons in the method entry signatures.
-      method = lookup_definition(method_index).replace(";", ":")
-      expanded_definition += ";" + method
+      method = lookup_definition(int(token)).replace(";", ":")
+      if len(expanded_definition) > 0:
+        expanded_definition += ";"
+      expanded_definition += method
+    # Add the thread/type/size as the top-most stack frame.
+    if len(expanded_definition) > 0:
+      expanded_definition += ";"
+    expanded_definition += thread_type_size
     definition = expanded_definition
   table[index] = definition
 
@@ -57,34 +66,42 @@ def lookup_definition(index):
   return table[index]
 
 traces = {}
-def record_stack_trace(string):
+def record_stack_trace(string, count_or_size):
   """
   Remembers one stack trace index in the list of stack traces we have seen.
   Remembering a stack trace increments a count associated with the trace.
   """
   index = int(string)
+  if count_or_size == "size":
+    weight = weights[index]
+  else:
+    weight = 1
   if index in traces:
     count = traces[index]
-    traces[index] = count + 1
+    traces[index] = count + weight
   else:
-    traces[index] = 1
+    traces[index] = weight
 
 def main(argv):
-  filename = argv[1]
+  count_or_size = argv[1]
+  filename = argv[2]
   pagefile = open(filename, "r")
   current_allocation_trace = ""
   for line in pagefile:
-    args = line.split()
     line = line.rstrip("\n")
     if line[0:1] == "=" or line[0:1] == "+":
       # definition.
       add_definition_to_table(line)
     else:
       # stack trace.
-      record_stack_trace(line)
+      record_stack_trace(line, count_or_size)
   # Dump all the traces, with count.
   for k, v in traces.items():
-    print(lookup_definition(k) + " " + str(v))
+    definition = lookup_definition(k)
+    if len(definition) == 0:
+      # Zero length samples are discarded.
+      return
+    print(definition + " " + str(v))
 
 if __name__ == '__main__':
   sys.exit(main(sys.argv))