stacktools.py -> stacks.folded ; flamegraph.pl stacks.folded > graph.svg

stacktools.py

import sys
import os
import signal 
from traceback import extract_stack 
from collections import OrderedDict
"""
MODE_SYSTRACE, ncalls
MODE_TIMER, capture stack every SAMPLE_INTERVAL
trace pure python code only, don't use syscall
"""

datafile = None 
stacks = {}

SAMPLE_INTERVAL = 0.00001
MODE_TIMER = 0x1 << 1
MODE_SYSTRACE = 0x1 << 2

profile_mode = None

def sigtimer_handler(signum, frame):
    folded = ";".join([x[2] for x in extract_stack()[:-1]])
    if folded in stacks:
        stacks[folded] += 1
    else:
        stacks[folded] = 1 

def _trace(frame, event, arg): 
    if event in "return line c_return": 
        return None 
    folded = ";".join([x[2] for x in extract_stack()[:-1]])
    if folded in stacks:
        stacks[folded] += 1
    else:
        stacks[folded] = 1


def start_trace(path=None, mode=MODE_SYSTRACE): 
    global datafile, profile_mode
    if path:
        datafile = path
    else:        
        datafile = "stacks.folded" 
    profile_mode = mode 
    if mode & MODE_SYSTRACE:
        sys.settrace(_trace)
    elif mode & MODE_TIMER: 
        signal.siginterrupt(signal.SIGVTALRM, False)
        signal.signal(signal.SIGVTALRM, sigtimer_handler)
        signal.setitimer(signal.ITIMER_VIRTUAL, SAMPLE_INTERVAL, SAMPLE_INTERVAL) 

def stop_trace():
    if profile_mode & MODE_SYSTRACE:
        sys.settrace(None) 
    if profile_mode & MODE_TIMER:
        signal.signal(signal.SIGVTALRM, signal.SIG_IGN)
    f = open(datafile, "w")
    for k, v in OrderedDict(sorted(stacks.items(),
        key=lambda t: t[1])).items():
        f.write("\x20".join((k, str(v))) + "\n")
    f.close()