hprof

#!/usr/bin/env python

"""Convert a java hprof cpu profile dump to a useful format.

Usage:
java -agentlib:hprof=cpu=samples,depth=128 com.friendfeed.ClassName <args>*
hprof|less
"""

from __future__ import with_statement

import os
import pprint
import sys
import re
import optparse


class Frame(object):
    def __init__(self, method, loc="<unknown>:-1"):
        self.method = method
        self.loc = loc


class Node(object):
    frame = Frame(None, -1)

    def __init__(self, trace=(), parent=None):
        trace = trace or []
        self.trace = trace
        if trace:
            self.frame = trace[-1]
        self.count = 0
        self.traces = 0
        self.parent = parent
        self.children = {}
        self.method = self.frame.method


def parse(input):

    traces = {}

    ignore_classes = frozenset([
        "sun.reflect.NativeMethodAccessorImpl",
        "sun.reflect.DelegatingMethodAccessorImpl",
        ])

    def ignore_method(method):
        cls, _, mname = method.rpartition(".")
        if cls in ignore_classes: return True
        if cls.startswith("sun.reflect.GeneratedMethodAccessor"): return True
        return False

    line = input.readline()
    while line != "--------\n":
        line = input.readline()

    line = input.readline()
    while line:
        if line.startswith("TRACE "):
            # TRACE 300591:
            #         sun.nio.ch.EPollArrayWrapper.epollWait(EPollArrayWrapper.java:Unknown line)
            #         sun.nio.ch.EPollArrayWrapper.poll(EPollArrayWrapper.java:215)

            m = re.match(r"^TRACE (\d+):", line)
            trace_num = int(m.group(1))
            trace = []

            line = input.readline()
            while line and line.startswith("\t"):
                line = line.strip()
                pos = line.find("(")
                if pos > 0:
                    method = line[:pos]
                    loc = line[pos+1:-1]
                    if not ignore_method(method):
                        trace.append(Frame(method=method, loc=loc))
                line = input.readline()

            traces[trace_num] = trace

        elif line.startswith("CPU SAMPLES BEGIN"):
            # CPU SAMPLES BEGIN (total = 17984) Tue Nov 18 16:38:08 2008
            # rank   self  accum   count trace method
            #    1 23.91% 23.91%    4300 300591 sun.nio.ch.EPollArrayWrapper.epollWait
            # CPU SAMPLES END

            m = re.match(r"CPU SAMPLES BEGIN \(total = (\d+)\)(.*)", line)
            total_count = int(m.group(1))
            date = m.group(2)

            trace_counts = {}

            input.readline()  # read column headers
            line = input.readline()
            while line and line != "CPU SAMPLES END\n":
                (rank, pct, cpct, count, tnum, method) = line.strip().split()
                trace_counts[int(tnum)] = int(count)
                line = input.readline()

            root = build_tree(traces, trace_counts)
            write_profile(root)

            break

        else:
            line = input.readline()


def build_tree(traces, trace_counts):

    ignore_set = set([
        'java.net.SocketInputStream.socketRead0',
        'java.net.SocketOutputStream.socketWrite0',
        'java.net.PlainSocketImpl.socketAvailable',
        'java.net.PlainSocketImpl.socketAccept',
        'sun.nio.ch.EPollArrayWrapper.epollWait',
    ])

    def resolve(tnum):
        trace = traces[tnum]

        # filter out traces that don't involve friendfeed code
        def include_trace(trace):
            for frame in trace:
                if frame.method.startswith('com.friendfeed.'):
                    return True
            return False

        if False and not include_trace(trace):
            return None

        # focus on one method
        if OPTIONS.method:
            for i, frame in enumerate(trace):
                m = frame.method
                #print OPTIONS.method, m
                if m.find(OPTIONS.method) != -1:
                    #trace = trace[:i+1]
                    trace = trace[i:]
                    break
            else:
                trace = []

        # filter out useless stack frames
        def accept_frame(frame):
            skip = ('java.', 'javax.', 'sun.',
                    #'com.friendfeed.io.Encoder.',
                    'com.friendfeed.jsonrpc.server.',
                    )
            return not frame.method.startswith(skip)
        #trace = [f for f in trace[:-1] if accept_frame(f)] + trace[-1:]

        return trace

#     for trace in traces.itervalues():
#         print len(trace), trace[-1].method

    root = Node()
    for tnum, count in trace_counts.iteritems():

        trace = resolve(tnum)
        if trace is None:
            continue

        if trace[0].method in ignore_set:
            continue
            #root.count -= count

        node = root
        node.count += count
        node.traces += 1

        if OPTIONS.tree:
            trace = list(reversed(trace))

        if not trace:
            continue

        for i, frame in enumerate(trace):
            parent = node
            if OPTIONS.lines:
                key = (frame.method, frame.loc)
            else:
                key = frame.method
            node = node.children.setdefault(key, Node(parent=parent, trace=trace[:i+1]))
            node.count += count
            node.traces += 1

    return root


def write_profile(root):
    write = sys.stdout.write
    total = root.count
    pct_mult = 100.0 / total
    stack = [root]

    def output(node, indent):
        parent = stack[-2] if len(stack) > 1 else None
        parent_method = parent.method if parent else None

        n2 = node
        p2 = parent

        if n2.method:
            count = node.count
            trace = [shorten_node(n2, p2)]

            while len(n2.children) == 1:
                p2 = n2
                n2 = n2.children.values()[0]
                trace.append(shorten_node(n2, p2))

            write('%s+%4.1f%% %s\n' % ('.  ' * indent, count * pct_mult, ' .. '.join(trace)))
            indent += 1

        children = n2.children.values()
        children.sort(key=lambda child: -child.count)
        for child in children:
            stack.append(child)
            output(child, indent)
            stack.pop()

    output(root, 0)

def shorten_node(node, parent):
    text = shorten(node.method, parent and parent.method)
    if OPTIONS.lines:
        text = '%s:%s' % (text, node.frame.loc.partition(':')[2])
    return text

def shorten(curr, prev=None):
    if OPTIONS.verbose:
        return curr
    curr = curr.split('.')
    prev = (prev or '').split('.')
    if curr[:-1] == prev[:-1]:
        return curr[-1]
    elif not OPTIONS.packages or curr[:-2] == prev[:-2]:
        return '.'.join(curr[-2:])
    else:
        return '.'.join([c[0] for c in curr[:-2]] + curr[-2:])



def init_gtk():
    import pygtk
    pygtk.require('2.0')
    import gtk


#     class TreeModel(object):
#         def __init__(self, root):
#             self.root = root

#             def label(node):
#                 node._child_list = node.children.values()
#                 for i, child in node._child_list:
#                     child._index = i
#                     label(child)
#             root._index = 0
#             label(root)

#         def on_get_flags(self):
#             return gtk.TREE_MODEL_ITERS_PERSIST

#         def on_get_n_columns(self):
#             return 1

#         def on_get_column_type(self, index):
#             return str

#         def on_get_iter(self, path):
#             node = self.root
#             for p in path:
#                 node = node._child_list[p]
#             return node

#         def on_get_path(self, rowref):
#             node = rowref
#             path = []
#             while node.parent is not None:
#                 path.append(node._index)
#                 node = node.parent
#             path.reverse()
#             return path

#         def on_get_value(self, rowref, column):
#             return rowref.method

#         def on_iter_next(self, rowref):
#             parent = rowref.parent


#         def on_iter_children(self, parent)
#         def on_iter_has_child(self, rowref)
#         def on_iter_n_children(self, rowref)
#         def on_iter_nth_child(self, parent, n)
#         def on_iter_parent(self, child)


    class ProfileView(object):
        def __init__(self, root):
            self.root = root
            self.window = window = gtk.Window(gtk.WINDOW_TOPLEVEL)
            window.set_title('Profile')
            window.set_size_request(600, 600)
            window.connect('delete_event', self._delete_event)

            import math
            total = root.count
            inv_total = 1.0 / total
            pct_mult = 100.0 * inv_total
            log_total = math.log(total)

            model_col_types = (object, str, int, str, float)
            (NODE, METHOD, CALLS, PERCENT_STR, PERCENT_FLOAT) = range(len(model_col_types))

            def make_row(node):
                percent_str = '%5.1f%%' % (node.count * pct_mult)
                percent_float = 100.0 * (node.count * inv_total) ** 0.5
                method = shorten_node(node)
                return (node, method, node.count, percent_str, percent_float)

            self.treestore = treestore = gtk.TreeStore(*model_col_types)

            def add(parent, node):
                if node.method:
                    iter = treestore.append(parent, row=make_row(node))
                else:
                    iter = None

                children = node.children.values()
                children.sort(key=lambda child: -child.count)
                for child in children:
                    add(iter, child)

            add(None, root)

            self.treeview = treeview = gtk.TreeView(treestore)

            font = 'normal 8'
            height = 6

            if True:
                col = gtk.TreeViewColumn('Method')
                col.set_expand(True)
                col.set_sort_column_id(CALLS)
                #col.set_sizing(gtk.TREE_VIEW_COLUMN_FIXED)

                # progress bar
                cell = gtk.CellRendererProgress()
                cell.set_property('text', '')
                cell.set_property('width', 50)
                cell.set_property('height', height)
                cell.set_property('ypad', 2)
                #cell.set_fixed_size(-1, 4)
                col.pack_start(cell, False)
                col.add_attribute(cell, 'value', PERCENT_FLOAT)
                #col.add_attribute(cell, 'xpad', PERCENT_FLOAT)
                #col.add_attribute(cell, 'width', PERCENT_FLOAT)

                # percentage text
                cell = gtk.CellRendererText()
                cell.set_property('xalign', 1.0)
                cell.set_property('font', font)
                cell.set_property('height', height)
                col.pack_start(cell, False)
                col.add_attribute(cell, 'text', PERCENT_STR)

                # method name
                cell = gtk.CellRendererText()
                cell.set_property('font', font)
                cell.set_property('height', height)
                #cell.set_fixed_size(-1, height)
                col.pack_start(cell, True)
                col.add_attribute(cell, 'text', METHOD)

                treeview.append_column(col)

            #treeview.set_fixed_height_mode(True)
            treeview.set_search_column(METHOD)
            treeview.connect('row-expanded', self._row_expanded)
            #treeview.expand_all()

            scroller = gtk.ScrolledWindow()
            scroller.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
            scroller.add(treeview)

            window.add(scroller)
            window.show_all()

        def _row_expanded(self, treeview, iter, path):
            num_children = self.treestore.iter_n_children(iter)
            #print 'row_expanded', treeview, iter, path, num_children
            if num_children == 1:
                #child_iter = self.treestore.iter_children(iter)
                #child_path = self.treestore.get_path(child_iter)
                child_path = path + (0,)
                self.treeview.expand_row(child_path, False)

        def _delete_event(self, widget, event, data=None):
            gtk.main_quit()
            return False

    global write_profile
    def write_profile(root):
        view = ProfileView(root)
        gtk.main()




import contextlib

@contextlib.contextmanager
def pager():
    proc = None
    try:
        out = sys.stdout
        if out.isatty():
            import subprocess
            pager = os.environ.get('PAGER') or 'less'
            proc = subprocess.Popen(pager, shell=True, stdin=subprocess.PIPE)
            sys.stdout = proc.stdin

        yield

    finally:
        if proc:
            sys.stdout = out
            proc.stdin.close()
            proc.wait()



def main():
    parser = optparse.OptionParser()
    parser.add_option('-t', '--tree', action='store_true', dest='tree')
    parser.add_option('--hot', action='store_false', dest='tree')
    parser.add_option('-p', '--packages', action='store_true')
    parser.add_option('-v', '--verbose', action='store_true')
    parser.add_option('-g', '--gui', action='store_true')
    parser.add_option('-l', '--lines', action='store_true')
    parser.add_option('-m', '--method')

    global OPTIONS
    (OPTIONS, args) = parser.parse_args()

    if OPTIONS.gui:
        init_gtk()

    if args:
        input = open(args[0])
    elif not sys.stdin.isatty():
        input = sys.stdin
    else:
        input = open('java.hprof.txt')

    parse(input)


if __name__ == '__main__':
    #with pager():
    main()