#! /usr/bin/env perl
##**************************************************************
##
## Copyright (C) 1990-2013, Condor Team, Computer Sciences Department,
## University of Wisconsin-Madison, WI.
## 
## Licensed under the Apache License, Version 2.0 (the "License"); you
## may not use this file except in compliance with the License.  You may
## obtain a copy of the License at
## 
##    http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.
##
##**************************************************************

use strict;
use CondorTest;
use Check::SimpleJob;

my $smap_details = 1; # set this to non-zero for more detailed output
my $smap_dump = 0;    # set this to non-zero to dump entire smap

my $testname = CondorTest::GetDefaultTestName();
my $temptest = CondorTest::TempFileName("$testname");

#testreq: personal
my $config = <<CONDOR_TESTREQ_CONFIG;
  DAEMON_LIST = MASTER,SCHEDD,COLLECTOR,NEGOTIATOR,STARTD
  SHADOW = $(SBIN)/condor_shadow_s
  ChirpPath = "$(LIBEXEC)/condor_chirp"
  STARTD_ATTRS = $(STARTD_ATTRS) ChirpPath
CONDOR_TESTREQ_CONFIG
#endtestreq

# run a job and look at shadow memory usage
my $append_submit_commands = '
want_io_proxy = true
executable = x_shadow_mem_usage.pl
output = ' . $testname . '.out
error  = ' . $testname . '.err
';

# expected memory for shadows by platform.
# these are lowercase so that we can do case-insensitive lookup
#
# for X86_64 shadows
my %memory_data64 = (
	"debian6"  => "652",
	"debian7"  => "688",
	"debian8"  => "692",
	"fedora20" => "668",
	"fedora21" => "668",
	"redhat5"  => "444",
	"redhat6"  => "492",
	"redhat7"  => "648",
	"sl6"      => "496",
	"sl7"      => "656",
	"ubuntu12" => "684",
	"ubuntu14" => "664",
	"ubuntu16" => "664",
);
# for X86 shadows
my %memory_data32 = (
	"redhat5" => "348",
	"redhat6" => "388",
	"sl6"     => "348",
);


my $opsys;
my $arch;
if (CondorUtils::is_windows()) {
   $arch = "x86"; # assume arch on windows.
   $opsys = "Windows";
} else {
   my @plat = `condor_config_val arch opsysandver`;
   chomp @plat;
   $arch = $plat[0];	 # should really be checking the bitness of the shadow here
   $opsys = $plat[1];
}

#my $nmi_platform = `condor_config_val CONDOR_BUILD_PLATFORM`; CondorUtils::fullchomp $nmi_platform;
#print "Detected platform from ccv as $arch, $opsys, $nmi_platform. NMI_PLATFORM=$ENV{NMI_PLATFORM}.\n";

my $random_addr = 0;
if (-f "/proc/sys/kernel/randomize_va_space") { $random_addr = `cat /proc/sys/kernel/randomize_va_space`; chomp($random_addr); }
print "Address randomization is $random_addr\n";

my $tolerance = 1.05;
my $expectedmemory = 0;
if ($arch =~ /64/) { 
   $expectedmemory = $memory_data64{lc($opsys)};
} else {
   $expectedmemory = $memory_data32{lc($opsys)};
}
my $allowedmemory = ComputeMemLimit($tolerance, $expectedmemory);
print "Platform is $arch $opsys, expecting shadow memory = $expectedmemory, will error if > $allowedmemory.\n";

# run a job and look at shadow memory usage
my $on_ulog = sub {
  my %info = @_; # contains cluster,job,date,time,ulog
  CondorTest::debug("Job $info{cluster} ulog: $info{ulog}.\n",1);
  my $text = $info{ulog};
  if ($text =~ /shadow_pid\s*=\s*(\d+)/) {
     my $pid = $1;
     CondorTest::debug("Got shadow PID: $pid.\n",1);
     CondorTest::runcmd("cat /proc/sys/kernel/randomize_va_space");
     if (-f "/proc/self/personality") { CondorTest::runcmd("cat /proc/self/personality"); }
     if ($smap_dump) { CondorTest::runcmd("cat /proc/$pid/smaps"); }
     smap_sum_verbose($pid);
  }
};

SimpleJob::RunCheck(
  append_submit_commands=>$append_submit_commands
  , on_ulog=>$on_ulog
);

CondorTest::EndTest();

# scan /proc/<PID>/smaps and add up dirty private pages.
#
sub smap_sum_verbose {

   my $pid = shift;

   my $lib = ""; my $perm; my $inode; my $exe;
   #my $fd; my $addr1; my $addr2;
   my %tot;
   $tot{Heap} = 0;
   $tot{Stack} = 0;
   $tot{LibData} = 0;
   my $crypto_data = 0;

   if ($smap_details) { printf "Private Prot Module\n"; }

   open (SMAP, "</proc/$pid/smaps");
   while (<SMAP>) {
      if ($_ =~ /^([0-9A-Fa-f]+)-([0-9A-Fa-f]+)\s+(r[w-][x-]p)\s+(\S+\s+\S+)\s+(.*)$/) {

         $perm = $3; # perm
         $inode = $5; $lib=""; if ($inode =~ /^(\S+)\s+(.*)$/) { $inode = $1; $lib = $2; }
         #$addr1 = $1; $addr2 = $2; $fd = $4;

         if ( ! ($lib =~ /\.so/) && $lib =~ /condor_/) { $exe = $lib; }
         #
      } elsif ($_ =~ /^(\w+):\s+(\d+)/) {

         my $n = $2;
         $tot{$1} += $n;
         if ($1 =~ /Private_Dirty/) { 
            if ($perm eq "r--p") { $tot{LibRO} += $n; }
            if ($lib eq "") { 
               $tot{Anon} += $n;
            } else {
               if ($lib eq "[heap]") { $tot{Heap} += $n; }
               elsif ($lib eq "[stack]") { $tot{Stack} += $n; }
               elsif ($perm eq "rw-p") { $tot{LibData} += $n; if ($lib =~ /libcrypto.so/) { $crypto_data = $n; }}
               
               if ($lib =~ /condor_util/) { $tot{Utils} += $n; }
               elsif ($lib =~ /condor_/)  { $tot{Exec} += $n; }
               elsif ($lib =~ /condor/)   { 
                  if ($lib =~ /globus|voms/) { $tot{Globus} += $n; }
                  elsif ($lib =~ /classad/)  { $tot{Classad} += $n; }
                  else                       { $tot{Other} += $n; }
               } elsif ($lib =~ /lib/) { $tot{Other} += $n; }
            }
            if ($smap_details && $n) {
               my $short_lib = $lib;
               $short_lib =~ s|^.+/userdir/|/.../userdir/|;
               printf "%7d %s %s\n", $n, $perm, $short_lib;
            }
         }
      }
   } #while

   my $lbl = $pid;
   if ($exe) { 
      $exe =~ s/.*\///g;
      $lbl = "$pid:$exe";
   }

   printf "\n";
   print 'PID:Process            Total Dynamic =(Heap+Stack+Anon) LibRO LibData =(Exec+Utils+Clasad+Globus+System) Clean' . "\n";
   printf "%-22.22s %5d %7d %6d %5d %5d %5d %7d %6d %5d %6d %6d %6d %6d\n", $lbl, 
          $tot{Private_Dirty}, $tot{Heap}+$tot{Stack}+$tot{Anon}, $tot{Heap}, $tot{Stack}, $tot{Anon},
          $tot{LibRO}, $tot{LibData}, $tot{Exec}, $tot{Utils}, $tot{Classad}, $tot{Globus}, $tot{System}+$tot{Other},
          $tot{Private_Clean};

   if ($allowedmemory) {
      my $usedmem = $tot{Private_Dirty};
      my $romem = $tot{LibRO};
      if ($crypto_data > 12) {
          $allowedmemory += $crypto_data;
          print "libcrypto.so using $crypto_data of data, increasing allowed memory to $allowedmemory\n";
      }
      my $overmem = $usedmem - $allowedmemory;
      if ($random_addr && $romem > 120) {
          $allowedmemory += $romem - 100;
      }
      if($usedmem > $allowedmemory) {
         die "Shadow private memory use is $usedmem which is $overmem more than allowed value of $allowedmemory for $arch $opsys\n";
      } elsif ($overmem > 0) {
         print "Shadow private memory use is $usedmem, which is $overmem more than expected, but address randomization is using $romem, so the test passes with a warning\n";
      }
   }
}

sub ComputeMemLimit
{
	my $changerate = shift;
	my $memstat = shift;
	my $newsize = ($changerate * $memstat);
	print "$changerate * $memstat = $newsize\n";
	return($newsize);
}
