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

2 files changed, 48 insertions, 47 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))
diff --git a/tools/jvmti-agents/ti-alloc-sample/ti_alloc_sample.cc b/tools/jvmti-agents/ti-alloc-sample/ti_alloc_sample.cc
index 862aa51b86..d719db5af8 100644
--- a/tools/jvmti-agents/ti-alloc-sample/ti_alloc_sample.cc
+++ b/tools/jvmti-agents/ti-alloc-sample/ti_alloc_sample.cc
@@ -293,31 +293,15 @@ static std::string formatAllocation(jvmtiEnv* jvmti,
 
 // Formatter for a method entry on a call stack.
 static std::string formatMethod(jvmtiEnv* jvmti, JNIEnv* jni, jmethodID method_id) {
-  jclass declaring_class;
-  jvmtiError err = jvmti->GetMethodDeclaringClass(method_id, &declaring_class);
-  std::string class_name = "";
-  if (err == JVMTI_ERROR_NONE) {
-    char* class_signature;
-    err = jvmti->GetClassSignature(declaring_class,
-                                   &class_signature,
-                                   /*generic_ptr*/ nullptr);
-    class_name = std::string(class_signature);
-    jvmti->Deallocate(reinterpret_cast<unsigned char*>(class_signature));
-    jni->DeleteLocalRef(declaring_class);
-  }
-  char *method_name;
-  char *method_signature;
-  char *generic_pointer;
-  err = jvmti->GetMethodName(method_id,
-                             &method_name,
-                             &method_signature,
-                             &generic_pointer);
-  std::string method = "METHODERROR";
-  if (err == JVMTI_ERROR_NONE) {
-    method = ((method_name == nullptr) ? "UNKNOWN" : method_name);
-    method += ((method_signature == nullptr) ? "(UNKNOWN)" : method_signature);
+  ScopedMethodInfo smi(jvmti, jni, method_id);
+  std::string method;
+  if (smi.Init(/*get_generic=*/false)) {
+    method = std::string(smi.GetDeclaringClassInfo().GetName()) +
+        "::" + smi.GetName() + smi.GetSignature();
+  } else {
+    method = "ERROR";
   }
-  return string_table->Intern("+", class_name + "::" + method);
+  return string_table->Intern("+", method);
 }
 
 static int sampling_rate;