automation/server/server.py

#!/usr/bin/python2
#
# Copyright 2010 Google Inc. All Rights Reserved.

import logging
import optparse
import pickle
import signal
from SimpleXMLRPCServer import SimpleXMLRPCServer
import sys

from automation.common import logger
from automation.common.command_executer import CommandExecuter
from automation.server import machine_manager
from automation.server.job_group_manager import JobGroupManager
from automation.server.job_manager import JobManager


class Server(object):
  """Plays a role of external interface accessible over XMLRPC."""

  def __init__(self, machines_file=None, dry_run=False):
    """Default constructor.

    Args:
      machines_file: Path to file storing a list of machines.
      dry_run: If True, the server only simulates command execution.
    """
    CommandExecuter.Configure(dry_run)

    self.job_manager = JobManager(
        machine_manager.MachineManager.FromMachineListFile(
            machines_file or machine_manager.DEFAULT_MACHINES_FILE))

    self.job_group_manager = JobGroupManager(self.job_manager)

    self._logger = logging.getLogger(self.__class__.__name__)

  def ExecuteJobGroup(self, job_group, dry_run=False):
    job_group = pickle.loads(job_group)
    self._logger.info('Received ExecuteJobGroup(%r, dry_run=%s) request.',
                      job_group, dry_run)

    for job in job_group.jobs:
      job.dry_run = dry_run
    return self.job_group_manager.AddJobGroup(job_group)

  def GetAllJobGroups(self):
    self._logger.info('Received GetAllJobGroups() request.')
    return pickle.dumps(self.job_group_manager.GetAllJobGroups())

  def KillJobGroup(self, job_group_id):
    self._logger.info('Received KillJobGroup(%d) request.', job_group_id)
    self.job_group_manager.KillJobGroup(pickle.loads(job_group_id))

  def GetJobGroup(self, job_group_id):
    self._logger.info('Received GetJobGroup(%d) request.', job_group_id)

    return pickle.dumps(self.job_group_manager.GetJobGroup(job_group_id))

  def GetJob(self, job_id):
    self._logger.info('Received GetJob(%d) request.', job_id)

    return pickle.dumps(self.job_manager.GetJob(job_id))

  def GetMachineList(self):
    self._logger.info('Received GetMachineList() request.')

    return pickle.dumps(self.job_manager.machine_manager.GetMachineList())

  def StartServer(self):
    self.job_manager.StartJobManager()

  def StopServer(self):
    self.job_manager.StopJobManager()
    self.job_manager.join()


def GetServerOptions():
  """Get server's settings from command line options."""
  parser = optparse.OptionParser()
  parser.add_option('-m',
                    '--machines-file',
                    dest='machines_file',
                    help='The location of the file '
                    'containing the machines database',
                    default=machine_manager.DEFAULT_MACHINES_FILE)
  parser.add_option('-n',
                    '--dry-run',
                    dest='dry_run',
                    help='Start the server in dry-run mode, where jobs will '
                    'not actually be executed.',
                    action='store_true',
                    default=False)
  return parser.parse_args()[0]


def Main():
  logger.SetUpRootLogger(filename='server.log', level=logging.DEBUG)

  options = GetServerOptions()
  server = Server(options.machines_file, options.dry_run)
  server.StartServer()

  def _HandleKeyboardInterrupt(*_):
    server.StopServer()
    sys.exit(1)

  signal.signal(signal.SIGINT, _HandleKeyboardInterrupt)

  try:
    xmlserver = SimpleXMLRPCServer(
        ('localhost', 8000),
        allow_none=True,
        logRequests=False)
    xmlserver.register_instance(server)
    xmlserver.serve_forever()
  except Exception as ex:
    logging.error(ex)
    server.StopServer()
    sys.exit(1)


if __name__ == '__main__':
  Main()