#!/usr/bin/perl

#
# memstats lists processes memory usage
# Copyright 2014  Grégory Soutadé
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

use POSIX ();
use Getopt::Long;

# Convert int to kilobytes, megabytes and gigabytes
sub printPacked {
    my ($size) = @_;
    my $suffix = "";

    if($size < 1024*1024)
    {
	$suffix = "k";
	$size = int($size/1024);
    }
    elsif($size < 1024*1024*1024)
    {
	$suffix = "m";
	$size = int($size/(1024*1024));
    }
    else
    {
	$suffix = "g";
	$size /= 1024*1024*1024;
	if ($size != int($size))
	{
	    $size = sprintf "%.2f", $size;
	}
    }

    return "$size$suffix";
}

# Convert seconds in hours:minutes:seconds
sub printTime {
    my ($time) = @_;
    my $res = "";

    $seconds = $time % 60;
    $minutes = 0;
    $hours = 0;
    if ($time < 60)
    {
	$seconds = $time;
	$minutes = 0;
	$hours = 0;
    }
    else
    {
	$seconds = $time % 60;
	$time -= $seconds;
	$time /= 60;
	if ($time < 60)
	{
	    $minutes = $time;
	    $hours = 0;
	}
	else
	{
	    $minutes = $time % 60;
	    $time -= $minutes;
	    $time /= 60;
	    $hours = $time;
	}
    }

    $res = sprintf "%02d:%02d:%02d", $hours, $minute, $seconds;

    return $res;
}

# PID OWNER USER VIRT RES S TIME COMMAND
my @sizes ;
my @lines ;

# Put line to column writer
sub putLine {
    my ($line) = @_;

    push @lines, $line;

    @s = split(/ /, $line);
    if (scalar(@sizes) == 0)
    {
	foreach $v (@s) {push @sizes, length($v);}
    }
    else
    {
	for($i=0; $i<scalar(@sizes); $i++)
	{
	    $l = length($s[$i]);
	    $sizes[$i] = $l if ($l > $sizes[$i]);
	}
    }
}

# Compute column format
sub computeFormat {
    my $format = "";

    for($i=0; $i<scalar(@sizes); $i++)
    {
	if ($i > 0)
	{
	    $format .= sprintf " %%%ds", $sizes[$i];
	}
	else
	{
	    $format .= sprintf "%%%ds", $sizes[$i];
	}
    }
    $format .= "\n";

    return $format;
}

# Get current start time in seconds
sub getStartTime {
    open(FIC,"</proc/self/stat") or die $!;
    my @infos = split(/ /, <FIC>);
    close (FIC);

    return $infos[21];
}

# Constants
my $clock_ticks = POSIX::sysconf(&POSIX::_SC_CLK_TCK );
my $page_size   = POSIX::sysconf(&POSIX::_SC_PAGESIZE);
my $starttime   = getStartTime();
my $NB_PROCESSES  = 30;
my $help;
my $all;
my $virt;

# Process paramters
my $program_name = $0;

GetOptions ("nb_processes=i" => \$NB_PROCESSES,
	    "all"          => \$all,
	    "virt"         => \$virt,
	    "help"         => \$help);

if ($help)
{
    printf "Display processes memory usage\nusage: %s [--all|--nb_processes=NB_PROCESSES] [--virt]\n", $program_name;
    print "\t--all:\t\tList all processes\n";
    print "\t--nb_processes:\tLimit list to NB_PROCESSES (default: 30)\n";
    print "\t--virt:\t\tSort by virtual memory (RES by default)\n";

    exit(0);
}

$NB_PROCESSES = -1 if ($all);

############ START ############

my @processes;

# Analyze all processes
opendir (DIR, "/proc") or die $!;

while (my $pid = readdir(DIR)) {
    next if ($pid !~ /[0-9]+/);

    # Get owner name
    my $uid = 0;
    open(FIC,"</proc/$pid/status") or next; # Processus may have stopped between readdir and analyze
    while( defined( $l = <FIC> ) ) {
	if ($l =~ /Uid:\s+([0-9]+)\s+.*/)
	{
	    $uid = $1;
	    last;
	}
    }
    close (FIC);
    
    ($owner_name,$passwd,$uid,$gid,
     $quota,$comment,$gcos,$dir,$shell,$expire) = getpwuid($uid);

    # Get infos
    open(FIC,"</proc/$pid/stat") or next; # Processus may have stopped between readdir and analyze
    my @infos = split(/ /, <FIC>);
    close (FIC);

    push @processes,
    	{"pid"       => $infos[0],
	 "owner"     => $owner_name,
    	 "comm"      => $infos[1],
    	 "state"     => $infos[2],
    	 "starttime" => $infos[21],
    	 "vsize"     => $infos[22],
    	 "rss"       => $infos[23],
    	};
}

# Sort in reverse order by memory usage
if ($virt)
{
    @p = sort {$b->{vsize} <=> $a->{vsize}} @processes;
}
else
{
    @p = sort {$b->{rss} <=> $a->{rss}} @processes;
}

# Compute lines
putLine("PID OWNER VIRT RES S TIME COMMAND");

foreach $e (@p)
{
    $vsize = printPacked($$e{vsize});
    $rss   = printPacked($$e{rss}*$page_size);
    $time  = int(($starttime-$$e{starttime}) / $clock_ticks);
    $time  = printTime $time;
    $name  = $$e{comm};
    $name  =~ s/\(//;
    $name  =~ s/\)//;

    last if ($NB_PROCESSES > 0 && scalar(@lines) > $NB_PROCESSES);

    putLine("$$e{pid} $$e{owner} $vsize $rss $$e{state} $time $name");
}

my $format = computeFormat;

# Start display
foreach $l (@lines)
{
    printf $format, split(/ /, $l);
}