Support per-slice daily bandwidth caps

diff --git a/pl_mom.pl b/pl_mom.pl
index 4e611a4..8ed352f 100755 (executable)
--- a/pl_mom.pl
+++ b/pl_mom.pl
@@ -5,12 +5,10 @@ use Sys::Syslog;
 use Sys::Hostname;
 #use LWP::Simple;
 
-$debug = 0;
+$debug = 1;
 $proc = "pl_mom";
 $alias_addr = "pl-mom\@planet-lab.org";
 $from_addr = "support\@planet-lab.org";
-#$bwcap = "1.5Mbit";
-$bwcap = "off";
 
 if (! $debug) {
     $kill_thresh = 90;
@@ -18,10 +16,13 @@ if (! $debug) {
     $log_thresh = 85;
     $change_thresh = 5;
     $min_thresh = 10;
-    $byte_cutoff = 16000000000;    # 16GB
+    #$bwcap_default = "off";
+    $bwcap_default = "1.5Mbit";
+    $cutoff_default = "16200000000";  # 16GB, for 1.5Mbit cap
     $bwmon_sleep = 900;
 
     $sendmail = "/usr/sbin/sendmail -t -f$from_addr";
+    $vservers = "/etc/vservers";
     $pidfile = "/var/run/$proc.pid";
     $rebootfile = "/var/lib/misc/pl_mom.reboot";
     $daily_log = "/var/lib/misc/pl_mom.daily";
@@ -33,17 +34,19 @@ if (! $debug) {
     $reboot_thresh = 20;
     $log_thresh = 2;
     $change_thresh = 5;
-    $min_thresh = 10;
-    $byte_cutoff = 16000;
+    $min_thresh = 2;
+    $bwcap_default = "1Kbit";
+    $cutoff_default = "10800";
     $bwmon_sleep = 10;
 
     $sendmail = "cat";
+    $vservers = "./debug";
     $pidfile = "./$proc.pid";
-    $rebootfile = "./pl_mom.reboot";
-    $daily_log = "./pl_mom.daily";
-    $daily_stamp = "./pl_mom.stamp";
-    $configfile = "./pl_mom.conf";
-    $capfile = "./pl_mom.oldcaps";
+    $rebootfile = "./debug/pl_mom.reboot";
+    $daily_log = "./debug/pl_mom.daily";
+    $daily_stamp = "./debug/pl_mom.stamp";
+    $configfile = "./debug/pl_mom.conf";
+    $capfile = "./debug/pl_mom.oldcaps";
 }
 
 $sleep = 30;
@@ -73,8 +76,8 @@ if (! $pid) {
 $pid = fork();
 if (! $pid) {
     syslog ("info", "pl_mom: Launching bandwidth monitor");
-    if ($bwcap =~ /off/) {
-       syslog("info", "pl_mom: Bandwidth capping is off");
+    if ($bwcap_default =~ /off/) {
+       syslog("info", "pl_mom: Max rate unlimited by default");
     }
     reset_bandwidth_caps();
     bandwidth_monitor();
@@ -169,59 +172,37 @@ sub bandwidth_monitor {
            syslog("info", "pl_mom: Beginning bandwidth monitoring for $now");
        }
 
-        # Get baseline counts
-       `touch $daily_log`;
-       open (BASE, "+<$daily_log") ||
-           print "Cannot open $daily_log; $!\n";
-       while (<BASE>) {
-           my ($sliceid, $bytecount) = split(/ /);
-           $Start{$sliceid} = $bytecount;
-       }
-       
-       $status = `tc -s -d qdisc show`;
-       @lines = split(/\n/, $status);
-       for ($i = 0; $i < @lines; $i++) {
-           if ($lines[$i] =~ /qdisc pfifo/) {
-               $lines[$i] =~ s/^ +//;
-               @fields = split(/ /, $lines[$i]);
-               $slice = $fields[2];
-               $slice =~ s/://;
-
-               if ($slice != 9999) {
-                   $lines[$i+1] =~ s/^ +//;
-                   @fields = split(/ /, $lines[$i+1]);
-                   $bytes = $fields[1];
-                   #if ($bytes) {print "Slice $slice sent $bytes bytes\n";}
-                   
-                   if (! defined($Start{$slice})) {
-                       print BASE "$slice $bytes\n";
-                       $Start{$slice} = $bytes;
-                   }
-                   $Now{$slice} = $bytes;
+       get_slice_names();
+       get_baseline_counts();
+       get_slice_limits();
+
+       foreach $slice ( sort (keys %Start) ) {
+           if (defined $Now{$slice}) {
+               $today = $Now{$slice} - $Start{$slice};
+               if (! (defined ($Cutoff{$slice})||$bwcap_default =~ /off/)) {
+                   $Cutoff{$slice} = $cutoff_default;
+                   $Maxrate{$slice} = $bwcap_default;
                }
-           }
-       }
-       close (BASE);
-       
-       if (!($bwcap =~ /off/)) {
-           foreach $slice ( sort (keys %Start) ) {
-               if (defined $Now{$slice}) {
-                   $today = $Now{$slice} - $Start{$slice};
-                   if ($today >= $byte_cutoff && ! defined($Cap{$slice})) {
-                       $Cap{$slice} = "sent";
-                       $slicename = get_slice_name($slice);
-                       if ($slicename) {
-                           bw_cap_mail($slicename);
-                           log_bandwidth_cap($slicename, $bwcap);
-                           cap_bandwidth($slicename, $bwcap);
-                       } else {
-                           syslog("warning", "pl_mom: Could not find slice ".
-                                  "name for slice ID $slice");
-                       }
+               if ($debug) {
+                   if ($today) {
+                       $cutoff = defined($Cutoff{$slice}) 
+                           ? $Cutoff{$slice} : "<none>"; 
+                       print "Slice $slice sent $today bytes; ".
+                           "cutoff $cutoff\n";
                    }
-               } else {
-                   # Token bucket for this slice is gone!
                }
+               if (defined ($Cutoff{$slice}) && 
+                   $today >= $Cutoff{$slice} && 
+                   ! defined($Cap{$slice})) {
+                   $Cap{$slice} = "sent";
+                   bw_cap_mail($slice);
+                   if (! $debug) {
+                       log_bandwidth_cap($slice, $bwcap_default);
+                       cap_bandwidth($slice, $bwcap_default);
+                   } 
+               }
+           } else {
+               # Token bucket for this slice is gone!
            }
        }
 
@@ -245,6 +226,84 @@ sub read_config_file {
     }
 }
 
+sub get_slice_names {
+    # Read slice names from /etc/passwd
+    if (defined (%Name)) { undef %Name; }
+    open (PASSWD, "</etc/passwd") ||
+       print "Cannot open /etc/passwd; $!\n";
+    while (<PASSWD>) {
+       my ($slicename, $passwd, $sliceid) = split(/:/);
+       $Name{$sliceid} = $slicename;
+    }
+    close PASSWD;
+}
+
+sub get_baseline_counts {
+    `touch $daily_log`;
+    open (BASE, "+<$daily_log") ||
+       print "Cannot open $daily_log; $!\n";
+    while (<BASE>) {
+       my ($slice, $bytecount) = split(/ /);
+       $Start{$slice} = $bytecount;
+    }
+
+    my $status = `tc -s -d qdisc show`;
+    my $sliceid = "9999";
+    @Lines = split(/\n/, $status);
+    foreach $line ( @Lines ) {
+       if ($line =~ /qdisc pfifo (.*):/) {
+           $sliceid = $1;
+       } else {
+           if ($line =~ /Sent (.*) bytes/) {
+               my $bytes = $1;
+               if ($sliceid != 9999) {
+                   my $slice = $Name{$sliceid};
+                   if ($debug && $bytes) {
+                       print "Slice: $slice ($sliceid), bytes $bytes\n";
+                   }
+                   if (! defined($Start{$slice})) {
+                       print BASE "$slice $bytes\n";
+                       $Start{$slice} = $bytes;
+                   }
+                   $Now{$slice} = $bytes;
+               }
+           }
+       }
+    }
+    close (BASE);
+}
+
+sub get_slice_limits {
+    if (defined %Maxrate) { undef %Maxrate; }
+    if (defined %Cutoff)  { undef %Cutoff; }
+    if (-e $vservers) {
+       my $result = `grep -H "^BWMAXRATE" $vservers/*.conf`;
+       chomp ($result);
+       my @Lines = split(/\n/,$result);
+       foreach $line ( @Lines ) {
+           if ($line =~ /\/([^\/]*).conf:BWMAXRATE=(.*)[Mm]bit/) {
+               $slice = $1;
+               $limit = $2."Mbit";
+               $cutoff = ($2 * 1000000 * 86400)/8;
+           } else {
+               if ($line =~ /\/([^\/]*).conf:BWMAXRATE=(.*)[Kk]bit/) {
+                   $slice = $1;
+                   $limit = $2."Kbit";
+                   $cutoff = ($2 * 1000 * 86400)/8;
+               } else {
+                   die "Could not parse line $line";
+               }
+           }
+           $Maxrate{$slice} = $limit;
+           $Cutoff{$slice} = $cutoff;
+           if ($debug) {
+               print "Slice $slice, maxrate $Maxrate{$slice}, ".
+                   "cutoff $Cutoff{$slice}\n";
+           }
+       }
+    }
+}
+
 sub reset_bandwidth_caps {
     if (-e $capfile) {
        open(CAP, "<$capfile") or die "Cannot open $capfile: $!";
@@ -300,28 +359,12 @@ sub get_date {
     return $date;
 }
 
-sub get_slice_name {
-    # Arg 0: slice ID
-
-    # Need to map slice id to slice name; is there a sensor?
-    # For now, get it from /etc/passwd
-    my $name = "";
-    open (PASSWD, "</etc/passwd") ||
-       print "Cannot open /etc/passwd; $!\n";
-    while (<PASSWD>) {
-       my ($slicename, $passwd, $sliceid) = split(/:/);
-       if ($sliceid == $_[0]) {
-           $name = $slicename;
-       }
-    }
-    close PASSWD;
-    return $name;
-}
-
 sub bw_cap_mail {
+    my ($slicename) = @_;
     my $hostname = hostname();
     my $date = get_date();
-    my $sent = int($byte_cutoff/1000000000);
+    my $sent = int($Cutoff{$slicename}/1000000000);
+    my $bwcap = $Maxrate{$slicename};
 
     send_mail("$alias_addr, $slicename\@slices.planet-lab.org",
              "$proc capped bandwidth of slice $slicename on $hostname",