os_setup.in

#!/usr/bin/perl -wT

#
# EMULAB-COPYRIGHT
# Copyright (c) 2000-2003 University of Utah and the Flux Group.
# All rights reserved.
#

use English;
use Getopt::Std;
require 'ctime.pl';

#
# Reboot the nodes in an experiment. The nodes table will already contain
# all the information. This script deals with possible disk reloading,
# rebooting, and waiting for nodes to come back alive before allowing
# experiment creation to continue.
#
# TODO: Reload disk images.
# 
# usage: os_setup <pid> <eid>
#
# errorcode:  0 - all reboots succeeded.
#             1 - some/all reboots failed; retry may help.
#            -1 - failure; retry is inappropriate.
#

sub usage()
{
    print STDERR "Usage: os_setup <pid> <eid>\n";
    exit(-1);
}
my  $optlist = "d";

#
# Used to die with a -1 return code, to indicate to caller (tbswap)
# that the failure is not likely to be fixed with another attempt.
#
sub die_noretry($)
{
    my ($mesg) = shift;
    print STDERR "$mesg\n";
    exit(-1);
}

#
# Configure variables
#
my $TB		= "@prefix@";
my $DBNAME	= "@TBDBNAME@";
my $TBOPS       = "@TBOPSEMAIL@";
my $TESTMODE    = @TESTMODE@;
my $TFTP	= "/tftpboot";

#
# Testbed Support libraries
#
use lib "@prefix@/lib";
use libdb;
use libtestbed;

my $nodereboot	= "$TB/bin/node_reboot";
my $os_load	= "$TB/bin/os_load";
my $vnode_setup = "$TB/sbin/vnode_setup";
my $osselect    = "$TB/bin/os_select";
my $dbg		= 0;
my $failed      = 0;
my %nodes       = ();
my %vnodes      = ();
my %osids       = ();
my %canfail     = ();
my $db_result;
my @row;

#
# Ah, Frisbee works so lets do auto reloading for nodes that do not have
# the proper OS loaded on it. This will be a hash of lists; for each
# imageid, a list of the nodes to pass to os_load for that imageid. 
#
my %reloads     = ();
my %reboots	= ();
my $doautoload  = 1;
my $dolastload  = 1;
    
# un-taint path
$ENV{'PATH'} = '/bin:/usr/bin:/usr/local/bin';
delete @ENV{'IFS', 'CDPATH', 'ENV', 'BASH_ENV'};

$| = 1; #Turn off line buffering on output

#
# Parse command arguments. Once we return from getopts, all that should be
# left are the required arguments.
#
%options = ();
if (! getopts($optlist, \%options)) {
    usage();
}
if (@ARGV != 2) {
    usage();
}
if (defined($options{"d"})) {
    $dbg = 1;
}

my $pid = $ARGV[0];
my $eid = $ARGV[1];

#
# Untaint args.
#
if ($pid =~ /^([-\@\w]+)$/) {
    $pid = $1;
}
else {
    die_noretry("Bad data in pid: $pid.");
}
if ($eid =~ /^([-\@\w]+)$/) {
    $eid = $1;
}
else {
    die_noretry("Bad data in eid: $eid.");
}

#
# Figure out who called us. Only root, people with admin status
# in the DB, or the owner of the experiment can run this script.
#
if ($UID && !TBAdmin($UID) &&
    !TBExptAccessCheck($UID, $pid, $eid, TB_EXPT_MODIFY)) {
    die_noretry("*** $0:\n".
		"    You do not have permission to swap this experiment!");
}

TBDebugTimeStamp("os_setup started");

#
# Get the set of nodes, as well as the nodes table information for them.
#
$db_result =
    DBQueryFatal("select n.*,l.pid from reserved as r ".
		 "left join nodes as n on n.node_id=r.node_id ".
		 "left join last_reservation as l on n.node_id=l.node_id ".
		 "where r.pid='$pid' and r.eid='$eid'");

if ($db_result->numrows < 1) {	
    print "There are no nodes in experiment '$eid' in project '$pid'.\n";
    exit 0;
}

while (my %row = $db_result->fetchhash()) {
    my $node     = $row{'node_id'};
    my $osid     = $row{'def_boot_osid'};
    my $type     = $row{'type'};
    my $bootpath = 0;
    my $jailnode = 0;

    #
    # VIRTNODE HACK: Virtual nodes are special. Jailed vnodes can do quite
    # a bit, and so run them through the checks below.
    #
    if (TBIsNodeVirtual($node, \$jailed)) {
	$vnodes{$node} = $jailed;
	if (! $jailed) {
	    next;
	}
	$jailnode = 1;
    }
    else {
	my $nodeAllocState;
	TBGetNodeAllocState( $node, \$nodeAllocState );
	$nodes{$node}  = $node;
	$nodeAllocStates{$node} = $nodeAllocState;
	# only reboot node if assign_wrapper just pulled it into expt.
	# (e.g. it isnt ALLOCSTATE_RES_READY)
	#if (($nodeAllocState eq TBDB_ALLOCSTATE_RES_INIT_DIRTY()) ||
	#    ($nodeAllocState eq TBDB_ALLOCSTATE_RES_INIT_CLEAN())) {
	if ($nodeAllocState ne TBDB_ALLOCSTATE_RES_READY()) {
	    $reboots{$node} = 1;
	}
    }
    $osids{$node} = $osid;

    #
    # Make sure the files specified in the paths exist. We mount the
    # user tftp directory on boss node, so we can ignore the IP address,
    # and just check the path directly. 
    #
    if (defined($row{'def_boot_path'})) {
	my $path = $row{'def_boot_path'};

	if ($path ne "") {
	    my $ip   = 0;

	    # Split out IP address if it exists.
	    if ($path =~ /^([0-9\.]+):(\/.*)$/) {
		$ip   = $1;
		$path = $2;
	    }

	    # Path must begin with $TFTP
	    if (! ($path =~ /^\/$TFTP\//)) {
		die_noretry(
		    "*** File $path for node $node must reside in $TFTP");
	    }

	    if (! -f $path) {
		die_noretry("*** File $path for node $node does not exist!");
	    }
	    $bootpath = 1;
	}
    }
    if (defined($row{'next_boot_path'})) {
	my $path = $row{'next_boot_path'};

	if ($path ne "") {
	    my $ip   = 0;

	    # Split out IP address if it exists.
	    if ($path =~ /^([0-9\.]+):(\/.*)$/) {
		$ip   = $1;
		$path = $2;
	    }

	    # Path must begin with $TFTP
	    if (! ($path =~ /^\/$TFTP\//)) {
		die_noretry(
		    "*** File $path for node $node must reside in $TFTP");
	    }

	    if (! -f $path) {
		die_noretry("*** File $path for node $node does not exist!");
	    }
	}
    }

    #
    # XXX - Check for existence of the delta files. We do this here
    # cause its easier than looking for a failure later, when the node
    # tries to install the delta. Not a general solution though. Needs
    # more thought.
    #
    foreach my $delta (split(":", $row{'deltas'})) {
	if (! -f $delta) {
	    die_noretry("*** Delta file $delta for node $node does not exist!");
	}
    }
    #
    # XXX - Ditto for RPMs.
    #
    foreach my $rpm (split(":", $row{'rpms'})) {
	if (! -f $rpm) {
	    die_noretry("*** RPM $rpm for node $node does not exist!");
	}
    }
    
    #
    # XXX - Ditto for tarfiles.
    #
    foreach my $tarspec (split(":", $row{'tarballs'})) {
	my ($dir, $tar) = split(" ", $tarspec);
	
	if (! -f $tar) {
	    die_noretry("*** Tarfile $tar for node $node does not exist!");
	}
    }

    #
    # If there is a path specified, then we don't worry anymore about it.
    # The user must know what is going on. The OSID might have a path
    # associated with it, which means the same thing; we don't worry about
    # it. 
    #
    if (!$bootpath && !$jailnode) {
	#
	# These checks are not necessary if the front end and web page
	# are doing the right thing, but lets be careful anyway.
	# 
	if (! $osid) {
	    die_noretry(
	        "*** $node has no bootpath and no def_boot_osid set!");
	}

	#
	# Grab the info for this OSID. This is part of the image check.
	#
	my $osid_result =
	    DBQueryFatal("select * from os_info where osid='$osid'");
	
	if ($osid_result->numrows == 0) {
	    die_noretry("*** No such OSID $osid is defined!");
	}
	
	my %osid_row   = $osid_result->fetchhash();

	#
	# If there is an actual path, its an OSKit kernel not an image.
	# 
	if (! defined($osid_row{'path'}) || $osid_row{'path'} eq "") {
	    #
	    # Not an OSKit kernel.
	    # Make sure this OSID is actually loaded on the machine. 
	    #
	    my $p_result =
		DBQueryFatal("select * from partitions ".
			     "where node_id='$node' and osid='$osid'");

	    #
	    # If not loaded, then see if the user was looking for the generic
	    # name of the OS that is loaded. 
	    # 
	    if ($p_result->numrows == 0) {
		#
		# Check to see if a non specific version specified.
		#
		if (! defined($osid_row{'version'}) ||
		    $osid_row{'version'} eq "") {

		    #
		    # A non-specific version. There needs to be a way to
		    # map it to another osid. 
		    #
		    if (!defined($osid_row{'nextosid'})) {
			die_noretry(
			    "*** $0:\n".
			    "    No mapping can be made for $osid ($node)!");
		    }
		    my $nextosid = $osid_row{'nextosid'};
		    
		    #
		    # See if the nextosid is already on the disk. If not,
		    # it needs to be loaded.
		    # 
		    my $o_result =
			DBQueryFatal("select osid from partitions as p ".
				     "where p.node_id='$node' and ".
				     "      p.osid='$nextosid'");

		    if (! $o_result->numrows) {
			#
			# User wants a specific version of an OS, but its not
			# loaded on the machine. 
			#
			print "Mapping $osid on $node to $nextosid ".
			    "and setting up a reload.\n";
			
			SetupReload($node, $nextosid, $type);
			$osids{$node} = $nextosid;
		    }
		    else {
			#
			# Already loaded. 
			# 
			print "Mapping $osid on $node to $nextosid.\n";

			if ($dolastload &&
			    defined($row{'pid'}) && $row{'pid'} ne $pid) {
			    SetupReload($node, $nextosid, $type);
			}
			else {
			    system("$osselect $nextosid $node") and
				die_noretry("*** Could not set boot OS to ".
				    "$nextosid for $node");
			}
			$osids{$node} = $nextosid;
		    }
		}
		else {
		    #
		    # User wants a specific version of an OS, but its not
		    # loaded on the machine. 
		    #
		    SetupReload($node, $osid, $type);
		}
	    }
	    else {
		#
		# OSID is loaded, but might need to be cleaned. 
		#
		if ($dolastload &&
		    defined($row{'pid'}) && $row{'pid'} ne $pid) {
		    SetupReload($node, $osid, $type);
		}
	    }
	}
    }
    
    #
    # Set the canfail bit. Currently, sharks are always canfail=1.
    # Will come from DB at some point.
    #
    if ($row{'type'} eq "dnard") {
	$canfail{$node} = 1;
    }
    else {
	$canfail{$node} = 0;
    }
    
    print STDERR "$node - $osids{$node} - $canfail{$node}\n"
	if $dbg;
}

#
# Now do osid setup for jailed nodes. We waited until the physnodes were
# done above so that we can set the osid for the vnodes to the same
# as the physnode. We rely on the fact that the user is not allowed to
# set the OS for jailed nodes or for the physnodes that are hosting
# jailed nodes, and the node_types table has the right stuff. Non-jailed
# nodes do not need to be done.
#
foreach my $vnode (keys(%vnodes)) {
    my $jailed = $vnodes{$vnode};
    my $pnode;

    if (! $jailed) {
	next;
    }

    if (! TBPhysNodeID($vnode, \$pnode)) {
	die_noretry("*** $0:\n".
	    "    Cannot determine phys_nodeid for $vnode!");
    }
    my $n_osid;
    
    if (!defined($nodes{$pnode})) {
	#
	# Typical on remote nodes; we do not allocate the underlying
	# phys node to the experiment.
	#
	next;
    }
    else {
	$n_osid = $osids{$pnode};
    }
    
    system("$osselect $n_osid $vnode") and
	die_noretry("*** Could not set boot OS to $n_osid for $vnode");

    $osids{$vnode} = $n_osid;
}

#
# We need to issue the reboots and the reloads in parallel.
#
TBDebugTimeStamp("rebooting/reloading started");
if (!$TESTMODE) {
    my %pids  = ();
    my $count = 0;
    my $cmd;

    if (keys(%reboots)) {
	foreach my $node (keys(%reboots)) {
	    if ($nodeAllocStates{$node} eq TBDB_ALLOCSTATE_RES_INIT_CLEAN()) {
		TBSetNodeAllocState($node, TBDB_ALLOCSTATE_RES_REBOOT_CLEAN());
		$nodeAllocStates{$node} = TBDB_ALLOCSTATE_RES_REBOOT_CLEAN();
	    } else {
		TBSetNodeAllocState($node, TBDB_ALLOCSTATE_RES_REBOOT_DIRTY());
		$nodeAllocStates{$node} = TBDB_ALLOCSTATE_RES_REBOOT_DIRTY();
	    }
	}

	$cmd = "$nodereboot " . join(" ", keys(%reboots));
	$pids{$cmd} = ForkCmd($cmd);
    }

    foreach my $imageid ( keys(%reloads) ) {
	my @list = @{ $reloads{$imageid} };

	foreach my $node (@list) {
	    TBSetNodeAllocState( $node, TBDB_ALLOCSTATE_RES_RELOAD() );
	    $nodeAllocStates{$node} = TBDB_ALLOCSTATE_RES_RELOAD();
	}

	sleep(5);
	$pids{"$os_load -m $imageid @list"} =
	    ForkCmd("$os_load -m $imageid @list");
    }

    foreach $cmd ( keys(%pids) ) {
	my $pid = $pids{$cmd};

	waitpid($pid, 0);
	if ($?) {
	    $failed++;
	    print "*** Failed: $cmd\n";
	}
    }
}
TBDebugTimeStamp("rebooting/reloading finished");
sleep(2);

#
# XXX What happens if something above fails? We could exit, but some nodes
# that *are* rebooting would be caught in the middle. For the nodes that
# were reloaded, we can check the state right away (and avoid the wait
# below as well); they should be in the ISUP state when os_load is
# finished.  If not, thats a failure and we can save some time below.  For
# plain reboot failures, nothing to do but find out below after the wait.
# I do not want to exit right away cause we might end up with a lot more
# power cycles since the nodes are very likely to be in a non responsive
# state if just rebooted!
#
foreach my $imageid ( keys(%reloads) ) {
    my @list = @{ $reloads{$imageid} };

    foreach my $node ( @list ) {
	my $mode;
	
	if (!TBGetNodeOpMode($node, \$mode)) {
	    print "*** Error getting operational mode for $node!\n";
	    $failed++;
	    delete($nodes{$node});
	}
	if ($mode eq TBDB_NODEOPMODE_RELOAD) {
	    print "*** Not waiting for $node since its reload failed!\n";
	    $failed++;
	    delete($nodes{$node});
	}
    }
}
# Remaining nodes we need to wait for.
my @nodelist = keys(%nodes);

#
# Now lets wait for them to come back alive. Set up a retry list though
# so that we can give each node at least 1 second chance. Avoids pointless
# experiment failures.
#
if (@nodelist) {
    print "Waiting for local testbed nodes to finish rebooting ...\n";
}

my %retries;
my %waitstart;
foreach my $node ( @nodelist ) {
    $retries{$node} = 1;
    $waitstart{$node} = time;
}

TBDebugTimeStamp("Local node waiting started");
while ( @nodelist ) {
    my $node   = shift(@nodelist);
    my $wstart = $waitstart{$node};

    if (!TBNodeStateWait($node, TBDB_NODESTATE_ISUP, $wstart, (60*7))) {
	print "$node is alive and well\n";
	SetNodeBootStatus($node, NODEBOOTSTATUS_OKAY);
	TBSetNodeAllocState( $node, TBDB_ALLOCSTATE_RES_READY() );
	$nodeAllocStates{$node} = TBDB_ALLOCSTATE_RES_READY();	
	next;
    }

    if ($retries{$node}) {
	$retries{$node} -= 1;

	print "*** Rebooting $node and waiting again ...\n";
	
	if (system("$nodereboot $node") == 0) {
	    push(@nodelist, $node);
	    $waitstart{$node} = time;
	    next;
	}
	# Fall through on failure.
    }

    SetNodeBootStatus($node, NODEBOOTSTATUS_FAILED);
	
    print "*** WARNING: $node may be down.\n".
	  "    This has been reported to testbed-ops.\n";
	
    if ($canfail{$node}) {
	# Send mail to testbed-ops and to the user about it.
	my ($user) = getpwuid($UID);
	
	SENDMAIL($user, "Node $node is down",
		 "Node $node in pid/eid $pid/$eid appears to be dead.\n\n".
		 "Your experiment will continue to run since this failure\n".
		 "is nonfatal, although you might encounter other problems\n".
		 "if your experiment depends explicitly on this node.\n".
		 "You should terminate this experiment if it cannot ".
		 "tolerate this failure.\n\n".
		 "Testbed Operations has also been notified so they can ".
		 "investigate.\n\n".
		 "Thanks\n".
		 "Testbed Operations\n",
		 0,
		 "Cc: $TBOPS");

	print "*** Continuing with experiment setup anyway ...\n";
	next;
    }

    # Reserve it to down experiment.
    MarkNodeDown($node);
    TBSetNodeAllocState( $node, TBDB_ALLOCSTATE_DOWN() );
    $nodeAllocStates{$node} = TBDB_ALLOCSTATE_DOWN();	

    # Send mail to testbed-ops about it
    SENDMAIL($TBOPS, "Node $node is down",
	     "Node $node in pid/eid $pid/$eid appears to be dead.\n\n".
	     "$node has been taken out of the pool until this matter ".
	     "is resolved.\n");

    $failed++;
}
TBDebugTimeStamp("Local node waiting finished");

#
# Now deal with virtual nodes.
#
# We do this in a sub script since nodes are not owned by the user
# and so must be setuid root so that ssh will work.
#
# XXX - Don't bother if something above failed. A waste of time and
# usually leads to cascading errors.
#
my @vnodelist = keys(%vnodes);

if ($failed && @vnodelist) {
    print "*** Skipping virtual node setup since there were previous ".
	"failures!\n";
}
elsif (@vnodelist) {
    print "Setting up virtual testbed nodes ...\n";
    system("$vnode_setup $pid $eid");
    
    if ($?) {
	die_noretry("*** $0:\n".
	    "    Vnode setup failed!");
    }

    foreach my $node (@vnodelist) {
	$waitstart{$node} = time;
    }
    my $maxwait = 120 + (30 * scalar(@vnodelist));

    print "Waiting for virtual testbed nodes to finish setting up ...\n";

    TBDebugTimeStamp("Virtual node waiting started");
    while ( @vnodelist ) {
	my $node   = shift(@vnodelist);
	my $wstart = $waitstart{$node};

	if (!TBNodeStateWait($node, TBDB_NODESTATE_ISUP, $wstart, $maxwait)) {
	    print "$node is alive and well\n";
	    SetNodeBootStatus($node, NODEBOOTSTATUS_OKAY);
	    next;
	}

	SetNodeBootStatus($node, NODEBOOTSTATUS_FAILED);
	
	print "*** WARNING: $node may be down.\n".
	      "    This has been reported to testbed-ops.\n";
	
	if ($canfail{$node}) {
	    # Send mail to testbed-ops and to the user about it.
	    my ($user) = getpwuid($UID);
	
	    SENDMAIL($user, "Virtual Node $node is down",
		 "Node $node in pid/eid $pid/$eid appears to be dead.\n\n".
		 "Your experiment will continue to run since this failure\n".
		 "is nonfatal, although you might encounter other problems\n".
		 "if your experiment depends explicitly on this node.\n".
		 "You should terminate this experiment if it cannot ".
		 "tolerate this failure.\n\n".
		 "Testbed Operations has also been notified so they can ".
		 "investigate.\n\n".
		 "Thanks\n".
		 "Testbed Operations\n",
		 0,
		 "Cc: $TBOPS");

	    print "*** Continuing with experiment setup anyway ...\n";
	    next;
	}

	# Reserve it to down experiment. 
	# MarkNodeDown($node);

	# Send mail to testbed-ops about it
	SENDMAIL($TBOPS, "Virtual Node $node is down",
		 "Virtual node $node in pid/eid $pid/$eid appears to be ".
		 "unresponsive.\n\n");

	print "*** Experiment will be terminated automatically.\n";
	$failed++;
    }
    TBDebugTimeStamp("Virtual node waiting finished");
}

print "OS Setup Done. There were $failed failed nodes.\n";
TBDebugTimeStamp("os_setup finished");

if ($failed > 0) {
    exit 1;
} else {
    exit 0;
}

#
# Map an OSID to an imageid for a node type.
#
sub TBMapOSIDtoImageID($$)
{
    my ($osid, $type) = @_;

    my $query_result =
	DBQueryFatal("select imageid from osidtoimageid ".
		     "where type='$type' and osid='$osid'");

    if ($query_result->numrows == 0) {
	return 0;
    }
    my ($imageid) = $query_result->fetchrow_array();

    return $imageid;
}

#
# Setup a reload of a node if we can find an image.
# This goo constructs a hashed array of lists.
#
sub SetupReload($$$)
{
    my ($node, $osid, $type) = @_;

    if ((my $imageid = TBMapOSIDtoImageID($osid, $type))) {
	if (! defined($reloads{$imageid})) {
	    $reloads{$imageid} = [ $node ];
	}
	else {
	    push(@{ $reloads{$imageid} }, $node);
	}
	delete $reboots{$node};
    }
    else {
	die_noretry("*** $0:\n".
	    "    No image can be found for $osid on $node!");
    }
}

#
# Fork a process to exec a command. Return the pid to wait on.
# 
sub ForkCmd($) {
    my ($cmd) = @_;
    my($mypid);

    $mypid = fork();
    if ($mypid) {
	return $mypid;
    }

    if ($dbg) {
	print STDERR "Forking command: $cmd\n";
    }

    system($cmd);
    exit($? >> 8);
}