#!/usr/bin/perl -w
#
# vs-mem.pl -- create and view memory limits on vservers
# Copyright (C) 2010 KKoncepts Ltd. - <support@kkoncepts.net>
#
# vs-mem.pl 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.
#
# vs-mem.pl 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 can view the full GPL at <http://www.gnu.org/licenses/>.

use strict;
$|++;

&usage unless $ARGV[0];

my $binary = '/usr/sbin/vserver';

my $vs = $ARGV[0];
&usage if $vs =~ /^-/;
my $confdir = "/etc/vservers/$vs";

die qq("$vs" is not a vserver on this system.\n) unless (-d $confdir);
my $stopped = system("$binary $vs running");

my $rssdir = "$confdir/rlimits";
my $rssfile = "$rssdir/rss";

if (defined $ARGV[1]) {
  my $rss = $ARGV[1];
  die qq("$rss" is not a valid amount of RAM to assign.  Please specify an integer for RAM limit.\n) unless ($rss =~ /^\d+$/);
  mkdir "$rssdir" unless (-d "$rssdir"); 
  my $pages = $rss * 256 ;  #convert to 4k pages 
  if ($rss) {
    if (open RSS, ">$rssfile") {
      print RSS "$pages\n";
      close RSS;
      unless ($stopped) {
        system ("vlimit --xid $vs --rss $pages");
        system ("vlimit --xid $vs -S --rss $pages");
      }
    } else {
      print qq(Cannot write to the "$rssfile" file to set the memory limits.  Please check the permissions and try again.\n);
    } 
  } else {
    system qq(rm -rf "$rssdir") if (-d "$rssdir");
    unless ($stopped) {
      system ("vlimit --xid $vs --rss inf");
      system ("vlimit --xid $vs -S --rss inf");
    }
  }
}  

unless ($stopped) {
    my @memstats = grep /^Mem:/, qx($binary $vs exec free -m -o);
    my (undef,$rl,$ru,$rf,undef) = split /\s+/, $memstats[0];
    $rl = 0 unless -f "$rssfile";
    if ($rl) { print qq("$vs" memory limit:\t$rl Megabytes.\n); } 
    else { print qq("$vs" memory limit:\tNOT ASSIGNED.\n); }
    print qq("$vs" memory  used:\t$ru Megabytes.\n);
    print qq("$vs" memory  free:\t$rf Megabytes.\n) if ($rl);
} else {
    if ((-f $rssfile) and (open RSS, "$rssfile")) {
      my $limit = <RSS>;
      $limit /= 256;
      print qq("$vs" memory limit:\t$limit Megabytes.\n);
    } else {
      print qq("$vs" memory limit:\tNOT ASSIGNED.\n); 
    }
    print qq("$vs" memory  used:\t0 Megabytes (not running).\n);
}

exit 0;

sub usage {
  print qq(USAGE:\n);
  print qq("vs-mem.pl NAME" to see the memory statics for vserver "NAME".\n);
  print qq("vs-mem.pl NAME LIMIT" to set a memory limit of "LIMIT" Megabytes on vserver "NAME".\n);
  print qq("vs-mem.pl NAME 0" to remove any memory limits from vserver "NAME".\n);
  exit 0;
}