python-tornado-tracemalloc.py

import logging
import tracemalloc

from tornado.web import RequestHandler, HTTPError


class MemTraceHandler(RequestHandler):
    def post(self) -> None:
        action = self.get_argument('action')
        nframe = self.get_argument('nframe', "6")
        try:
            nframe = int(nframe)
        except ValueError:
            raise HTTPError(400, "invalid value for nframe %r" % nframe)
        if nframe < 1 or nframe > 40:
            raise HTTPError(400, "nframe=%d out of range" % nframe)
        if action not in ("start", "stop"):
            raise HTTPError(400, "invalid value for action %r" % action)

        is_tracing = tracemalloc.is_tracing()
        if is_tracing != (action == "stop"):
            raise HTTPError(400, "action=%s but is_tracing=%r" % (action, is_tracing))

        if action == "start":
            tracemalloc.start(nframe)
        else:
            tracemalloc.stop()

    def get(self) -> None:
        top = self.get_argument('top', "8")
        try:
            top = int(top)
        except ValueError:
            raise HTTPError(400, "invalid value for top")
        if top < 1 or top > 40:
            raise HTTPError(400, "top out of range")
        if not tracemalloc.is_tracing():
            raise HTTPError(400, "not tracing")

        current_mem, peak_mem = tracemalloc.get_traced_memory()
        overhead = tracemalloc.get_tracemalloc_memory()
        stats = tracemalloc.take_snapshot().statistics('traceback')[:top]

        summary = "traced memory: %d KiB  peak: %d KiB  overhead: %d KiB" % (
            current_mem // 1024, peak_mem // 1024, overhead // 1024
        )
        logging.info("%s", summary)

        out_lines = [summary]
        for trace in stats:
            out_lines.append("---")
            out_lines.append("%d KiB in %d blocks" % (trace.size // 1024, trace.count))
            out_lines.extend(trace.traceback.format())
        out_lines.append('')

        self.set_header('Content-Type', "text/plain")
        self.set_header('Connection', "close")
        self.write('\n'.join(out_lines).encode('utf-8'))