#!/bin/bash
#
# Copyright (c) 2016 by Cisco Systems, Inc.
# All rights reserved.

VMIGNORE="VmLck|VmPin|VmSwap|VmPTE|VmPeak|VmSize|VmStk|VmHWM"
WATCH_PROCS=
headers=1
count=10
attribute=rss

function ltracesize
{
    local pid=$1
    local total=0
    local i

    if [[ ! -d  /proc/$pid ]]; then
        echo 0
    fi
    if [[ ! -f  /proc/$pid/smaps ]]; then
        echo 0
    fi

    local ltvals=$(grep -C2 "shm.*ltrace" /proc/$pid/smaps  2> /dev/null | grep Rss | awk '{print $2}')
    for i in $ltvals; do
        let "total=$total+$i"
    done
    echo $total
}

function psssize
{
    local pid=$1
    local total=0
    local i

    if [[ ! -d  /proc/$pid ]]; then
        echo 0
    fi
    if [[ ! -f  /proc/$pid/smaps ]]; then
        echo 0
    fi

    local pssvals=$(grep "Pss:" /proc/$pid/smaps  2> /dev/null | awk '{print $2}')
    for i in $pssvals; do
        let "total=$total+$i"
    done
    echo $total
}

function heapsize
{
    local pid=$1
    local total=0

    if [[ ! -d  /proc/$pid ]]; then
        echo 0
    fi
    if [[ ! -f  /proc/$pid/smaps ]]; then
        echo 0
    fi

    local heapsz=$(grep -C1 "\[heap\]" /proc/$pid/smaps 2> /dev/null | grep Size | awk '{print $2}')
    for i in $heapsz; do
        let "total=$total+$i"
    done
    echo $total
}

ltracetotal=0
psstotal=0
heaptotal=0

function vmstat
{
    local pid=$1
    local name=$(cat /proc/$pid/cmdline)

    if [[ -z $name ]]; then
        name=$(printf "pid:%d" $pid)
    fi
    if [[ "$headers" == "1" ]]; then
        printf "%-20.20s%6s %9s%9s%9s" "Process" "PID" "Heap:" "Ltrace:" "Pss:"
        grep ^Vm /proc/$pid/status| egrep -v $VMIGNORE | awk '{printf "%8s", $1}' | paste - - - - - - - - -  - - -
    fi
    headers=0

    local heapsz=$(heapsize $pid)
    if [[ $heapsz == "0" ]]; then
        return;
    fi

    local ltacesz=$(ltracesize $pid)
    local psssz=$(psssize $pid)
    let "ltracetotal = $ltracetotal + $ltacesz"
    let "psstotal = $psstotal + $psssz"
    let "heaptotal = $heaptotal + $heapsz"
    printf "%-20.20s %6d %8d %8d %8d" "$name" $pid $heapsz $ltacesz $psssz
    grep ^Vm  /proc/$pid/status | egrep -v $VMIGNORE | awk '{printf "%8d", $2}' | paste - - - - - - - - - - - -
}

function topproc
{
    local count=$1
    local vmregion=$2
    local pid
    local pids
    local proc
    local pidtemp
    local attribsz

    case $vmregion in
    VmRSS|VmData|VmStk|VmExe|VmLib)
        list=$(cd /proc; find [1-9]* -ignore_readdir_race -name status 2> /dev/null | \
           xargs grep ${vmregion}: 2> /dev/null | grep -v task | \
           awk '{print $1, $2 }' | sort -n -k 2 | \
           sed -e "s;/status:${vmregion}:;;" | tail -${count} | \
           awk '{print $1}')
        ;;
    Heap|Pss|Ltrace)
        list=$(cd /proc; find [1-9]* -maxdepth 0 2> /dev/null)
        for pid in $list; do
            if [[ $vmregion == "Heap" ]]; then
                attribsz=$(heapsize $pid)
            elif [[ $vmregion == "Ltrace" ]]; then
                attribsz=$(ltracesize $pid)
            elif [[ $vmregion == "Pss" ]]; then
                attribsz=$(psssize $pid)
            else
                return
            fi
            echo $pid $attribsz >> /tmp/attrlist.$$
        done
        list=$(cat /tmp/attrlist.$$ | sort -n -k 2 | awk '{print $1}' | tail -${count})
        rm -f /tmp/attrlist.$$
        ;;
    esac


    for pid in $list; do
        vmstat $pid $vmregion
    done
    printf "%-20.20s %6s %8d %8d %8d\n" "Total" ""  $heaptotal $ltracetotal $psstotal

    if [[ ! -z $WATCH_PROCS ]]; then
        echo ---- Watch Procs ----
        local psfile=/tmp/topproc.$$
        ps -e > $psfile
        for proc in $WATCH_PROCS; do
            pidtemp=$(grep $proc $psfile | awk '{print $1}')
            pids="$pids $pidtemp"
        done
        pidtemp=$pids
        pids=$(echo $pidtemp | tr " " "\n" | sort | uniq)
        for pid in $pids; do
            vmstat $pid $vmregion
        done
        rm -f $psfile
    fi
}

function showhelp
{
    echo Display processes sorted by various attributes - see -a - default $attribute
    echo $0 [-c count] [-a attribute] [-w proc_name]* [-d]
    echo Where:
    echo \-a attribute - attribute is one of: rss data exe lib pss heap ltrace
    echo \-c count     - number of rows to display \(default $count\)
    echo \-w proc_name - process to watch. *multiple processes allowed, separated by -w
    echo \-d           - ltrace directory usage
}

while [ $# -gt 0 ]; do
    case "$1" in
    -c)
        shift
        count=$1
        shift
        ;;
    -a)
        shift
        attribute=$1
        shift
        ;;
    -w)
        shift
        WATCH_PROCS="$WATCH_PROCS $1"
        shift
        ;;
    -d)
        attribute="du"
        shift
        ;;
    -h|help)
        showhelp
        exit 0
        ;;
    *)
        showhelp
        exit 1
        ;;
    esac
done

case $attribute in
*rss)
    region=VmRSS
    ;;
*data)
    region=VmData
    ;;
*exe)
    region=VmExe
    ;;
*lib)
    region=VmLib
    ;;
*pss)
    region=Pss
    ;;
*heap)
    region=Heap
    ;;
*ltrace)
    region=Ltrace
    ;;
*du)
    du -k /dev/shmem/ltrace | sort -n | tail -$count
    exit $?
    ;;
*)
    showhelp
    exit 1
    ;;
esac
topproc $count $region
