Parse goroutine dumps

goroutine_histogram.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import argparse
import operator
from collections import defaultdict
import re

def parseFile(f):
  with open(f, 'r') as opened:
    return opened.read().split("\n\n")

def parseStack(stack, includePaths=False):
  lines = stack.split("\n")
  # skip first line for now
  # sanitize any memory addresses
  return ";".join(re.sub(r'0x[\da-f]+', '0x0', line) for line in lines[1:])

def getStackCounts(input):
  counts = defaultdict(lambda: 0)
  stacks = parseFile(input)
  for stack in stacks:
    cleanedStack = parseStack(stack)
    counts[cleanedStack] += 1
  return counts

def comparotor(x, y):
  if x[1] == y[1]:
    if x[0] < y[0]:
      return -1
    elif x[0] == y[0]:
      return 0
    else:
      return 1
  elif x[1] < y[1]:
    return -1
  else:
    return 1

def writeStackCounts(counts, output):
  sorted_counts = sorted(counts.items(), cmp=comparotor, reverse=True)
  first = True
  for item in sorted_counts:
    stack = "\n".join(item[0].split(";"))
    count = item[1]
    if not first:
      output.write("\n")
    output.write("Count: %d\n" % count)
    output.write(stack)
    output.write("\n")
    first = False

def main():
    parser = argparse.ArgumentParser(
        description='Tool to create counts for goroutine dumps')

    parser.add_argument('-o', '--output', help='output target [default: stdout]')

    parser.add_argument('input', help='file containint goroutine stack dump')

    args = parser.parse_args()
    output = args.output

    input = args.input
    if not input:
        print >>sys.stderr, "No input specified"
        sys.exit(1)

    counts = getStackCounts(input)

    output = sys.stdout
    if args.output:
      with open(args.output, 'w') as output:
        writeStackCounts(counts, output)
    else:
        writeStackCounts(counts, output)

if __name__ == '__main__':
    main()