nohang/oom-sort

#!/usr/bin/env python3
"""
sort processes by oom_score
"""


from operator import itemgetter
from os import listdir
from argparse import ArgumentParser

##########################################################################


# define funtcions


def pid_to_oom_score(pid):
    with open('/proc/' + pid + '/oom_score') as f:
        return f.readline()[:-1]


def pid_to_oom_score_adj(pid):
    with open('/proc/' + pid + '/oom_score_adj') as f:
        return f.readline()[:-1]


def pid_to_cmdline(pid):
    """
    Get process cmdline by pid.

    pid: str pid of required process
    returns string cmdline
    """
    with open('/proc/' + pid + '/cmdline') as f:
        return f.read().replace('\x00', ' ').rstrip()


def pid_to_status_units(pid):

    with open('/proc/' + pid + '/status', 'rb') as f:
        f_list = f.read().decode('utf-8', 'ignore').split('\n')

        for i in range(len(f_list)):

            if i is 1:
                name = f_list[0].split('\t')[1]

            if i is uid_index:
                uid = f_list[i].split('\t')[2]

            if i is vm_rss_index:
                vm_rss = f_list[i].split('\t')[1][:-3]

            if i is vm_swap_index:
                vm_swap = f_list[i].split('\t')[1][:-3]

    return name, uid, vm_rss, vm_swap


def get_max_pid_len():
    with open('/proc/sys/kernel/pid_max') as file:
        for line in file:
            return len(line.strip())


##########################################################################


sort_dict = {
    'PID': 0,
    'oom_score': 1,
    'oom_score_adj': 2,
    'cmdline': 3,
    'Name': 4,
    'UID': 5,
    'VmRSS': 6,
    'VmSwap': 7
}


##########################################################################


# parse CLI args

# @TODO: user input validation

parser = ArgumentParser()

parser.add_argument(
    '--num',
    '-n',
    help="""max number of lines; default: 99999""",
    default=99999,
    type=str
)

parser.add_argument(
    '--len',
    '-l',
    help="""max cmdline length; default: 99999""",
    default=99999,
    type=int
)


parser.add_argument(
    '--sort',
    '-s',
    help="""sort by unit; default: oom_score""",
    default='oom_score',
    type=str
)


args = parser.parse_args()

display_cmdline = args.len

num_lines = args.num

sort_by = args.sort


if sort_by not in sort_dict:
    print('Invalid -s/--sort value. Valid values are:\nPID\noom_scor'
          'e [default value]\noom-sore_adj\nUID\nName\ncmdline\nVmR'
          'SS\nVmSwap')
    exit()


##########################################################################

# find VmRSS, VmSwap and UID positions in /proc/*/status for further
# searching positions of UID, VmRSS and VmSwap in each process

with open('/proc/self/status') as file:
    status_list = file.readlines()

status_names = []
for s in status_list:
    status_names.append(s.split(':')[0])

uid_index = status_names.index('Uid')
vm_rss_index = status_names.index('VmRSS')
vm_swap_index = status_names.index('VmSwap')


##########################################################################

# get sorted list with pid, oom_score, oom_score_adj, cmdline
# get status units: name, uid, rss, swap


oom_list = []

for pid in listdir('/proc'):

    # skip non-numeric entries and PID 1
    if pid.isdigit() is False or pid == '1':
        continue

    try:

        oom_score = pid_to_oom_score(pid)

        oom_score_adj = pid_to_oom_score_adj(pid)

        cmdline = pid_to_cmdline(pid)
        if cmdline == '':
            continue

        name, uid, vm_rss, vm_swap = pid_to_status_units(pid)

    except FileNotFoundError:
        continue

    except ProcessLookupError:
        continue

    oom_list.append((
        int(pid), int(oom_score), int(oom_score_adj), cmdline,
        name, int(uid), int(vm_rss), int(vm_swap)))

# list sorted by oom_score
oom_list_sorted = sorted(
    oom_list, key=itemgetter(int(sort_dict[sort_by])), reverse=True)


##########################################################################

# find width of columns


max_pid_len = get_max_pid_len()


max_uid_len = len(str(sorted(
    oom_list, key=itemgetter(5), reverse=True)[0][5]))


max_vm_rss_len = len(str(round(
    sorted(oom_list, key=itemgetter(6), reverse=True)[0][6] / 1024.0)))

if max_vm_rss_len < 5:
    max_vm_rss_len = 5


##########################################################################


# print output

if display_cmdline == 0:

    print(
        'oom_score oom_score_adj{}UID{}PID Name       {}VmRSS   VmSwap'.format(
            ' ' * (max_uid_len - 2),
            ' ' * (max_pid_len - 2),
            ' ' * max_vm_rss_len
        )
    )

    print(
        '--------- ------------- {} {} --------------- {}-- --------'.format(
            '-' * max_uid_len,
            '-' * max_pid_len,
            '-' * max_vm_rss_len
        )
    )

else:

    print(
        'oom_score oom_score_adj{}UID{}PID Name       {}VmRSS   VmSwa'
        'p   cmdline'.format(
            ' ' * (max_uid_len - 2),
            ' ' * (max_pid_len - 2),
            ' ' * max_vm_rss_len
        )
    )

    print(
        '--------- ------------- {} {} --------------- {}-- ------'
        '-- -------'.format(
            '-' * max_uid_len,
            '-' * max_pid_len,
            '-' * max_vm_rss_len
        )
    )

# print processes stats sorted by sort_dict[sort_by]

for i in oom_list_sorted[:int(num_lines)]:

    pid = i[0]
    oom_score = i[1]
    oom_score_adj = i[2]
    cmdline = i[3]
    name = i[4]
    uid = i[5]
    vm_rss = i[6]
    vm_swap = i[7]

    print(
        '{} {} {} {} {} {} M {} M {}'.format(
            str(oom_score).rjust(9),
            str(oom_score_adj).rjust(13),
            str(uid).rjust(max_uid_len),
            str(pid).rjust(max_pid_len),
            name.ljust(15),
            str(round(vm_rss / 1024.0)).rjust(max_vm_rss_len, ' '),
            str(round(vm_swap / 1024.0)).rjust(6, ' '),
            cmdline[:display_cmdline]
        )
    )