Flamegraph django middleware

flames.py

"""
Django middleware for generating request flame graphs.
Requires the flamegraph.pl perl script:
https://github.com/brendangregg/FlameGraph/blob/master/flamegraph.pl

Installation:
1. Create a directory for flame graphs
2. Copy the flamegraph.pl script to it
3. Add the FLAMES_DIR django setting
4. Add the flames.FlamesMiddleware to MIDDLEWARE_CLASSES

Usage:
To generate a flame graph just append ?flames to the requested url.
Middleware will create an svg in the FLAMES_DIR with the current timestamp.

Uncomment line 88 to automatically open the svg in a new google chrome tab.
"""

import os
import subprocess
import sys
import threading
import time
import traceback
from datetime import datetime
from xml.dom.minidom import Text

from django.conf import settings

FLAMES_DIR = os.path.abspath(settings.FLAMES_DIR)
FLAMEGRAPH_SCRIPT_PATH = os.path.join(FLAMES_DIR, "flamegraph.pl")


def get_module_name(module_path):
    for path in sys.path:
        path = path or os.getcwd()

        if module_path.startswith(path):
            rel_path = module_path[len(path) + 1 :]

            return (
                rel_path.replace(u"/__init__.py", u"")
                .replace(u"/", ".")
                .replace(u".py", u"")
                .strip()
            )

    return module_path


def write_samples(file, samples):
    for _, frame in samples:
        stack = traceback.extract_stack(frame)
        frame_strings = (
            func_name + u"@" + get_module_name(module_path)
            for module_path, _, func_name, _ in stack
        )
        stack_string = u";".join(frame_strings)
        file.write(u"{} {}\n".format(stack_string, 1))


def write_flames(logger, title="flames"):
    base_path = os.path.join(FLAMES_DIR, datetime.now().strftime("%Y-%m-%dT%H:%M:%S"),)

    out_txt_path = base_path + ".txt"
    out_svg_path = base_path + ".svg"

    title_element = Text()
    title_element.data = title
    title_xml = title_element.toxml()

    with open(out_txt_path, "w") as out_txt:
        write_samples(out_txt, logger.samples)

    with open(out_svg_path, "w") as out_svg:
        subprocess.call(
            ["perl", FLAMEGRAPH_SCRIPT_PATH, out_txt_path, "--title", title_xml,],
            stdout=out_svg,
        )

    # subprocess.call(['google-chrome', 'file://' + out_svg_path])


class StackLogger(object):
    def __init__(self, thread_id, interval=0.001):
        super(StackLogger, self).__init__()

        self.thread_id = thread_id
        self.interval = interval
        self.should_stop = False
        self.samples = []

        self.started = threading.Event()
        self.thread = threading.Thread(target=self.run)

    def run(self):
        start_time = time.time()
        self.samples = []
        self.started.set()

        while not self.should_stop:
            frame_start = time.perf_counter()

            timestamp = time.time() - start_time
            frame = sys._current_frames()[self.thread_id]
            self.samples.append((timestamp, frame))

            frame_dt = time.perf_counter() - frame_start
            time.sleep(max(0, self.interval - frame_dt))

    def start(self):
        self.started.clear()
        self.should_stop = False
        self.thread.start()
        self.started.wait()

    def stop(self):
        self.should_stop = True
        self.thread.join()


class FlamesMiddleware(object):
    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        if settings.DEBUG and "flames" in request.GET:
            logger = StackLogger(threading.current_thread().ident)
            logger.start()
            request._stack_logger = logger

        response = self.get_response(request)

        if settings.DEBUG and hasattr(request, "_stack_logger"):
            logger = request._stack_logger
            logger.stop()

            url = request.build_absolute_uri()
            write_flames(logger, title=url)

        return response