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)
"""

from common.check_version import check_python_version
check_python_version()

import sys, os, re, server_job, hosts.site_host, utils, traceback, signal

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()


pid_file_manager = PidFileManager()


# 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)


usage = """\
usage: autoserv
	[-h, --help]               # This help message
	[-m machine,[machine,...]] # list of machines to pass to control file
	[-M machines_file]         # list of machines (from a file)
	[-c]                       # control file is a client side control
	[-r resultsdir]            # specify results directory (default '.')
	[-i]                       # reinstall machines before running the job
	[-I]                       # reinstall machines after running the job
	[-b]                       # reboot all specified machines after the job
	[-l label]                 # label for the job (arbitrary string)
	[-u user]                  # username for the job (email address)
	[-v]                       # verify the machines only
	[-R]                       # repair the machines
	[-n]                       # no teeing the status to stdout/stderr
	[-p]                       # write pidfile (.autoserv_execute)
	<control file>             # name of the control file to run
	[args ...]                 # args to pass through to the control file
"""

args = sys.argv[1:]
parser = utils.AutoservOptionParser(args)

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

machines = parser.parse_opts_param('-m', None, split = ',')
machines_file = parser.parse_opts_param('-M', None)
results  = parser.parse_opts_param('-r', os.path.abspath('.'))
results  = os.path.abspath(results)
label    = parser.parse_opts_param('-l', '')
user     = parser.parse_opts_param('-u', realuser)
client   = parser.parse_opts('-c')
reboot   = parser.parse_opts('-b')
install_before = parser.parse_opts('-i')
install_after  = parser.parse_opts('-I')
verify   = parser.parse_opts('-v')
repair   = parser.parse_opts('-R')
no_tee   = parser.parse_opts('-n')
help     = parser.parse_opts('-h') or parser.parse_opts('--help')
write_pidfile = parser.parse_opts('-p')

if help is True:
  print usage
  sys.exit(0)

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

if machines_file:
	machines = []
	for m in open(machines_file, 'r').readlines():
		m = re.sub('#.*', '', m).strip()   # remove comments, spaces
		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()

# 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

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

if write_pidfile:
	pid_file_manager.open_pid_file(results)

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

if getattr(job, 'aborted', False):
	exit_code = 1
pid_file_manager.close_pid_file(exit_code)

sys.exit(exit_code)