#!/bin/bash

# This program will calculate memory usage for each process and generate a
# comma-separated value (CSV) output file named output.csv in the current
# directory.  The output will consist of 2 lines.  The first is a comma-
# separated list of process names.  The second is a list of comma-separated
# memory usage values (expressed in bytes). Here is an abbrieviated example of
# the output:
# python(9),/lib/systemd/systemd-journald,/usr/bin/python,/sbin/init,
# phosphor-hwmon-readd(4),ipmid,phosphor-ledcontroller(4)
# 57896960,11789312,4434944,2893824,1900544,1764352

program_name=${0##*/}
temp_file_path_1=/tmp/${program_name}_results_1
temp_file_path_2=/tmp/${program_name}_results_2
temp_file_path_3=/tmp/${program_name}_results_3

temp_file_list="${temp_file_path_1} ${temp_file_path_2} ${temp_file_path_3}"
csv_file_path="output.csv"

# Description of argument(s):
# pid  The pid for which you desire statistics. If this is not specified,
# statistics will be gathered for all active pids.

function get_parms {

  # Get program parms.

  pids="${1:-$(ls /proc | grep -v [A-Za-z])}" ; shift

  return 0

}


function exit_function {

  # Used to clean up tmp files.

  rm -f ${temp_file_list}
  return

}


function validate_parms {

  # Validate program parameters.

  # Making sure only root can run our script.
  if [ "${USER}" != "root" ] ; then
    echo "This script must be run as root" 1>&2
    return 1
  fi

  trap "exit_function $signal \$?" EXIT
  return 0

}


function get_process_mem {

  local pid="${1}" ; shift
  # Count memory statistic for passed pid.

  # Description of argument(s):
  # pid  The process ID for which you desire statistics.
  [ ! -f /proc/${pid}/status -o ! -f /proc/${pid}/smaps ] && return 0

  # pss_total      Total proportional set size of a process.
  # private_total  Total number of clean and dirty private pages in the
  #                mapping.
  # shared_total   The difference between pss_total and private_total.

  local pss_total private_total shared_total sum name
  pss_total=$(grep -e "^Pss:" /proc/${pid}/smaps | awk '{print $2}' | awk '{sum += $1} END {print sum}')
  private_total=$(grep -e "^Private" /proc/${pid}/smaps | awk '{print $2}' | awk '{sum += $1} END {print sum}')

  [ -z "${pss_total}" -o -z "${private_total}" ] && return 0
  (( shared_total=pss_total-private_total ))
  name=$(cut -d "" -f 1 /proc/${pid}/cmdline)
  (( sum=shared_total+private_total ))
  echo -e "${private_total} + ${shared_total} = ${sum} ${name}"

}


function mainf {

  get_parms "$@" || return 1

  validate_parms || return 1

  # To ensure temp files not exist.
  rm -f ${temp_file_list}

  for pid in ${pids} ; do
    get_process_mem ${pid} >> ${temp_file_path_1}
  done

  # This is used to sort results by memory usage.
  sort -gr -k 5 ${temp_file_path_1} > ${temp_file_path_2}

  # Find duplicates in the process list output and combine them.  In such
  # cases, adjust the process name by including a (<count>) suffix.  In the
  # following example of output, 4 instances of "phosphor-hwmon-readd" have
  # been combined.
  # 974848 + 925696 = 1900544       phosphor-hwmon-readd(4)
  for proc_name in $(awk '{print $6}' ${temp_file_path_2} | sort -u) ; do
    count=$(awk -v src=${proc_name} '{if ($6==src) {print $6}}' ${temp_file_path_2} | wc -l)
    [ "${count}" = "1" ] && count_string="" || count_string="(${count})"
    vmsize_in_kb=$(awk -v src=${proc_name} '{if ($6==src) {print $1}}' ${temp_file_path_2} | awk '{sum += $1} END {print sum}')
    vmrss_in_kb=$(awk -v src=${proc_name} '{if ($6==src) {print $3}}' ${temp_file_path_2} | awk '{sum += $1} END {print sum}')
    total=$(awk '{print $5}' ${temp_file_path_2} | awk '{sum += $1} END {print sum}')
    (( sum=vmrss_in_kb+vmsize_in_kb ))
    echo -e "${vmsize_in_kb}  + ${vmrss_in_kb} = ${sum} \t ${proc_name}${count_string}" >> ${temp_file_path_3}
  done

  # Sort once more.
  sort -gr -k 5 ${temp_file_path_3} > ${temp_file_path_1}

  # Read results from temp file and convert it to csv.
  csv_line1=""
  csv_line2=""
  while read line ; do
    while read private plus_operator shared equal_sign sum name ; do
      (( sum == 0 )) && continue
      csv_line1+=",${name}"
      csv_line2+=",${sum}"
    done<<<${line}
  done < ${temp_file_path_1}

  # Strip leading commas.
  csv_line1="${csv_line1#,}"
  csv_line2="${csv_line2#,}"
  { echo "${csv_line1}" ; echo "${csv_line2}" ; } >> ${csv_file_path}

  return 0

}

# Main

  mainf "${@}"
  rc="${?}"
  exit "${rc}"