server/autoserv

#!/usr/bin/python -u
#
# Copyright 2007 Google Inc. Released under the GPL v2

"""
Run an control file through the server side engine
"""

__author__ = """\
mbligh@google.com (Martin J. Bligh)
"""

import sys, os, re, traceback, signal

import common
from autotest_lib.server import server_job, utils
from autotest_lib.server import autoserv_parser


class PidFileManager(object):
    pid_file = None

    def open_pid_file(self, results_dir):
        pid_file_path = os.path.join(results_dir, '.autoserv_execute')
        assert not os.path.exists(pid_file_path)
        self.pid_file = open(pid_file_path, 'w')
        self.pid_file.write(str(os.getpid()) + '\n')
        self.pid_file.flush()


    def close_pid_file(self, exit_code, signal_code=0):
        if not self.pid_file:
            return
        real_exit_code = (exit_code << 8) | (signal_code & 0xFF)
        self.pid_file.write(str(real_exit_code) + '\n')
        self.pid_file.close()
        self.pid_file = None


def run_autoserv(pid_file_manager, results, parser):
    # send stdin to /dev/null
    dev_null = os.open(os.devnull, os.O_RDONLY)
    os.dup2(dev_null, sys.stdin.fileno())
    os.close(dev_null)

    # Create separate process group
    os.setpgrp()

    # Implement SIGTERM handler
    def handle_sigint(signum, frame):
        pid_file_manager.close_pid_file(1, signal.SIGTERM)
        os.killpg(os.getpgrp(), signal.SIGKILL)

    # Set signal handler
    signal.signal(signal.SIGTERM, handle_sigint)

    # Get a useful value for running 'USER'
    realuser = os.environ.get('USER')
    if not realuser:
        realuser = 'anonymous'

    machines = parser.options.machines.split(',')
    machines_file = parser.options.machines_file
    label    = parser.options.label
    user     = parser.options.user
    client   = parser.options.client
    reboot   = parser.options.reboot
    install_before = parser.options.install_before
    install_after  = parser.options.install_after
    verify   = parser.options.verify
    repair   = parser.options.repair
    no_tee   = parser.options.no_tee
    parse_job = parser.options.parse_job
    host_protection = parser.options.host_protection
    ssh_user = parser.options.ssh_user
    ssh_port = parser.options.ssh_port
    ssh_pass = parser.options.ssh_pass

    if len(parser.args) < 1 and not verify and not repair:
        print parser.parser.print_help()
        sys.exit(1)

    # We have a control file unless it's just a verify/repair job
    if len(parser.args) > 0:
        control = parser.args[0]
    else:
        control = None

    if machines_file:
        machines = []
        for m in open(machines_file, 'r').readlines():
            # remove comments, spaces
            m = re.sub('#.*', '', m).strip()
            if m:
                machines.append(m)
        print "Read list of machines from file: %s" % machines_file
        print ','.join(machines)

    if machines:
        for machine in machines:
            if not machine or re.search('\s', machine):
                print "Invalid machine %s" % str(machine)
                sys.exit(1)
        machines = list(set(machines))
        machines.sort()

    job = server_job.server_job(control, parser.args[1:], results, label,
                                user, machines, client, parse_job,
                                ssh_user, ssh_port, ssh_pass)
    debug_dir = os.path.join(results, 'debug')
    stdout = os.path.join(debug_dir, 'autoserv.stdout')
    stderr = os.path.join(debug_dir, 'autoserv.stderr')
    if no_tee:
        job.stdout.redirect(stdout)
        job.stderr.redirect(stderr)
    else:
        job.stdout.tee_redirect(stdout)
        job.stderr.tee_redirect(stderr)

    # perform checks
    job.precheck()

    # run the job
    exit_code = 0
    try:
        if repair:
            job.repair(host_protection)
        elif verify:
            job.verify()
        else:
            try:
                job.run(reboot, install_before, install_after)
            finally:
                job.cleanup_parser()
    except:
        job.aborted = True
        traceback.print_exc()

    if getattr(job, 'aborted', False):
        sys.exit(1)


def main():
    pid_file_manager = PidFileManager()

    # grab the parser
    parser = autoserv_parser.autoserv_parser

    if len(sys.argv) == 1:
        parser.parser.print_help()
        sys.exit(1)

    results  = parser.options.results
    results  = os.path.abspath(results)
    write_pidfile = parser.options.write_pidfile
    if write_pidfile:
        pid_file_manager.open_pid_file(results)

    exit_code = 0
    try:
        try:
            run_autoserv(pid_file_manager, results, parser)
        except SystemExit, e:
            exit_code = e.code
        except:
            traceback.print_exc()
            # If we don't know what happened, we'll classify it as
            # an 'abort' and return 1.
            exit_code = 1
    finally:
        pid_file_manager.close_pid_file(exit_code)
    sys.exit(exit_code)


if __name__ == '__main__':
    main()