source: fedd/federation/proxy_emulab_segment.py @ 1d913e13

#!/usr/local/bin/python

import sys, os
import re

import tempfile
import subprocess
import logging 
import time
import signal

class proxy_emulab_segment:
    class ssh_cmd_timeout(RuntimeError): pass

    def __init__(self, log=None, keyfile=None, debug=False):
        self.log = log or logging.getLogger(\
                'fedd.access.proxy_emulab_segment')
        self.ssh_privkey_file = keyfile
        self.debug = debug
        self.ssh_exec="/usr/bin/ssh"
        self.scp_exec = "/usr/bin/scp"
        self.ssh_cmd_timeout = proxy_emulab_segment.ssh_cmd_timeout

    def scp_file(self, file, user, host, dest=""):
        """
        scp a file to the remote host.  If debug is set the action is only
        logged.
        """

        scp_cmd = [self.scp_exec, '-o', 'IdentitiesOnly yes', 
                '-o', 'StrictHostKeyChecking yes', '-i', 
                self.ssh_privkey_file, file, 
                "%s@%s:%s" % (user, host, dest)]
        rv = 0

        try:
            dnull = open("/dev/null", "w")
        except IOError:
            self.log.debug("[ssh_file]: failed to open " + \
                    "/dev/null for redirect")
            dnull = Null

        self.log.debug("[scp_file]: %s" % " ".join(scp_cmd))
        if not self.debug:
            rv = subprocess.call(scp_cmd, stdout=dnull, 
                    stderr=dnull, close_fds=True, close_fds=True)

        return rv == 0

    def ssh_cmd(self, user, host, cmd, wname=None, timeout=None):
        """
        Run a remote command on host as user.  If debug is set, the action
        is only logged.  Commands are run without stdin, to avoid stray
        SIGTTINs.
        """
        sh_str = ("%s -n -o 'IdentitiesOnly yes' -o " + \
                "'StrictHostKeyChecking yes' -i %s %s@%s %s") % \
                (self.ssh_exec, self.ssh_privkey_file, 
                        user, host, cmd)

        try:
            dnull = open("/dev/null", "w")
        except IOError:
            self.log.debug("[ssh_cmd]: failed to open /dev/null " + \
                    "for redirect")
            dnull = Null

        self.log.debug("[ssh_cmd]: %s" % sh_str)
        if not self.debug:
            if dnull:
                sub = subprocess.Popen(sh_str, shell=True, stdout=dnull,
                        stderr=dnull, close_fds=True)
            else:
                sub = subprocess.Popen(sh_str, shell=True, close_fds=True)
            if timeout:
                i = 0
                rv = sub.poll()
                while i < timeout:
                    if rv is not None: break
                    else:
                        time.sleep(1)
                        rv = sub.poll()
                        i += 1
                else:
                    self.log.debug("Process exceeded runtime: %s" % sh_str)
                    os.kill(sub.pid, signal.SIGKILL)
                    raise self.ssh_cmd_timeout();
                return rv == 0
            else:
                return sub.wait() == 0
        else:
            if timeout == 0:
                self.log.debug("debug timeout raised on %s " % sh_str)
                raise self.ssh_cmd_timeout()
            else:
                return True

class start_segment(proxy_emulab_segment):
    def __init__(self, log=None, keyfile=None, debug=False):
        proxy_emulab_segment.__init__(self, log=log, 
                keyfile=keyfile, debug=debug)
        self.null = """
set ns [new Simulator]
source tb_compat.tcl

set a [$ns node]

$ns rtproto Session
$ns run 
"""

    def get_state(self, user, host, pid, eid):
        # command to test experiment state
        expinfo_exec = "/usr/testbed/bin/expinfo"  
        # Regular expressions to parse the expinfo response
        state_re = re.compile("State:\s+(\w+)")
        no_exp_re = re.compile("^No\s+such\s+experiment")
        swapping_re = re.compile("^No\s+information\s+available.")
        state = None    # Experiment state parsed from expinfo
        # The expinfo ssh command.  Note the identity restriction to use
        # only the identity provided in the pubkey given.
        cmd = [self.ssh_exec, '-o', 'IdentitiesOnly yes', '-o', 
                'StrictHostKeyChecking yes', '-i', 
                self.ssh_privkey_file, "%s@%s" % (user, host), 
                expinfo_exec, pid, eid]

        dev_null = None
        try:
            dev_null = open("/dev/null", "a")
        except IOError, e:
            self.log.error("[get_state]: can't open /dev/null: %s" %e)

        if self.debug:
            state = 'swapped'
            rv = 0
        else:
            self.log.debug("Checking state")
            status = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                    stderr=dev_null, close_fds=True)
            for line in status.stdout:
                m = state_re.match(line)
                if m: state = m.group(1)
                else:
                    for reg, st in ((no_exp_re, "none"),
                            (swapping_re, "swapping")):
                        m = reg.match(line)
                        if m: state = st
            rv = status.wait()

        # If the experiment is not present the subcommand returns a
        # non-zero return value.  If we successfully parsed a "none"
        # outcome, ignore the return code.
        if rv != 0 and state != 'none':
            raise service_error(service_error.internal,
                    "Cannot get status of segment:%s/%s" % (pid, eid))
        elif state not in ('active', 'swapped', 'swapping', 'none'):
            raise service_error(service_error.internal,
                    "Cannot get status of segment:%s/%s" % (pid, eid))
        else:
            self.log.debug("State is %s" % state)
            return state


    def __call__(self, parent, eid, pid, user, tclfile, tmpdir, timeout=0):
        """
        Start a sub-experiment on a federant.

        Get the current state, modify or create as appropriate, ship data
        and configs and start the experiment.  There are small ordering
        differences based on the initial state of the sub-experiment.
        """
        # ops node in the federant
        host = "%s%s" % (parent.ops, parent.domain)
        # Configuration directories on the remote machine
        proj_dir = "/proj/%s/exp/%s/tmp" % (pid, eid)
        softdir = "/proj/%s/software/%s" % (pid, eid)
        # Local software dir
        lsoftdir = "%s/software" % tmpdir

        state = self.get_state(user, host, pid, eid)

        if not self.scp_file(tclfile, user, host):
            return False
        
        if state == 'none':
            # Create a null copy of the experiment so that we capture any
            # logs there if the modify fails.  Emulab software discards the
            # logs from a failed startexp
            try:
                f = open("%s/null.tcl" % tmpdir, "w")
                print >>f, self.null
                f.close()
            except IOError, e:
                raise service_error(service_error.internal,
                        "Cannot stage tarfile/rpm: %s" % e.strerror)

            if not self.scp_file("%s/null.tcl" % tmpdir, user, host):
                return False
            self.log.info("[start_segment]: Creating %s" % eid)
            timedout = False
            try:
                if not self.ssh_cmd(user, host,
                        ("/usr/testbed/bin/startexp -i -f -w -p %s " + 
                        "-e %s null.tcl") % (pid, eid), "startexp",
                        timeout=60 * 10):
                    return False
            except self.ssh_cmd_timeout:
                timedout = True

            if timedout:
                state = self.get_state(user, host, pid, eid)
                if state != "swapped":
                    return False
        
        # Open up a temporary file to contain a script for setting up the
        # filespace for the new experiment.
        self.log.info("[start_segment]: creating script file")
        try:
            sf, scriptname = tempfile.mkstemp()
            scriptfile = os.fdopen(sf, 'w')
        except IOError:
            return False

        scriptbase = os.path.basename(scriptname)

        # Script the filesystem changes
        print >>scriptfile, "/bin/rm -rf %s" % proj_dir
        # Clear and create the software directory
        print >>scriptfile, "/bin/rm -rf %s/*" % softdir
        print >>scriptfile, 'mkdir -p %s' % proj_dir
        if os.path.isdir(lsoftdir):
            print >>scriptfile, 'mkdir -p %s' % softdir
        print >>scriptfile, "rm -f %s" % scriptbase
        scriptfile.close()

        # Move the script to the remote machine
        # XXX: could collide tempfile names on the remote host
        if self.scp_file(scriptname, user, host, scriptbase):
            os.remove(scriptname)
        else:
            return False

        # Execute the script (and the script's last line deletes it)
        if not self.ssh_cmd(user, host, "sh -x %s" % scriptbase):
            return False

        for f in os.listdir(tmpdir):
            if not os.path.isdir("%s/%s" % (tmpdir, f)):
                if not self.scp_file("%s/%s" % (tmpdir, f), user, host,
                        "%s/%s" % (proj_dir, f)):
                    return False
        if os.path.isdir(lsoftdir):
            for f in os.listdir(lsoftdir):
                if not os.path.isdir("%s/%s" % (lsoftdir, f)):
                    if not self.scp_file("%s/%s" % (lsoftdir, f), 
                            user, host, "%s/%s" % (softdir, f)):
                        return False
        # Stage the new configuration (active experiments will stay swapped
        # in now)
        self.log.info("[start_segment]: Modifying %s" % eid)
        try:
            if not self.ssh_cmd(user, host,
                    "/usr/testbed/bin/modexp -r -s -w %s %s %s" % \
                            (pid, eid, tclfile.rpartition('/')[2]),
                    "modexp", timeout= 60 * 10):
                return False
        except self.ssh_cmd_timeout:
            self.log.error("Modify command failed to complete in time")
            # There's really no way to see if this succeeded or failed, so
            # if it hangs, assume the worst.
            return False
        # Active experiments are still swapped, this swaps the others in.
        if state != 'active':
            self.log.info("[start_segment]: Swapping %s" % eid)
            timedout = False
            try:
                if not self.ssh_cmd(user, host,
                        "/usr/testbed/bin/swapexp -w %s %s in" % (pid, eid),
                        "swapexp", timeout=10*60):
                    return False
            except self.ssh_cmd_timeout:
                timedout = True
           
            # If the command was terminated, but completed successfully,
            # report success.
            if timedout:
                self.log.debug("[start_segment]: swapin timed out " +\
                        "checking state")
                state = self.get_state(user, host, pid, eid)
                self.log.debug("[start_segment]: state is %s" % state)
                return state == 'active'
        # Everything has gone OK.
        return True

class stop_segment(proxy_emulab_segment):
    def __init__(self, log=None, keyfile=None, debug=False):
        proxy_emulab_segment.__init__(self,
                log=log, keyfile=keyfile, debug=debug)

    def __call__(self, parent, user, pid, eid):
        """
        Stop a sub experiment by calling swapexp on the federant
        """
        host = "%s%s" % (parent.ops, parent.domain)
        self.log.info("[stop_segment]: Stopping %s" % eid)
        rv = False
        try:
            # Clean out tar files: we've gone over quota in the past
            self.ssh_cmd(user, host, "rm -rf /proj/%s/software/%s" % \
                    (pid, eid))
            rv = self.ssh_cmd(user, host,
                    "/usr/testbed/bin/swapexp -w %s %s out" % (pid, eid))
        except self.ssh_cmd_timeout:
            rv = False
        return rv