#!/usr/bin/env python3
#
# this walks all the tasks on the system and prints out a stack trace
# of any tasks waiting in D state.  If you pass -a, it will print out
# the stack of every task it finds.
#
# It also makes a histogram of the common stacks so you can see where
# more of the tasks are.  Usually when we're deadlocked, we care about
# the least common stacks.
#
import sys
import os
import argparse

parser = argparse.ArgumentParser(description='Show kernel stacks')
parser.add_argument('-a', '--all_tasks', action='store_true', help='Dump all stacks')
parser.add_argument('-p', '--pid', type=str, help='Filter on pid')
parser.add_argument('-c', '--command', type=str, help='Filter on command name')
options = parser.parse_args()

stacks = {}

# parse the units from a number and normalize into KB
def parse_number(s):
    try:
        words = s.split()
        unit = words[-1].lower()
        number = int(words[1])
        tag = words[0].lower().rstrip(':')

        # we store in kb
        if unit == "mb":
            number = number * 1024
        elif unit == "gb":
            number = number * 1024 * 1024
        elif unit == "tb":
            number = number * 1024 * 1024

        return (tag, number)
    except:
        return (None, None)

# read /proc/pid/stack and add it to the hashes
def add_stack(path, pid, cmd, status):
    global stacks

    try:
        stack = open(os.path.join(path, "stack"), 'r').read()
    except:
        return

    if (status != "D" and not options.all_tasks):
        return

    print("%s %s %s" % (pid, cmd, status))
    print(stack)
    v = stacks.get(stack)
    if v:
        v += 1
    else:
        v = 1
    stacks[stack] = v


# worker to read all the files for one individual task
def run_one_task(path):

    try:
        stat = open(os.path.join(path, "stat"), 'r').read()
    except:
        return
    words = stat.split()
    pid, cmd, status = words[0:3]

    cmd = cmd.lstrip('(')
    cmd = cmd.rstrip(')')

    if options.command and options.command != cmd:
        return

    add_stack(path, pid, cmd, status)

def print_usage():
    sys.stderr.write("Usage: %s [-a]\n" % sys.argv[0])
    sys.exit(1)

# for a given pid in string form, read the files from proc
def run_pid(name):
    try:
        pid = int(name)
    except:
        return

    p = os.path.join("/proc", name, "task")
    if not os.path.exists(p):
        return

    try:
        for t in os.listdir(p):
            run_one_task(os.path.join(p, t))
    except:
        pass

if options.pid:
    run_pid(options.pid)
else:
    for name in os.listdir("/proc"):
        run_pid(name)

values = {}
for stack, count in stacks.items():
    l = values.setdefault(count, [])
    l.append(stack)

counts = list(values.keys())
counts.sort(reverse=True)
if counts:
    print("-----\nstack summary\n")

for x in counts:
    if x == 1:
        print("1 hit:")
    else:
        print("%d hits: " % x)

    l = values[x]
    for stack in l:
        print(stack)
        print("-----")